"""Hardware-aware deployment planning for H3 Studio.

The planner is intentionally read-only.  It probes the local machine, checks
the files already present in the workspace, and produces a beginner-friendly
installation/configuration plan.  The web API may persist the user's selected
profile, but dependency downloads remain an explicit user action.
"""

from __future__ import annotations

import json
import os
import platform
import shutil
import subprocess
import sys
import ctypes
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

try:
    from .runtime import clear_runtime_cache, discover_runtime
except ImportError:
    from runtime import clear_runtime_cache, discover_runtime


ROOT = Path(__file__).resolve().parents[1]
COMFY_ROOT = ROOT / "ComfyUI"
CONFIG_FILE = Path(__file__).with_name("deployment_config.json")
UNIVERSAL_FINISH = ROOT / "tools" / "universal_video_finish.py"
FL2VA_RUNNER = ROOT / "tools" / "run_h3_736p_smoke.ps1"
R2V_RUNNER = ROOT / "tools" / "run_h3_r2v.ps1"
REDRAW_RUNNER = ROOT / "tools" / "run_h3_redraw_rtx.ps1"


PROFILES = {
    "rtx_creator": {
        "label": "RTX 创作者完整档",
        "badge": "推荐",
        "description": "完整 H3、SageAttention、Turbo、R2V、视频重绘和 Maxine 1.5× 终修。",
        "finish_engine": "maxine",
        "acceleration": "sageattention",
        "default_resolution": "standard",
        "h3_support": "full",
        "direct2k": True,
        "r2v": True,
        "redraw": True,
    },
    "rtx_low_vram": {
        "label": "RTX 低显存实验档",
        "badge": "实验",
        "description": "以 608×352、短时长和重度卸载优先；完整功能可以尝试，但会按运行时能力自动降级。",
        "finish_engine": "maxine",
        "acceleration": "sageattention",
        "default_resolution": "draft",
        "h3_support": "experimental",
        "direct2k": False,
        "r2v": True,
        "redraw": True,
    },
    "cuda_universal": {
        "label": "非 RTX CUDA 兼容档",
        "badge": "兼容",
        "description": "H3 生成按显存能力运行；终修改用通用时序降噪和 Lanczos 高质量缩放。",
        "finish_engine": "universal",
        "acceleration": "pytorch",
        "default_resolution": "draft",
        "h3_support": "full",
        "direct2k": False,
        "r2v": True,
        "redraw": True,
    },
    "non_nvidia_experimental": {
        "label": "AMD / Intel 实验档",
        "badge": "需适配",
        "description": "检测到 ROCm、XPU、DirectML 或兼容 PyTorch 后端时开放 H3；终修始终可回退到通用实现。",
        "finish_engine": "universal",
        "acceleration": "pytorch",
        "default_resolution": "draft",
        "h3_support": "experimental",
        "direct2k": False,
        "r2v": True,
        "redraw": True,
    },
    "media_finish_only": {
        "label": "通用媒体终修档",
        "badge": "基础",
        "description": "不承诺本地 H3 生成；提供视频降噪、音轨保留和高质量缩放。",
        "finish_engine": "universal",
        "acceleration": "cpu",
        "default_resolution": "draft",
        "h3_support": "unavailable",
        "direct2k": False,
        "r2v": False,
        "redraw": False,
    },
}


def _run(command: list[str], timeout: int = 8) -> str:
    try:
        result = subprocess.run(
            command,
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=timeout,
            check=False,
            creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0,
        )
        return result.stdout.strip() if result.returncode == 0 else ""
    except (OSError, subprocess.SubprocessError):
        return ""


def _windows_video_controllers() -> list[dict[str, Any]]:
    if os.name != "nt":
        return []
    script = (
        "Get-CimInstance Win32_VideoController | "
        "Select-Object Name,AdapterRAM,DriverVersion | ConvertTo-Json -Compress"
    )
    raw = _run(["powershell.exe", "-NoProfile", "-Command", script])
    if not raw:
        return []
    try:
        value = json.loads(raw)
        rows = value if isinstance(value, list) else [value]
        return [row for row in rows if isinstance(row, dict)]
    except json.JSONDecodeError:
        return []


