"""Environment preflight checks for conversion and golden export."""

from __future__ import annotations

import importlib.util
import json
import subprocess
import sys
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Iterable


@dataclass(frozen=True)
class PreflightResult:
    name: str
    status: str
    detail: str


def _module_available(module_name: str) -> bool:
    return importlib.util.find_spec(module_name) is not None


def check_python_modules(modules: Iterable[str]) -> list[PreflightResult]:
    results: list[PreflightResult] = []
    for module_name in modules:
        if _module_available(module_name):
            results.append(PreflightResult(module_name, "ok", "module import spec found"))
        else:
            results.append(PreflightResult(module_name, "missing", "module import spec not found"))
    return results


def check_mlx_import(timeout_seconds: int = 20) -> PreflightResult:
    code = (
        "import mlx.core as mx; "
        "print({'default_device': str(mx.default_device()), "
        "'has_metal': hasattr(mx, 'metal')})"
    )
    try:
        completed = subprocess.run(
            [sys.executable, "-c", code],
            check=False,
            capture_output=True,
            text=True,
            timeout=timeout_seconds,
        )
    except subprocess.TimeoutExpired:
        return PreflightResult("mlx_import", "failed", f"timed out after {timeout_seconds}s")
    detail = (completed.stdout + completed.stderr).strip()
    if completed.returncode == 0:
        return PreflightResult("mlx_import", "ok", detail)
    return PreflightResult("mlx_import", "failed", detail)


def check_paths(paths: Iterable[Path]) -> list[PreflightResult]:
    results: list[PreflightResult] = []
    for path in paths:
        if path.exists():
            results.append(PreflightResult(str(path), "ok", "path exists"))
        else:
            results.append(PreflightResult(str(path), "missing", "path does not exist"))
    return results


def run_preflight(include_mlx_import: bool = True) -> list[PreflightResult]:
    results = check_python_modules(["mlx", "transformers", "PIL"])
    if include_mlx_import:
        results.append(check_mlx_import())
    results.extend(
        check_paths(
            [
                Path("split_files/text_encoders/qwen_3_06b_base-mlx.safetensors"),
                Path("split_files/diffusion_models/anima-base-v1.0-mlx.safetensors"),
                Path("split_files/vae/qwen_image_vae-mlx.safetensors"),
                Path("tokenizers/qwen25_tokenizer"),
                Path("tokenizers/t5_tokenizer"),
            ]
        )
    )
    return results


def write_preflight_report(path: Path, results: list[PreflightResult]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps([asdict(result) for result in results], indent=2) + "\n", encoding="utf-8")