from __future__ import annotations

import json
import shutil
import subprocess
import threading
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Any

import fcntl

from ...config import settings

READINESS_CACHE_SECONDS = 300
READINESS_FAILURE_CACHE_SECONDS = 30
_readiness_cache: tuple[float, dict[str, Any]] | None = None
_readiness_thread_lock = threading.Lock()


@contextmanager
def readiness_probe_lock():
    settings.project_root.joinpath("var").mkdir(parents=True, exist_ok=True)
    lock_path = settings.project_root / "var" / "runtime_readiness.lock"
    with _readiness_thread_lock, lock_path.open("w") as lock_file:
        fcntl.flock(lock_file, fcntl.LOCK_EX)
        try:
            yield
        finally:
            fcntl.flock(lock_file, fcntl.LOCK_UN)


def runtime_environment_specs() -> list[dict[str, Any]]:
    return [
        {
            "role": "backend_task",
            "name": settings.task_conda_env,
            "label": "Backend, SegModel, YOLO",
            "env_file": "envs/seg_smp.yml",
            "required_imports": [
                {"module": "fastapi", "package": "fastapi"},
                {"module": "uvicorn", "package": "uvicorn"},
                {"module": "torch", "package": "torch"},
                {"module": "cv2", "package": "opencv-python"},
                {"module": "segmentation_models_pytorch", "package": "segmentation-models-pytorch"},
                {"module": "ultralytics", "package": "ultralytics"},
                {"module": "albumentations", "package": "albumentations"},
                {"module": "mmengine", "package": "mmengine"},
                {"module": "mmseg", "package": "mmsegmentation"},
            ],
        },
        {
            "role": "mmseg_full",
            "name": settings.mmseg_conda_env,
            "label": "MMSeg full mmcv runtime",
            "env_file": "envs/seg_mmcv.yml",
            "required_imports": [
                {"module": "torch", "package": "torch"},
                {"module": "cv2", "package": "opencv-python"},
                {"module": "mmcv", "package": "mmcv"},
                {"module": "mmcv._ext", "package": "mmcv"},
                {"module": "mmengine", "package": "mmengine"},
                {"module": "mmseg", "package": "mmsegmentation"},
            ],
        },
    ]


def parse_nvidia_smi_csv(output: str) -> list[dict]:
    gpus: list[dict] = []
    for line in output.splitlines():
        if not line.strip():
            continue
        parts = [part.strip() for part in line.split(",")]
        if len(parts) < 7:
            continue
        index, name, total, used, free, util, temp = parts[:7]
        try:
            gpus.append(
                {
                    "index": int(index),
                    "name": name,
                    "memory_total_mb": int(total),
                    "memory_used_mb": int(used),
                    "memory_free_mb": int(free),
                    "utilization_gpu_percent": int(util),
                    "temperature_c": int(temp),
                }
            )
        except ValueError:
            continue
    return gpus


def get_gpus() -> dict:
    cmd = [
        "nvidia-smi",
        "--query-gpu=index,name,memory.total,memory.used,memory.free,utilization.gpu,temperature.gpu",
        "--format=csv,noheader,nounits",
    ]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return {"available": True, "gpus": parse_nvidia_smi_csv(result.stdout)}
    except Exception as exc:
        return {"available": False, "gpus": [], "error": str(exc)}


def get_conda_envs() -> dict:
    try:
        result = subprocess.run(["conda", "env", "list"], capture_output=True, text=True, check=True)
    except Exception as exc:
        return {"available": False, "envs": [], "error": str(exc)}
    envs = []
    for line in result.stdout.splitlines():
        raw = line.strip()
        if not raw or raw.startswith("#"):
            continue
        marker = "*" in raw.split()
        parts = raw.replace("*", " ").split()
        if len(parts) >= 2:
            envs.append({"name": parts[0], "path": parts[-1], "active": marker})
    return {"available": True, "envs": envs, "task_default": settings.task_conda_env, "mmseg_default": settings.mmseg_conda_env}


def probe_code(required_imports: list[dict[str, str]]) -> str:
    imports_json = json.dumps(required_imports, ensure_ascii=False)
    return f"""
import importlib
import importlib.metadata as metadata
import json
import platform
import sys

required = {imports_json}
checks = []
for item in required:
    module_name = item["module"]
    package_name = item.get("package") or module_name
    try:
        module = importlib.import_module(module_name)
        version = getattr(module, "__version__", None)
        if version is None:
            try:
                version = metadata.version(package_name)
            except Exception:
                version = None
        checks.append({{
            "module": module_name,
            "package": package_name,
            "passed": True,
            "version": str(version) if version is not None else None,
        }})
    except Exception as exc:
        checks.append({{
            "module": module_name,
            "package": package_name,
            "passed": False,
            "error": f"{{type(exc).__name__}}: {{exc}}",
        }})

extra = {{"python": sys.version.split()[0], "platform": platform.platform()}}
try:
    import torch
    extra["torch_cuda_available"] = bool(torch.cuda.is_available())
    extra["torch_cuda"] = getattr(torch.version, "cuda", None)
except Exception:
    pass
print(json.dumps({{"checks": checks, "extra": extra}}, ensure_ascii=False))
"""