def _classify_vendor(name: str) -> str:
    upper = name.upper()
    if "NVIDIA" in upper:
        return "nvidia"
    if "AMD" in upper or "RADEON" in upper:
        return "amd"
    if "INTEL" in upper or "ARC" in upper:
        return "intel"
    return "other"


def probe_hardware() -> dict[str, Any]:
    gpus: list[dict[str, Any]] = []
    query = _run([
        "nvidia-smi",
        "--query-gpu=name,memory.total,driver_version",
        "--format=csv,noheader,nounits",
    ])
    for line in query.splitlines():
        parts = [part.strip() for part in line.split(",")]
        if len(parts) >= 3:
            try:
                memory_gb = round(float(parts[1]) / 1024.0, 1)
            except ValueError:
                memory_gb = 0.0
            gpus.append({
                "name": parts[0],
                "vendor": "nvidia",
                "vram_gb": memory_gb,
                "driver": parts[2],
                "rtx": "RTX" in parts[0].upper(),
                "source": "nvidia-smi",
            })
    known_names = {gpu["name"] for gpu in gpus}
    for row in _windows_video_controllers():
        name = str(row.get("Name") or "Unknown GPU")
        if name in known_names or "MICROSOFT BASIC" in name.upper():
            continue
        adapter_ram = row.get("AdapterRAM") or 0
        try:
            memory_gb = round(int(adapter_ram) / (1024**3), 1)
        except (TypeError, ValueError):
            memory_gb = 0.0
        gpus.append({
            "name": name,
            "vendor": _classify_vendor(name),
            "vram_gb": memory_gb,
            "driver": str(row.get("DriverVersion") or ""),
            "rtx": "RTX" in name.upper(),
            "source": "system",
        })

    if os.name == "nt":
        class MemoryStatus(ctypes.Structure):
            _fields_ = [
                ("length", ctypes.c_ulong),
                ("memory_load", ctypes.c_ulong),
                ("total_physical", ctypes.c_ulonglong),
                ("available_physical", ctypes.c_ulonglong),
                ("total_page_file", ctypes.c_ulonglong),
                ("available_page_file", ctypes.c_ulonglong),
                ("total_virtual", ctypes.c_ulonglong),
                ("available_virtual", ctypes.c_ulonglong),
                ("available_extended_virtual", ctypes.c_ulonglong),
            ]
        status = MemoryStatus()
        status.length = ctypes.sizeof(status)
        ram_gb = round(status.total_physical / (1024**3), 1) if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)) else 0.0
    else:
        try:
            ram_gb = round(os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") / (1024**3), 1)
        except (AttributeError, ValueError):
            ram_gb = 0.0

    disk = shutil.disk_usage(ROOT.anchor or ROOT)
    return {
        "os": platform.platform(),
        "platform": platform.system().lower(),
        "cpu": platform.processor() or os.environ.get("PROCESSOR_IDENTIFIER", "Unknown CPU"),
        "ram_gb": ram_gb,
        "disk_free_gb": round(disk.free / (1024**3), 1),
        "disk_total_gb": round(disk.total / (1024**3), 1),
        "gpus": gpus,
        "python": sys.version.split()[0],
        "project_root": str(ROOT),
    }


def select_profile(hardware: dict[str, Any], override: dict[str, Any] | None = None) -> tuple[str, dict[str, Any]]:
    override = override or {}
    detected = hardware.get("gpus") or []
    gpu = dict(detected[0]) if detected else {"name": "CPU / 未检测到独立显卡", "vendor": "cpu", "vram_gb": 0, "rtx": False}
    if override.get("vendor"):
        gpu.update({
            "name": override.get("name") or f"手动选择：{override['vendor']}",
            "vendor": override["vendor"],
            "vram_gb": float(override.get("vram_gb") or 0),
            "rtx": override["vendor"] == "nvidia_rtx",
            "source": "manual",
        })
        if gpu["vendor"] == "nvidia_rtx":
            gpu["vendor"] = "nvidia"
    vendor = gpu.get("vendor", "cpu")
    vram = float(gpu.get("vram_gb") or 0)
    if vendor == "nvidia" and gpu.get("rtx") and vram >= 24:
        profile_id = "rtx_creator"
    elif vendor == "nvidia" and gpu.get("rtx") and vram >= 16:
        profile_id = "rtx_low_vram"
    elif vendor == "nvidia" and vram >= 24:
        profile_id = "cuda_universal"
    elif vendor in {"amd", "intel"} and vram >= 16:
        profile_id = "non_nvidia_experimental"
    else:
        profile_id = "media_finish_only"
    return profile_id, gpu


def load_configuration() -> dict[str, Any]:
    if not CONFIG_FILE.is_file():
        return {}
    try:
        value = json.loads(CONFIG_FILE.read_text(encoding="utf-8"))
        return value if isinstance(value, dict) else {}
    except (OSError, json.JSONDecodeError):
        return {}


def save_configuration(value: dict[str, Any]) -> None:
    CONFIG_FILE.write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding="utf-8")


