"""Pre-flight validation checks for bilge doctor."""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Literal

from bilgeai.core.constants import CONFIG_FILE, PROJECT_DIR


@dataclass
class CheckResult:
    """Result of a single pre-flight check."""

    name: str
    status: Literal["pass", "warn", "fail"]
    message: str
    suggestion: str | None = None


@dataclass
class PreflightReport:
    """Complete pre-flight validation report."""

    checks: list[CheckResult] = field(default_factory=list)

    @property
    def passed(self) -> bool:
        return all(c.status != "fail" for c in self.checks)

    @property
    def num_pass(self) -> int:
        return sum(1 for c in self.checks if c.status == "pass")

    @property
    def num_warn(self) -> int:
        return sum(1 for c in self.checks if c.status == "warn")

    @property
    def num_fail(self) -> int:
        return sum(1 for c in self.checks if c.status == "fail")


def run_preflight_checks(base_path: Path | None = None) -> PreflightReport:
    """Run all pre-flight validation checks."""
    base = base_path or Path.cwd()
    report = PreflightReport()

    report.checks.append(_check_config_exists(base))
    report.checks.append(_check_project_dir(base))
    report.checks.append(_check_hardware_profile(base))
    report.checks.append(_check_dataset_exists(base))
    report.checks.append(_check_disk_space(base))
    report.checks.append(_check_gpu_available())
    report.checks.append(_check_python_version())
    report.checks.append(_check_dependencies())
    report.checks.append(_check_llm_api_key(base))

    return report


def _check_config_exists(base: Path) -> CheckResult:
    config_path = base / CONFIG_FILE
    if config_path.exists():
        return CheckResult("config", "pass", f"{CONFIG_FILE} found")
    return CheckResult(
        "config",
        "fail",
        f"{CONFIG_FILE} not found",
        suggestion="Run 'bilge init' to create a configuration file.",
    )


def _check_project_dir(base: Path) -> CheckResult:
    project_dir = base / PROJECT_DIR
    if project_dir.exists():
        return CheckResult("project_dir", "pass", f"{PROJECT_DIR}/ directory exists")
    return CheckResult(
        "project_dir",
        "warn",
        f"{PROJECT_DIR}/ directory not found",
        suggestion="Run 'bilge init' to initialize the project.",
    )


def _check_hardware_profile(base: Path) -> CheckResult:
    profile_path = base / PROJECT_DIR / "hardware_profile.json"
    if profile_path.exists():
        return CheckResult("hardware_profile", "pass", "Hardware profile cached")
    return CheckResult(
        "hardware_profile",
        "warn",
        "No cached hardware profile",
        suggestion="Run 'bilge profile' to analyze your hardware.",
    )


def _check_dataset_exists(base: Path) -> CheckResult:
    config_path = base / CONFIG_FILE
    if not config_path.exists():
        return CheckResult("dataset", "warn", "Cannot check dataset (no config)")

    try:
        from bilgeai.core.config import BilgeConfig

        config = BilgeConfig.load(config_path)
        data_path = Path(config.data.path)
        if not data_path.is_absolute():
            data_path = base / data_path

        if data_path.exists():
            return CheckResult("dataset", "pass", f"Dataset found: {data_path.name}")
        return CheckResult(
            "dataset",
            "fail",
            f"Dataset not found: {config.data.path}",
            suggestion="Check data.path in bilge.yaml or provide correct path.",
        )
    except Exception as e:
        return CheckResult("dataset", "warn", f"Could not check dataset: {e}")


