"""Library compatibility checks."""

from __future__ import annotations

from bilgeai.validation.preflight import CheckResult


def check_ml_dependencies() -> list[CheckResult]:
    """Check ML library versions for compatibility."""
    results = []

    # Check torch
    try:
        import torch

        version = torch.__version__
        major, _minor = [int(x) for x in version.split(".")[:2]]
        if major >= 2:
            results.append(CheckResult("torch", "pass", f"PyTorch {version}"))
        else:
            results.append(
                CheckResult(
                    "torch",
                    "warn",
                    f"PyTorch {version} (2.0+ recommended)",
                    suggestion="Upgrade: pip install torch>=2.0",
                )
            )
    except ImportError:
        results.append(
            CheckResult(
                "torch",
                "warn",
                "PyTorch not installed",
                suggestion="Install: pip install bilgeai[ml]",
            )
        )

    # Check transformers
    try:
        import transformers

        results.append(
            CheckResult("transformers", "pass", f"transformers {transformers.__version__}")
        )
    except ImportError:
        results.append(
            CheckResult(
                "transformers",
                "warn",
                "transformers not installed",
                suggestion="Install: pip install bilgeai[ml]",
            )
        )

    # Check peft
    try:
        import peft

        results.append(CheckResult("peft", "pass", f"peft {peft.__version__}"))
    except ImportError:
        results.append(
            CheckResult(
                "peft",
                "warn",
                "peft not installed (needed for LoRA)",
                suggestion="Install: pip install bilgeai[ml]",
            )
        )

    return results