def _runtime_module(runtime: dict[str, Any], name: str) -> bool:
    return bool((runtime.get("modules") or {}).get(name))


def _accelerated_backend(runtime: dict[str, Any]) -> bool:
    return runtime.get("backend") in {"cuda", "rocm", "xpu", "mps", "directml"}


def component_status(profile: dict[str, Any], runtime: dict[str, Any]) -> list[dict[str, Any]]:
    paths = {
        "comfyui": COMFY_ROOT / "main.py",
        "fl2va": COMFY_ROOT / "models" / "diffusion_models" / "minimax_h3_fl2va_pruned_int8_convrot.safetensors",
        "ref2va": COMFY_ROOT / "models" / "diffusion_models" / "minimax_h3_ref2va_pruned_int8_convrot.safetensors",
        "text_encoder": COMFY_ROOT / "models" / "text_encoders" / "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors",
        "video_vae": COMFY_ROOT / "models" / "vae" / "minimax_h3_video_vae_fp16.safetensors",
        "audio_vae": COMFY_ROOT / "models" / "vae" / "minimax_h3_audio_vae_fp32.safetensors",
        "turbo": COMFY_ROOT / "models" / "loras" / "minimax_h3_turbo_v4_step600_ema.safetensors",
        "concat": COMFY_ROOT / "custom_nodes" / "ComfyUI-PT_H3ConcatAVLatent",
        "rtx_nodes": COMFY_ROOT / "custom_nodes" / "Nvidia_RTX_Nodes_ComfyUI",
        "universal": UNIVERSAL_FINISH,
        "fl2va_runner": FL2VA_RUNNER,
        "r2v_runner": R2V_RUNNER,
        "redraw_runner": REDRAW_RUNNER,
    }
    sources = {
        "comfyui": "https://github.com/Comfy-Org/ComfyUI",
        "fl2va": "https://huggingface.co/Comfy-Org/MiniMax-H3/blob/main/diffusion_models/minimax_h3_fl2va_pruned_int8_convrot.safetensors",
        "ref2va": "https://huggingface.co/Comfy-Org/MiniMax-H3/blob/main/diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors",
        "text_encoder": "https://huggingface.co/Comfy-Org/MiniMax-H3/blob/main/text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors",
        "video_vae": "https://huggingface.co/Comfy-Org/MiniMax-H3/blob/main/vae/minimax_h3_video_vae_fp16.safetensors",
        "audio_vae": "https://huggingface.co/Comfy-Org/MiniMax-H3/blob/main/vae/minimax_h3_audio_vae_fp32.safetensors",
        "turbo": "https://huggingface.co/larryvrh/MiniMax-H3-Turbo-Lora",
        "concat": "https://github.com/ptmaster/ComfyUI-PT_H3ConcatAVLatent",
        "rtx_nodes": "https://developer.nvidia.com/maxine",
        "universal": "",
        "fl2va_runner": "",
        "r2v_runner": "",
        "redraw_runner": "",
    }
    specs = [
        ("comfyui", "ComfyUI 本地引擎", "core", profile["h3_support"] != "unavailable"),
        ("fl2va", "MiniMax H3 FL2VA INT8", "model", profile["h3_support"] != "unavailable"),
        ("ref2va", "MiniMax H3 Ref2VA INT8", "model", False),
        ("text_encoder", "Qwen3-VL H3 文本编码器", "model", profile["h3_support"] != "unavailable"),
        ("video_vae", "H3 Video VAE", "model", profile["h3_support"] != "unavailable"),
        ("audio_vae", "H3 Audio VAE", "model", profile["h3_support"] != "unavailable"),
        ("turbo", "MiniMax H3 Turbo LoRA", "model", profile["h3_support"] != "unavailable"),
        ("concat", "PT H3 Concat AV Latent", "plugin", False),
        ("rtx_nodes", "NVIDIA RTX / Maxine 节点", "plugin", False),
        ("universal", "通用时序降噪与 Lanczos 终修", "plugin", profile["finish_engine"] == "universal"),
        ("fl2va_runner", "H3 FL2VA 执行器", "backend", profile["h3_support"] != "unavailable"),
        ("r2v_runner", "H3 R2V 执行器", "backend", False),
        ("redraw_runner", "H3 视频重绘执行器", "backend", False),
    ]
    result = []
    for component_id, label, kind, required in specs:
        path = paths[component_id]
        installed = path.exists()
        size = path.stat().st_size if path.is_file() else 0
        result.append({
            "id": component_id,
            "label": label,
            "kind": kind,
            "required": required,
            "installed": installed,
            "path": str(path),
            "size": size,
            "source_url": sources[component_id],
        })
    result.extend([
        {
            "id": "python_runtime",
            "label": "独立或现有 Python 运行时",
            "kind": "runtime",
            "required": True,
            "installed": bool(runtime.get("usable")),
            "path": runtime.get("executable") or "未找到兼容 Python",
            "size": 0,
            "source_url": "",
        },
        {
            "id": "torch_backend",
            "label": "PyTorch H3 计算后端",
            "kind": "runtime",
            "required": profile["h3_support"] != "unavailable",
            "installed": _runtime_module(runtime, "torch") and _accelerated_backend(runtime),
            "path": f"{runtime.get('backend', 'none')} · {runtime.get('device_name', 'Unknown device')}",
            "size": 0,
            "source_url": "https://pytorch.org/get-started/locally/",
        },
        {
            "id": "sageattention",
            "label": "SageAttention CUDA 加速",
            "kind": "python",
            "required": False,
            "installed": _runtime_module(runtime, "sageattention") and _runtime_module(runtime, "triton"),
            "path": "Python package: sageattention",
            "size": 0,
            "source_url": "https://github.com/thu-ml/SageAttention",
        },
        {
            "id": "nvvfx",
            "label": "NVIDIA Maxine VFX Runtime",
            "kind": "runtime",
            "required": False,
            "installed": _runtime_module(runtime, "nvvfx") and runtime.get("backend") == "cuda",
            "path": "Python package/runtime: nvvfx",
            "size": 0,
            "source_url": "https://developer.nvidia.com/maxine",
        },
        {
            "id": "modelscope",
            "label": "ModelScope 国内模型下载",
            "kind": "python",
            "required": False,
            "installed": _runtime_module(runtime, "modelscope"),
            "path": "Python package: modelscope",
            "size": 0,
            "source_url": "https://www.modelscope.cn/",
        },
        {
            "id": "opencv",
            "label": "OpenCV 可选视频节点",
            "kind": "python",
            "required": False,
            "installed": _runtime_module(runtime, "cv2"),
            "path": "Python package: opencv-python",
            "size": 0,
            "source_url": "https://pypi.org/project/opencv-python/",
        },
    ])
    return result


