#!/usr/bin/env python3
"""Safe diagnostics and opt-in repairs for H3 Studio.

Diagnostics never download models or terminate processes. Repair actions are an
explicit allow-list and keep recoverable backups whenever local caches move.
"""

from __future__ import annotations

import os
import shutil
import subprocess
from datetime import datetime
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
CACHE_ROOT = ROOT / ".runtime-cache"
SETUP_COMMAND = "setup_h3_cn.cmd"
START_COMMAND = "start_h3_studio.cmd"
OPTIONAL_PACKAGES = ("opencv-python>=4.9", "modelscope>=1.28")


def _check(
    check_id: str,
    label: str,
    status: str,
    summary: str,
    detail: str,
    *,
    action: str | None = None,
    steps: list[str] | None = None,
    command: str = "",
) -> dict[str, Any]:
    return {
        "id": check_id,
        "label": label,
        "status": status,
        "summary": summary,
        "detail": detail,
        "action": action,
        "steps": steps or [],
        "command": command,
    }


def diagnose(
    deployment: dict[str, Any],
    *,
    service: str,
    service_message: str,
    comfy_reachable: bool,
    owns_comfy: bool,
    comfy_log: list[str],
) -> dict[str, Any]:
    """Return UI-friendly checks for the current machine and deployment plan."""
    checks: list[dict[str, Any]] = []
    runtime = deployment.get("runtime") or {}
    hardware = deployment.get("hardware") or {}
    profile = deployment.get("profile") or {}
    components = deployment.get("components") or []
    capabilities = deployment.get("capabilities") or {}
    modules = runtime.get("modules") or {}

    if runtime.get("usable"):
        checks.append(
            _check(
                "runtime",
                "独立运行时",
                "ok",
                f"Python {runtime.get('python_version') or '可用'} 已就绪",
                str(runtime.get("executable") or "项目运行时已通过探测"),
            )
        )
    else:
        checks.append(
            _check(
                "runtime",
                "独立运行时",
                "error",
                "没有找到兼容的 Python 运行时",
                "本平台不依赖 venv；一键安装器会在项目内创建 .runtime\\python。",
                steps=["关闭当前窗口", "回到项目根目录运行 setup_h3_cn.cmd", "安装结束后重新启动 H3 Studio"],
                command=SETUP_COMMAND,
            )
        )

    required_missing = [item for item in components if item.get("required") and not item.get("installed")]
    missing_models = [item for item in required_missing if item.get("kind") == "model"]
    missing_other = [item for item in required_missing if item.get("kind") != "model"]
    if not required_missing:
        checks.append(
            _check("core", "H3 核心文件", "ok", "所选部署档案的必需组件齐全", "模型、执行器与本地引擎路径均已找到。")
        )
    else:
        labels = "、".join(str(item.get("label")) for item in required_missing[:5])
        if len(required_missing) > 5:
            labels += f" 等 {len(required_missing)} 项"
        detail = f"缺少：{labels}。"
        if missing_models:
            detail += " 模型属于大文件，故障中心不会静默下载；安装器支持 ModelScope 国内断点续传。"
        checks.append(
            _check(
                "core",
                "H3 核心文件",
                "error",
                f"有 {len(required_missing)} 个必需组件未就位",
                detail,
                steps=["确保磁盘至少剩余 100GB", "再次运行 setup_h3_cn.cmd 续装", "不要移动 ComfyUI/models 内已下载的权重"],
                command=SETUP_COMMAND,
            )
        )

    missing_helpers = [name for name in ("modelscope", "cv2") if not modules.get(name)]
    if not missing_helpers:
        checks.append(
            _check("helpers", "下载与视频工具", "ok", "ModelScope 与 OpenCV 已就绪", "可使用国内模型下载与通用视频处理路径。")
        )
    elif runtime.get("usable"):
        checks.append(
            _check(
                "helpers",
                "下载与视频工具",
                "warning",
                "缺少 " + " / ".join(missing_helpers),
                "不会影响已经下载完成的模型，但会影响续装、视频节点或首次部署。",
                action="install_download_helpers",
                steps=["点击“自动修复”从清华 TUNA 安装缺失组件", "完成后重新检测", "如果网络失败，再运行一键安装器续装"],
            )
        )

    backend = str(runtime.get("backend") or "none")
    if capabilities.get("h3"):
        acceleration = str(profile.get("acceleration") or "pytorch")
        checks.append(
            _check(
                "compute",
                "GPU 计算后端",
                "ok",
                f"{backend.upper()} · {acceleration}",
                str(runtime.get("device_name") or "已按实际运行时启用 H3"),
            )
        )
    elif profile.get("h3_support") == "unavailable":
        checks.append(
            _check(
                "compute",
                "GPU 计算后端",
                "warning" if capabilities.get("upscale") else "error",
                "当前档案未启用本地 H3 生成",
                "终修仍可使用。若显卡应当支持 H3，请先更新显卡驱动，再运行一键安装器让其重新选择 PyTorch 后端。",
                steps=["确认显卡型号与显存", "安装厂商官方驱动", "重新运行 setup_h3_cn.cmd", "在部署向导中重新检测"],
                command=SETUP_COMMAND,
            )
        )

    if profile.get("acceleration") == "sageattention":
        checks.append(_check("attention", "注意力加速", "ok", "SageAttention 已启用", "Triton 与 SageAttention 均已通过运行时探测。"))
    elif backend == "cuda" and capabilities.get("h3"):
        checks.append(
            _check(
                "attention",
                "注意力加速",
                "info",
                "已安全回退到 PyTorch Attention",
                "这不会阻止生成，只会降低速度。SageAttention 必须与当前 Torch/CUDA 精确匹配，不建议盲装未知 wheel。",
                steps=["先保持 PyTorch 回退确认工作流可生成", "需要加速时重新运行一键安装器", "安装后运行 tools\\check_sageattention.py 验证"],
                command=r".runtime\python\python.exe tools\check_sageattention.py",
            )
        )

    finish_engine = str(profile.get("finish_engine") or "universal")
    if finish_engine == "maxine" and capabilities.get("maxine"):
        checks.append(_check("finish", "终修与超分", "ok", "NVIDIA Maxine / RTX 已就绪", "可执行时序降噪、1.5× VSR 与 NVENC 输出。"))
    else:
        checks.append(
            _check(
                "finish",
                "终修与超分",
                "info" if capabilities.get("upscale") else "error",
                "正在使用通用终修" if capabilities.get("upscale") else "没有可用的终修执行器",
                "非 RTX 显卡会自动采用保结构时序降噪与 Lanczos 超分；这是设计内的兼容路径，不属于故障。"
                if capabilities.get("upscale")
                else "请重新运行安装器补全通用终修依赖。",
                steps=[] if capabilities.get("upscale") else ["重新运行 setup_h3_cn.cmd", "完成后在部署向导中重新检测"],
                command="" if capabilities.get("upscale") else SETUP_COMMAND,
            )
        )

    if comfy_reachable:
        checks.append(
            _check(
                "engine",
                "本地引擎连接",
                "ok",
                "ComfyUI API 可以访问",
                "当前由 H3 Studio 启动。" if owns_comfy else "已连接到现有 ComfyUI 进程；故障中心不会结束外部进程。",
            )
        )
    elif profile.get("h3_support") != "unavailable":
        action = "restart_engine" if owns_comfy or service in {"error", "setup"} else None
        checks.append(
            _check(
                "engine",
                "本地引擎连接",
                "error",
                "ComfyUI API 暂时无法访问",
                service_message or "检查 8189 端口、启动日志和是否存在重复实例。",
                action=action,
                steps=["先查看工作流日志末尾", "关闭重复的 ComfyUI 窗口", "点击自动修复重启由 Studio 管理的引擎，或重新运行 start_h3_studio.cmd"],
                command=START_COMMAND,
            )
        )

    ram = float(hardware.get("ram_gb") or 0)
    disk = float(hardware.get("disk_free_gb") or 0)
    resource_status = "ok" if ram >= 48 and disk >= 100 else "warning"
    resource_summary = f"内存 {ram:g}GB · 可用磁盘 {disk:g}GB"
    resource_detail = "满足推荐容量。" if resource_status == "ok" else "H3 推荐至少 48GB 内存与 100GB 可用磁盘；不足时先缩短到 0.5 秒草图，并清理自己的输出文件。"
    checks.append(_check("resources", "内存与磁盘", resource_status, resource_summary, resource_detail))

    logs = "\n".join(comfy_log[-80:]).lower()
    known_log_issue = None
    if "database is locked" in logs or "database locked" in logs:
        known_log_issue = ("数据库被占用", "通常由重复启动 ComfyUI 引起。关闭其他 ComfyUI 窗口后重新启动 Studio。")
    elif "out of memory" in logs or "cuda error: out of memory" in logs:
        known_log_issue = ("显存不足", "关闭其他 GPU 程序，改用 608×352 / 0.5 秒预览或 RTX 低显存档；一次只运行一个 GPU 任务。")
    elif "no module named" in logs:
        known_log_issue = ("Python 组件缺失", "重新运行一键安装器续装依赖；不要把系统 Python 与项目运行时混用。")
    elif "triton" in logs and ("permission" in logs or "access is denied" in logs or "cache" in logs and "error" in logs):
        known_log_issue = ("Triton 缓存异常", "先停止引擎，再用自动修复重建项目内缓存；旧缓存会保留在 repair-backups。")
    if known_log_issue:
        action = "reset_runtime_caches" if known_log_issue[0] == "Triton 缓存异常" else None
        checks.append(
            _check(
                "logs",
                "日志特征",
                "error",
                known_log_issue[0],
                known_log_issue[1],
                action=action,
                steps=["按上方建议处理", "重新启动后点击“重新检测”", "若仍失败，复制日志末尾用于反馈"],
            )
        )
    else:
        checks.append(_check("logs", "日志特征", "ok", "未发现已知致命错误", "已检查最近的本地引擎日志。"))

    counts = {status: sum(1 for item in checks if item["status"] == status) for status in ("error", "warning", "info", "ok")}
    if counts["error"]:
        headline = f"发现 {counts['error']} 项需要处理的问题"
        overall = "error"
    elif counts["warning"]:
        headline = f"可运行，但有 {counts['warning']} 项建议处理"
        overall = "warning"
    else:
        headline = "当前环境通过常规检查"
        overall = "ok"
    return {
        "generated_at": datetime.now().astimezone().isoformat(timespec="seconds"),
        "overall": overall,
        "headline": headline,
        "counts": counts,
        "checks": checks,
        "service": service,
    }


