"""
YAML capability cache resolver.

Reads model capabilities from provider-specific YAML cache files.
These caches are generated by scripts/update_capabilities.py and contain
tested capability data for each provider's models.
"""

from __future__ import annotations

from datetime import datetime
from pathlib import Path

import yaml

from chuk_llm.registry.models import ModelCapabilities, ModelSpec, QualityTier
from chuk_llm.registry.resolvers.base import BaseCapabilityResolver


class YamlCapabilityResolver(BaseCapabilityResolver):
    """
    Resolves capabilities from YAML cache files.

    Reads provider-specific cache files (e.g., openai.yaml, anthropic.yaml)
    that contain tested capability data. These files are generated by
    scripts/update_capabilities.py and updated before releases.
    """

    def __init__(self, cache_dir: Path | None = None):
        """
        Initialize YAML capability resolver.

        Args:
            cache_dir: Directory containing provider cache files
                      (defaults to package capabilities directory)
        """
        if cache_dir is None:
            # Use the package's capabilities directory
            package_dir = Path(__file__).parent.parent
            cache_dir = package_dir / "capabilities"

        self.cache_dir = Path(cache_dir)
        self._caches: dict[str, dict] = {}  # provider -> cache data
        self._load_all_caches()

    def _load_all_caches(self) -> None:
        """Load all provider cache files."""
        if not self.cache_dir.exists():
            return

        # Load each .yaml file in the cache directory
        for yaml_file in self.cache_dir.glob("*.yaml"):
            provider = yaml_file.stem  # filename without extension
            try:
                with open(yaml_file) as f:
                    self._caches[provider] = yaml.safe_load(f) or {}
            except Exception:
                # Skip invalid YAML files
                continue

    async def get_capabilities(self, spec: ModelSpec) -> ModelCapabilities:
        """
        Get capabilities from YAML cache.

        Args:
            spec: Model specification

        Returns:
            Model capabilities from cache or family defaults
        """
        # Get provider cache
        provider_cache = self._caches.get(spec.provider, {})
        if not provider_cache:
            return self._empty_capabilities()

        # Check for exact model match
        models = provider_cache.get("models", {})
        if spec.name in models:
            model_data = models[spec.name]
            # Handle inheritance
            if "inherits_from" in model_data:
                parent_name = model_data["inherits_from"]
                if parent_name in models:
                    # Merge parent capabilities with child overrides
                    parent_caps = self._parse_model_capabilities(models[parent_name])
                    child_caps = self._parse_model_capabilities(model_data)
                    return parent_caps.merge(child_caps)
            return self._parse_model_capabilities(model_data)

        # Check if model inherits from a family
        for model_name, model_data in models.items():
            if model_name == spec.name:
                continue
            # If this model is in the same family, use as template
            if "inherits_from" in model_data and spec.family == model_data.get(
                "inherits_from"
            ):
                # Found a family member - use as baseline
                return self._parse_model_capabilities(model_data)

        # Check family defaults
        families = provider_cache.get("families", {})
        if spec.family and spec.family in families:
            return self._parse_model_capabilities(families[spec.family])

        # No match found
        return self._empty_capabilities()

    def _parse_model_capabilities(self, cap_data: dict) -> ModelCapabilities:
        """
        Parse model capabilities from YAML data - Pydantic native!

        The YAML cache structure matches ModelCapabilities fields exactly,
        so we can just validate and create the Pydantic model directly.
        """
        try:
            # Convert known_params from list to set if present
            if "known_params" in cap_data and isinstance(
                cap_data["known_params"], list
            ):
                cap_data["known_params"] = set(cap_data["known_params"])

            # Parse quality_tier enum if it's a string
            if "quality_tier" in cap_data and isinstance(cap_data["quality_tier"], str):
                cap_data["quality_tier"] = QualityTier(cap_data["quality_tier"])

            # Let Pydantic do all the validation and parsing!
            return ModelCapabilities(
                **cap_data,
                source="yaml_cache",
                last_updated=datetime.now().isoformat(),
            )
        except Exception:
            # Invalid data - return empty
            return self._empty_capabilities()