def build_plan(
    hardware: dict[str, Any],
    override: dict[str, Any] | None = None,
    selected_profile: str | None = None,
    runtime: dict[str, Any] | None = None,
) -> dict[str, Any]:
    recommended_id, gpu = select_profile(hardware, override)
    profile_id = selected_profile if selected_profile in PROFILES else recommended_id
    profile = {"id": profile_id, **PROFILES[profile_id]}
    desired = dict(profile)
    runtime = runtime or discover_runtime()
    modules = runtime.get("modules") or {}
    backend = str(runtime.get("backend") or "none")
    fallback_notes: list[str] = []

    h3_requested = profile["h3_support"] != "unavailable"
    h3_runtime_ready = bool(
        runtime.get("usable")
        and modules.get("torch")
        and backend in {"cuda", "rocm", "xpu", "mps", "directml"}
    )
    if h3_requested and not h3_runtime_ready:
        profile["h3_support"] = "unavailable"
        fallback_notes.append("所选 Python 尚未提供可用的 GPU PyTorch 后端，本地 H3 已暂时关闭。")
    elif h3_requested and backend in {"rocm", "xpu", "mps", "directml"}:
        profile["h3_support"] = "experimental"

    sage_ready = bool(backend == "cuda" and modules.get("sageattention") and modules.get("triton"))
    if profile["acceleration"] == "sageattention" and not sage_ready:
        profile["acceleration"] = "pytorch"
        fallback_notes.append("SageAttention 与当前运行时不匹配，已回退到 PyTorch Attention。")
    if profile["h3_support"] == "unavailable":
        profile["acceleration"] = "cpu"

    maxine_ready = bool(
        backend == "cuda"
        and gpu.get("rtx")
        and modules.get("nvvfx")
        and (COMFY_ROOT / "custom_nodes" / "Nvidia_RTX_Nodes_ComfyUI").exists()
    )
    if profile["finish_engine"] == "maxine" and not maxine_ready:
        profile["finish_engine"] = "universal"
        fallback_notes.append("NVIDIA VFX 与当前显卡或运行时不匹配，已切换到通用终修。")

    if fallback_notes:
        profile["description"] += " " + " ".join(fallback_notes)

    components = component_status(profile, runtime)
    status = {component["id"]: component for component in components}
    required = [component for component in components if component["required"]]
    installed_required = [component for component in required if component["installed"]]
    core_ready = len(installed_required) == len(required)

    if profile["h3_support"] != "unavailable" and not core_ready:
        component_note = "H3 核心组件尚未齐全，生成能力已关闭，但可用终修后端保持开放。"
        profile["h3_support"] = "unavailable"
        profile["acceleration"] = "cpu"
        profile["description"] += " " + component_note
        fallback_notes.append(component_note)

    h3_enabled = bool(profile["h3_support"] != "unavailable" and core_ready)
    profile["r2v"] = bool(
        desired["r2v"]
        and h3_enabled
        and status["ref2va"]["installed"]
        and status["r2v_runner"]["installed"]
    )
    profile["redraw"] = bool(
        desired["redraw"]
        and h3_enabled
        and status["concat"]["installed"]
        and status["redraw_runner"]["installed"]
    )
    profile["direct2k"] = bool(
        desired["direct2k"]
        and h3_enabled
        and backend == "cuda"
        and float(gpu.get("vram_gb") or 0) >= 23.5
    )
    profile["h3"] = h3_enabled

    warnings: list[str] = []
    warnings.extend(fallback_notes)
    if profile["h3_support"] == "experimental":
        warnings.append("当前硬件上的 H3 生成属于实验路径，建议先运行 608×352 / 0.5 秒检查。")
    if profile["h3_support"] == "unavailable":
        warnings.append("当前配置不建议运行本地 H3；画布仍可用于通用视频终修或连接远程 H3 引擎。")
    if float(hardware.get("ram_gb") or 0) < 48:
        warnings.append("系统内存低于 48GB，H3 模型卸载和 CPU offload 可能不稳定。")
    if float(hardware.get("disk_free_gb") or 0) < 100:
        warnings.append("建议至少保留 100GB 可用空间，用于约 60GB 模型、缓存和视频输出。")
    if profile["finish_engine"] == "universal":
        warnings.append("非 RTX 终修不会调用 Maxine；采用保结构的通用降噪与 Lanczos 缩放。")
    if runtime.get("is_virtualenv"):
        warnings.append("当前使用的是开发型虚拟环境；生产部署不依赖它，建议安装到项目内 .runtime/python。")
    if not modules.get("modelscope"):
        warnings.append("未安装 ModelScope；不影响已下载模型的推理，但国内模型自动下载不可用。")
    if desired["r2v"] and not profile["r2v"] and h3_enabled:
        warnings.append("Ref2VA 模型尚未就位，多参考生成已单独禁用，其他 H3 模式仍可运行。")
    if desired["redraw"] and not profile["redraw"] and h3_enabled:
        warnings.append("Concat AV 节点尚未就位，视频重绘已单独禁用，其他 H3 模式仍可运行。")

    runtime_summary = {
        **runtime,
        "supports_h3": h3_runtime_ready,
        "supports_sageattention": sage_ready,
        "supports_maxine": maxine_ready,
    }
    return {
        "hardware": hardware,
        "runtime": runtime_summary,
        "selected_gpu": gpu,
        "recommended_profile_id": recommended_id,
        "profile": profile,
        "profiles": [{"id": key, **value} for key, value in PROFILES.items()],
        "components": components,
        "readiness": {
            "installed": len(installed_required),
            "required": len(required),
            "percent": round(100 * len(installed_required) / max(len(required), 1)),
            "ready": core_ready,
        },
        "capabilities": {
            "h3": h3_enabled,
            "t2v": h3_enabled,
            "i2v": h3_enabled,
            "l2v": h3_enabled,
            "fl2v": h3_enabled,
            "r2v": profile["r2v"],
            "redraw": profile["redraw"],
            "direct2k": profile["direct2k"],
            "upscale": bool(runtime.get("usable") and UNIVERSAL_FINISH.is_file()),
            "maxine": maxine_ready,
        },
        "recommended_settings": {
            "resolution": profile["default_resolution"],
            "preview_duration": 0.5,
            "final_steps": 8 if profile["h3_support"] == "full" else 4,
            "finish_engine": profile["finish_engine"],
            "parallel_gpu_jobs": 1,
            "safe_duration": 10 if profile["h3_support"] == "full" else 0.5,
        },
        "warnings": warnings,
    }


