from __future__ import annotations

import base64
import json
import time
from typing import Any

from core.config import DEPLOY_REGION, TAT_METRIC_TIMEOUT_SECONDS, TENCENT_SECRET_ID, TENCENT_SECRET_KEY, utc_now


def decode_task_output(output: str) -> str:
    # 没有输出时直接返回空串。
    if not output:
        return ""

    # 把 TAT 返回的 base64 文本解码成普通字符串。
    return base64.b64decode(output).decode("utf-8", errors="replace")


def build_tat_client() -> tuple[Any | None, str]:
    # TAT 运行需要完整的腾讯云密钥和区域。
    if not TENCENT_SECRET_ID or not TENCENT_SECRET_KEY or not DEPLOY_REGION:
        return None, "missing TAT credentials in .env"

    try:
        # 运行时才导入 SDK，避免没装依赖时整个服务启动失败。
        from tencentcloud.common import credential
        # 腾讯云异常类型用于重试。
        from tencentcloud.common.profile.client_profile import ClientProfile
        # 指定 TAT API endpoint。
        from tencentcloud.common.profile.http_profile import HttpProfile
        # TAT 客户端。
        from tencentcloud.tat.v20201028 import tat_client
    except Exception as exc:  # noqa: BLE001
        # SDK 不存在时直接返回错误。
        return None, f"tencentcloud sdk unavailable: {exc}"

    # 创建凭据对象。
    cred = credential.Credential(TENCENT_SECRET_ID, TENCENT_SECRET_KEY)

    # 创建 TAT 客户端。
    client = tat_client.TatClient(
        cred,
        DEPLOY_REGION,
        ClientProfile(httpProfile=HttpProfile(endpoint="tat.tencentcloudapi.com")),
    )

    # 成功时返回客户端。
    return client, ""


def invoke_tat(action: Any, label: str, retries: int = 3, delay_seconds: int = 2) -> Any:
    # 这里延迟导入异常类型，避免 SDK 缺失时模块导入失败。
    from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException

    # 记录最后一次异常。
    last_error: Exception | None = None

    # 简单重试几次，防止瞬时 API 抖动。
    for attempt in range(1, retries + 1):
        try:
            # 成功时直接返回结果。
            return action()
        except TencentCloudSDKException as exc:
            # 保存最后一次异常。
            last_error = exc

            # 最后一次重试失败时不再等待。
            if attempt == retries:
                break

            # 做一个线性退避。
            time.sleep(delay_seconds * attempt)

    # 有异常时抛出最后一次异常。
    if last_error is not None:
        raise last_error

    # 理论上不会走到这里，这里只是兜底。
    raise RuntimeError(f"{label} failed without error")


def run_shell_command(
    instance_ids: list[str],
    script: str,
    *,
    command_name: str,
    timeout_seconds: int = TAT_METRIC_TIMEOUT_SECONDS,
    working_directory: str = "/root",
) -> dict[str, dict[str, Any]]:
    # 先初始化 TAT 客户端。
    client, error = build_tat_client()

    # 客户端不可用时直接抛错。
    if client is None:
        raise RuntimeError(error)

    # 这里延迟导入 TAT models。
    from tencentcloud.tat.v20201028 import models as tat_models

    # 创建命令执行请求。
    req = tat_models.RunCommandRequest()

    # 组装批量执行命令。
    req.from_json_string(
        json.dumps(
            {
                "CommandName": command_name,
                "CommandType": "SHELL",
                "Content": base64.b64encode(script.encode("utf-8")).decode("ascii"),
                "InstanceIds": instance_ids,
                "Timeout": timeout_seconds,
                "WorkingDirectory": working_directory,
            }
        )
    )

    # 发起远端命令执行。
    response = json.loads(invoke_tat(lambda: client.RunCommand(req), "RunCommand").to_json_string())
    # 拿到 invocation id，后面轮询执行结果。
    invocation_id = response["InvocationId"]
    # 设置轮询截止时间。
    deadline = time.time() + timeout_seconds + 60
    # 这里保存各实例的最终结果。
    results: dict[str, dict[str, Any]] = {}

    # 开始轮询任务结果。
    while time.time() < deadline:
        # 轮询间隔不必太短。
        time.sleep(2)

        # 创建查询结果请求。
        query = tat_models.DescribeInvocationTasksRequest()
        # 按 invocation id 查询。
        query.from_json_string(
            json.dumps(
                {
                    "HideOutput": False,
                    "Limit": 50,
                    "Filters": [{"Name": "invocation-id", "Values": [invocation_id]}],
                }
            )
        )

        # 拉取当前命令所有实例的执行状态。
        payload = json.loads(
            invoke_tat(lambda: client.DescribeInvocationTasks(query), "DescribeInvocationTasks").to_json_string()
        )
        # 任务列表可能暂时为空。
        tasks = payload.get("InvocationTaskSet") or []

        # 没任务时继续等待。
        if not tasks:
            continue

        # 遍历当前拿到的实例结果。
        for task in tasks:
            # 读取实例 ID。
            instance_id = task.get("InstanceId", "")
            # 读取任务状态。
            status = task.get("TaskStatus", "")

            # 只处理终态。
            if status not in {"SUCCESS", "FAILED", "TIMEOUT", "CANCELLED", "START_FAILED"}:
                continue

            # 解码远端输出。
            output = decode_task_output(task.get("TaskResult", {}).get("Output", ""))

            # 统一存成结构化结果。
            results[instance_id] = {
                "status": status,
                "output": output,
                "updatedAt": utc_now(),
            }

        # 收齐全部实例结果后结束轮询。
        if len(results) >= len(instance_ids):
            break

    # 为缺失结果的实例补超时占位。
    for instance_id in instance_ids:
        results.setdefault(
            instance_id,
            {
                "status": "TIMEOUT",
                "output": "",
                "updatedAt": utc_now(),
            },
        )

    # 返回按实例 ID 索引的结果。
    return results