def _runtime_command(deployment: dict[str, Any]) -> list[str]:
    command = [str(part) for part in (deployment.get("runtime", {}).get("command") or [])]
    if not command:
        raise RuntimeError("没有可用的 Python 运行时，请先运行 setup_h3_cn.cmd")
    executable = Path(command[0])
    if executable.is_absolute() and not executable.is_file():
        raise RuntimeError(f"Python 运行时不存在：{executable}")
    return command


def install_download_helpers(deployment: dict[str, Any]) -> dict[str, Any]:
    command = [
        *_runtime_command(deployment),
        "-m",
        "pip",
        "install",
        "--disable-pip-version-check",
        "-i",
        "https://pypi.tuna.tsinghua.edu.cn/simple",
        *OPTIONAL_PACKAGES,
    ]
    creationflags = subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0
    result = subprocess.run(
        command,
        cwd=ROOT,
        capture_output=True,
        text=True,
        encoding="utf-8",
        errors="replace",
        timeout=1200,
        creationflags=creationflags,
        check=False,
    )
    output = "\n".join(part.strip() for part in (result.stdout, result.stderr) if part.strip())
    if result.returncode:
        raise RuntimeError((output or f"pip 退出码 {result.returncode}")[-5000:])
    return {"message": "下载与视频辅助组件已安装；请重新检测。", "log": output[-5000:]}