def _check_disk_space(base: Path) -> CheckResult:
    import psutil

    try:
        usage = psutil.disk_usage(str(base))
        free_gb = usage.free / (1024**3)
        if free_gb > 20:
            return CheckResult("disk", "pass", f"{free_gb:.0f} GB free disk space")
        elif free_gb > 5:
            return CheckResult(
                "disk",
                "warn",
                f"Only {free_gb:.1f} GB free disk space",
                suggestion="Training may need 5-20 GB for checkpoints. Free up disk space.",
            )
        else:
            return CheckResult(
                "disk",
                "fail",
                f"Only {free_gb:.1f} GB free disk space",
                suggestion="Not enough space for training. Need at least 5 GB free.",
            )
    except Exception:
        return CheckResult("disk", "warn", "Could not check disk space")


def _check_gpu_available() -> CheckResult:
    try:
        from bilgeai.profiler.gpu import detect_gpus

        gpus = detect_gpus()
        if gpus:
            names = ", ".join(g.name for g in gpus)
            return CheckResult("gpu", "pass", f"GPU detected: {names}")
        return CheckResult(
            "gpu",
            "warn",
            "No GPU detected",
            suggestion="Training will be very slow on CPU. Consider using a GPU machine.",
        )
    except Exception:
        return CheckResult("gpu", "warn", "Could not detect GPU")


def _check_python_version() -> CheckResult:
    import sys

    version = f"{sys.version_info.major}.{sys.version_info.minor}"
    if sys.version_info >= (3, 10):  # noqa: UP036
        return CheckResult("python", "pass", f"Python {version}")
    return CheckResult(
        "python",
        "fail",
        f"Python {version} (requires 3.10+)",
        suggestion="Upgrade to Python 3.10 or later.",
    )


def _check_dependencies() -> CheckResult:
    missing = []
    for pkg in ["typer", "rich", "pydantic", "psutil", "yaml", "jinja2"]:
        try:
            __import__(pkg)
        except ImportError:
            missing.append(pkg)

    if not missing:
        return CheckResult("dependencies", "pass", "Core dependencies installed")
    return CheckResult(
        "dependencies",
        "fail",
        f"Missing: {', '.join(missing)}",
        suggestion="Run 'pip install bilgeai' to install dependencies.",
    )


def _check_llm_api_key(base: Path) -> CheckResult:
    """Check if LLM API key is configured and valid."""
    import os

    config_path = base / CONFIG_FILE
    if not config_path.exists():
        return CheckResult(
            "llm_api_key",
            "warn",
            "Cannot check API key (no config)",
            suggestion="Run 'bilge init' to configure your LLM provider.",
        )

    try:
        from bilgeai.core.config import BilgeConfig

        config = BilgeConfig.load(config_path)
        provider = config.llm.provider

        if provider == "local":
            return CheckResult("llm_api_key", "pass", f"Local provider ({config.llm.base_url})")

        api_key_env = config.llm.api_key_env
        if not api_key_env:
            return CheckResult(
                "llm_api_key",
                "fail",
                "No API key env var configured",
                suggestion="Run 'bilge init' to configure your LLM provider.",
            )

        api_key = os.environ.get(api_key_env)
        if not api_key:
            return CheckResult(
                "llm_api_key",
                "fail",
                f"{api_key_env} not set",
                suggestion=f"export {api_key_env}=your-api-key-here",
            )

        # Try to validate the key
        try:
            from bilgeai.llm.router import get_llm_client

            client = get_llm_client(
                provider=provider,
                model=config.llm.model,
                api_key_env=api_key_env,
                base_url=config.llm.base_url,
            )
            if client.validate_api_key():
                return CheckResult(
                    "llm_api_key",
                    "pass",
                    f"{provider} API key valid ({config.llm.model})",
                )
            return CheckResult(
                "llm_api_key",
                "warn",
                f"{api_key_env} set but validation failed",
                suggestion="Check your API key or network connection.",
            )
        except Exception:
            return CheckResult(
                "llm_api_key",
                "warn",
                f"{api_key_env} set (could not validate)",
                suggestion="API key found but could not verify. Check your connection.",
            )

    except Exception as e:
        return CheckResult("llm_api_key", "warn", f"Could not check API key: {e}")