def parse_probe_stdout(stdout: str) -> dict[str, Any]:
    for line in reversed(stdout.splitlines()):
        text = line.strip()
        if not text or not text.startswith("{"):
            continue
        return json.loads(text)
    raise ValueError("probe did not emit JSON")


def inspect_conda_env(env_name: str, required_imports: list[dict[str, str]], timeout: int = 45, retries: int = 1) -> dict[str, Any]:
    command = ["conda", "run", "-n", env_name, "python", "-c", probe_code(required_imports)]
    attempts = []
    result: subprocess.CompletedProcess[str] | None = None
    for attempt in range(retries + 1):
        result = subprocess.run(command, capture_output=True, text=True, timeout=timeout)
        attempts.append({"attempt": attempt + 1, "returncode": result.returncode})
        if result.returncode == 0 or result.returncode not in {139, -11} or attempt >= retries:
            break
        time.sleep(2)
    if result is None:
        raise RuntimeError("probe did not run")
    detail = {
        "command": command[:4] + ["..."],
        "returncode": result.returncode,
        "attempts": attempts,
        "stdout_tail": result.stdout[-2000:],
        "stderr_tail": result.stderr[-2000:],
    }
    if result.returncode != 0:
        return {"passed": False, "checks": [], "extra": {}, "detail": detail}
    try:
        parsed = parse_probe_stdout(result.stdout)
    except Exception as exc:
        return {"passed": False, "checks": [], "extra": {}, "detail": {**detail, "error": str(exc)}}
    checks = parsed.get("checks", [])
    return {
        "passed": bool(checks) and all(item.get("passed") for item in checks),
        "checks": checks,
        "extra": parsed.get("extra", {}),
        "detail": detail,
    }


def get_runtime_readiness(force: bool = False) -> dict[str, Any]:
    global _readiness_cache
    now = time.time()
    if not force and _readiness_cache:
        cache_ttl = READINESS_CACHE_SECONDS if _readiness_cache[1].get("passed") else READINESS_FAILURE_CACHE_SECONDS
        if now - _readiness_cache[0] < cache_ttl:
            cached = dict(_readiness_cache[1])
            cached["cached"] = True
            return cached

    with readiness_probe_lock():
        now = time.time()
        if not force and _readiness_cache:
            cache_ttl = READINESS_CACHE_SECONDS if _readiness_cache[1].get("passed") else READINESS_FAILURE_CACHE_SECONDS
            if now - _readiness_cache[0] < cache_ttl:
                cached = dict(_readiness_cache[1])
                cached["cached"] = True
                return cached

        conda = get_conda_envs()
        env_paths = {item["name"]: item["path"] for item in conda.get("envs", [])}
        envs: list[dict[str, Any]] = []
        for spec in runtime_environment_specs():
            env_name = spec["name"]
            exists = env_name in env_paths
            env_report: dict[str, Any] = {
                "role": spec["role"],
                "name": env_name,
                "label": spec["label"],
                "env_file": spec["env_file"],
                "path": env_paths.get(env_name),
                "exists": exists,
                "passed": False,
                "checks": [],
                "extra": {},
            }
            if exists:
                probe = inspect_conda_env(env_name, spec["required_imports"])
                env_report.update(probe)
            envs.append(env_report)

        payload = {
            "available": bool(conda.get("available")),
            "passed": bool(conda.get("available")) and all(item["passed"] for item in envs),
            "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(now)),
            "cache_seconds": READINESS_CACHE_SECONDS,
            "failure_cache_seconds": READINESS_FAILURE_CACHE_SECONDS,
            "cached": False,
            "envs": envs,
            "specs": {
                "bootstrap_script": "scripts/bootstrap_conda_envs.sh",
                "verify_script": "scripts/verify_runtime_envs.py",
                "env_files": [spec["env_file"] for spec in runtime_environment_specs()],
                "task_default": settings.task_conda_env,
                "mmseg_default": settings.mmseg_conda_env,
            },
        }
        _readiness_cache = (now, payload)
        return payload


def disk_usage() -> dict:
    usage = shutil.disk_usage(settings.source_root)
    return {
        "path": str(settings.source_root),
        "total": usage.total,
        "used": usage.used,
        "free": usage.free,
    }


def scan_results() -> list[dict]:
    roots = [
        settings.source_root / "DataSet_Public_outputs",
        settings.source_root / "BestMode_Predict_Results_DataSet_Public",
        settings.source_root / "Hardisk",
        settings.source_root / "Seg_All_In_One_Analysis",
    ]
    exts = {".csv", ".png", ".jpg", ".jpeg", ".svg", ".log", ".pth", ".pt"}
    results: list[dict] = []
    for root in roots:
        if not root.exists():
            continue
        for path in root.rglob("*"):
            if path.is_file() and path.suffix.lower() in exts:
                try:
                    stat = path.stat()
                    results.append(
                        {
                            "name": path.name,
                            "path": str(path.resolve()),
                            "relative_path": str(path.resolve().relative_to(settings.source_root)),
                            "size": stat.st_size,
                            "modified": stat.st_mtime,
                            "kind": path.suffix.lower().lstrip("."),
                        }
                    )
                except OSError:
                    continue
    results.sort(key=lambda item: item["modified"], reverse=True)
    return results[:1000]