def _is_link_or_junction(path: Path) -> bool:
    if path.is_symlink():
        return True
    isjunction = getattr(os.path, "isjunction", None)
    return bool(isjunction(path)) if isjunction else False


def reset_runtime_caches() -> dict[str, Any]:
    """Rotate only known project cache directories; never traverse links."""
    cache_root = CACHE_ROOT.resolve()
    CACHE_ROOT.mkdir(parents=True, exist_ok=True)
    if _is_link_or_junction(CACHE_ROOT):
        raise RuntimeError(".runtime-cache 是符号链接或联接点，已拒绝自动处理。")
    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    backup_root = CACHE_ROOT / "repair-backups" / stamp
    targets = [CACHE_ROOT / "triton", CACHE_ROOT / "torchinductor"]
    moved: list[str] = []
    for target in targets:
        if not target.exists():
            target.mkdir(parents=True, exist_ok=True)
            continue
        if _is_link_or_junction(target) or target.resolve().parent != cache_root:
            raise RuntimeError(f"缓存路径不安全，已停止：{target}")
        backup_root.mkdir(parents=True, exist_ok=True)
        destination = backup_root / target.name
        if destination.exists():
            raise RuntimeError(f"备份路径已存在：{destination}")
        shutil.move(str(target), str(destination))
        target.mkdir(parents=True, exist_ok=True)
        moved.append(target.name)
    return {
        "message": "运行时缓存已重建。旧缓存未删除，可从 repair-backups 恢复。",
        "backup": str(backup_root) if moved else "",
        "moved": moved,
    }


REPAIR_ACTIONS = {
    "install_download_helpers": install_download_helpers,
    "reset_runtime_caches": lambda deployment: reset_runtime_caches(),
}


def run_repair(action: str, deployment: dict[str, Any]) -> dict[str, Any]:
    handler = REPAIR_ACTIONS.get(action)
    if handler is None:
        raise ValueError("未知或不允许的修复操作")
    return handler(deployment)
