"""Discover and inspect a usable H3 Python runtime without requiring a venv.

The production layout is ``.runtime/python``.  Environment overrides and the
interpreter already running H3 Studio are supported so advanced users can use
an existing Conda, system Python, ROCm, XPU, or DirectML installation.
"""

from __future__ import annotations

import json
import os
import shutil
import subprocess
import sys
from functools import lru_cache
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]

_PROBE = r'''
import importlib.util
import json
import platform
import sys

module_names = (
    "aiohttp", "av", "numpy", "PIL", "safetensors", "torch",
    "torchvision", "torchaudio", "sageattention", "triton", "nvvfx",
    "modelscope", "cv2", "torch_directml",
)
modules = {name: importlib.util.find_spec(name) is not None for name in module_names}
result = {
    "version": platform.python_version(),
    "implementation": platform.python_implementation(),
    "prefix": sys.prefix,
    "base_prefix": sys.base_prefix,
    "is_virtualenv": sys.prefix != sys.base_prefix,
    "modules": modules,
    "backend": "cpu",
    "device_name": "CPU",
    "device_count": 0,
    "torch_version": None,
    "torch_runtime": None,
}
if modules["torch"]:
    try:
        import torch
        result["torch_version"] = str(torch.__version__)
        hip = getattr(torch.version, "hip", None)
        cuda = getattr(torch.version, "cuda", None)
        if torch.cuda.is_available():
            result["backend"] = "rocm" if hip else "cuda"
            result["torch_runtime"] = str(hip or cuda or "")
            result["device_count"] = int(torch.cuda.device_count())
            result["device_name"] = str(torch.cuda.get_device_name(0))
        elif hasattr(torch, "xpu") and torch.xpu.is_available():
            result["backend"] = "xpu"
            result["device_count"] = int(torch.xpu.device_count())
            result["device_name"] = str(torch.xpu.get_device_name(0))
        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            result["backend"] = "mps"
            result["device_count"] = 1
            result["device_name"] = "Apple Metal"
        elif modules["torch_directml"]:
            result["backend"] = "directml"
            result["device_count"] = 1
            result["device_name"] = "DirectML device"
    except Exception as error:
        result["torch_error"] = f"{type(error).__name__}: {error}"
print("H3_RUNTIME_JSON=" + json.dumps(result, ensure_ascii=False))
'''


def _candidate_paths(preferred: str | None = None) -> list[tuple[str, str]]:
    values: list[tuple[str, str]] = []
    if preferred:
        values.append((preferred, "configuration"))
    for name in ("H3_ENGINE_PYTHON", "H3_PYTHON", "H3_STUDIO_PYTHON"):
        value = os.environ.get(name)
        if value:
            values.append((value, f"environment:{name}"))

    if os.name == "nt":
        values.append((str(ROOT / ".runtime" / "python" / "python.exe"), "portable"))
    else:
        values.extend(
            [
                (str(ROOT / ".runtime" / "python" / "bin" / "python3"), "portable"),
                (str(ROOT / ".runtime" / "python" / "bin" / "python"), "portable"),
            ]
        )

    # The current interpreter may be Conda, system Python, a portable runtime,
    # or a developer environment.  It is a fallback, not a venv requirement.
    values.append((sys.executable, "current"))
    for command in ("python3", "python"):
        resolved = shutil.which(command)
        if resolved:
            values.append((resolved, "system"))

    result: list[tuple[str, str]] = []
    seen: set[str] = set()
    for raw, source in values:
        cleaned = str(raw).strip().strip('"')
        if not cleaned:
            continue
        resolved = shutil.which(cleaned) or cleaned
        path = Path(resolved).expanduser()
        try:
            key = os.path.normcase(str(path.resolve()))
        except OSError:
            key = os.path.normcase(str(path))
        if key in seen:
            continue
        seen.add(key)
        result.append((str(path), source))
    return result


@lru_cache(maxsize=16)
def probe_python(executable: str, source: str = "unknown") -> dict[str, Any]:
    path = Path(executable)
    if not path.is_file():
        return {
            "usable": False,
            "executable": executable,
            "command": [executable],
            "source": source,
            "reason": "Python executable not found",
            "modules": {},
            "backend": "none",
        }
    try:
        completed = subprocess.run(
            [str(path), "-I", "-c", _PROBE],
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=30,
            check=False,
            creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0,
        )
    except (OSError, subprocess.SubprocessError) as error:
        return {
            "usable": False,
            "executable": str(path),
            "command": [str(path)],
            "source": source,
            "reason": f"Runtime probe failed: {error}",
            "modules": {},
            "backend": "none",
        }
    marker = "H3_RUNTIME_JSON="
    line = next((item for item in reversed(completed.stdout.splitlines()) if item.startswith(marker)), "")
    if completed.returncode != 0 or not line:
        detail = (completed.stderr or completed.stdout).strip()[-1200:]
        return {
            "usable": False,
            "executable": str(path),
            "command": [str(path)],
            "source": source,
            "reason": detail or f"Runtime exited with code {completed.returncode}",
            "modules": {},
            "backend": "none",
        }
    try:
        result = json.loads(line[len(marker):])
    except json.JSONDecodeError as error:
        return {
            "usable": False,
            "executable": str(path),
            "command": [str(path)],
            "source": source,
            "reason": f"Invalid runtime report: {error}",
            "modules": {},
            "backend": "none",
        }
    modules = result.get("modules") or {}
    result.update(
        {
            "usable": bool(modules.get("aiohttp") and modules.get("av") and modules.get("numpy")),
            "executable": str(path.resolve()),
            "command": [str(path.resolve())],
            "source": source,
        }
    )
    if not result["usable"]:
        missing = [name for name in ("aiohttp", "av", "numpy") if not modules.get(name)]
        result["reason"] = "Missing Studio runtime modules: " + ", ".join(missing)
    else:
        result["reason"] = ""
    return result


def discover_runtime(preferred: str | None = None) -> dict[str, Any]:
    attempts: list[dict[str, Any]] = []
    for executable, source in _candidate_paths(preferred):
        report = dict(probe_python(executable, source))
        attempts.append(
            {
                "executable": report.get("executable", executable),
                "source": source,
                "usable": bool(report.get("usable")),
                "reason": report.get("reason", ""),
            }
        )
        if report.get("usable"):
            report["attempts"] = attempts
            return report
    return {
        "usable": False,
        "executable": "",
        "command": [],
        "source": "none",
        "version": "",
        "backend": "none",
        "device_name": "No compatible runtime",
        "device_count": 0,
        "modules": {},
        "reason": "No Python runtime with aiohttp, av, and numpy was found",
        "attempts": attempts,
    }


def clear_runtime_cache() -> None:
    """Refresh probes after the installer or deployment wizard changes packages."""

    probe_python.cache_clear()