def current_deployment(refresh_runtime: bool = False) -> dict[str, Any]:
    hardware = probe_hardware()
    config = load_configuration()
    if refresh_runtime:
        clear_runtime_cache()
    runtime = discover_runtime(str(config.get("runtime_python") or "") or None)
    return build_plan(
        hardware,
        override=config.get("hardware_override"),
        selected_profile=config.get("profile_id"),
        runtime=runtime,
    ) | {"configuration": config}


def apply_plan(payload: dict[str, Any]) -> dict[str, Any]:
    hardware = probe_hardware()
    override = payload.get("hardware_override") if isinstance(payload.get("hardware_override"), dict) else None
    requested = str(payload.get("profile_id") or "")
    existing = load_configuration()
    runtime_python = str(payload.get("runtime_python") or existing.get("runtime_python") or "").strip()
    clear_runtime_cache()
    runtime = discover_runtime(runtime_python or None)
    plan = build_plan(hardware, override=override, selected_profile=requested, runtime=runtime)
    config = {
        "version": 2,
        "profile_id": plan["profile"]["id"],
        "hardware_override": override or {},
        "default_resolution": plan["profile"]["default_resolution"],
        "configured_at": datetime.now(timezone.utc).isoformat(),
    }
    if runtime_python:
        config["runtime_python"] = runtime_python
    save_configuration(config)
    return plan | {"configuration": config}
