"""
Model Registry for Fal.ai MCP Server

Provides dynamic model discovery from the Fal.ai Platform API with TTL-based caching.
Supports both full model IDs (e.g., 'fal-ai/flux-pro') and friendly aliases (e.g., 'flux_schnell').
"""

import asyncio
import os
import time
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Union

import httpx
from loguru import logger


@dataclass
class FalModel:
    """Represents a model from the Fal.ai platform."""

    id: str  # e.g., "fal-ai/flux-pro/v1.1-ultra"
    name: str  # Human-readable name
    description: str  # Model description
    category: str  # e.g., "text-to-image", "image-to-video", "audio"
    owner: str = ""  # e.g., "fal-ai"
    thumbnail_url: Optional[str] = None
    highlighted: bool = False  # Featured/recommended model
    group_key: Optional[str] = None  # Model family grouping key
    group_label: Optional[str] = None  # Model family display name
    status: str = "active"  # Model status (active, deprecated)
    tags: Optional[List[str]] = None  # Model tags for categorization


@dataclass
class ModelCache:
    """Cached model data with TTL."""

    models: Dict[str, FalModel]  # model_id -> FalModel
    aliases: Dict[str, str]  # alias -> model_id
    by_category: Dict[str, List[str]]  # category -> list of model_ids
    fetched_at: float  # timestamp
    ttl_seconds: int = 3600  # 1 hour default


@dataclass
class SearchResult:
    """Result from search_models with fallback indicator."""

    models: List[FalModel]
    used_fallback: bool = False  # True if results came from local cache
    fallback_reason: Optional[str] = None  # Why fallback was used


@dataclass
class RecommendationResult:
    """Result from recommend_models with fallback indicator."""

    recommendations: List[Dict[str, Any]]
    used_fallback: bool = False  # True if search fell back to cache
    fallback_reason: Optional[str] = None  # Why fallback was used


class ModelRegistry:
    """
    Manages model discovery from Fal.ai API with caching.

    Features:
    - Lazy loading on first use
    - TTL-based cache with configurable expiration
    - Backward-compatible friendly aliases
    - Auto-detection of full IDs vs aliases
    - Thread-safe async operations
    """

    FAL_API_BASE = "https://api.fal.ai/v1"
    DEFAULT_TTL = 3600  # 1 hour

    # Backward-compatible aliases for existing models
    LEGACY_ALIASES: Dict[str, str] = {
        # Image models
        "flux_schnell": "fal-ai/flux/schnell",
        "flux_dev": "fal-ai/flux/dev",
        "flux_pro": "fal-ai/flux-pro",
        "sdxl": "fal-ai/fast-sdxl",
        "stable_diffusion": "fal-ai/stable-diffusion-v3-medium",
        # Video models
        "svd": "fal-ai/stable-video-diffusion",
        "animatediff": "fal-ai/fast-animatediff",
        "kling": "fal-ai/kling-video",
        # Audio models
        "musicgen": "fal-ai/lyria2",  # Updated: musicgen-medium no longer exists
        "lyria2": "fal-ai/lyria2",
        "stable_audio": "fal-ai/stable-audio-25/text-to-audio",
        "bark": "fal-ai/bark",
        "whisper": "fal-ai/whisper",
    }

    # Category mappings from Fal API categories to our simplified categories
    CATEGORY_MAPPING: Dict[str, str] = {
        "text-to-image": "image",
        "image-to-image": "image",
        "text-to-video": "video",
        "image-to-video": "video",
        "video-to-video": "video",
        "audio": "audio",
        "text-to-audio": "audio",
        "speech-to-text": "audio",
        "text-to-speech": "audio",
        "audio-to-audio": "audio",
    }

    # Legacy alias category mappings for fallback cache
    LEGACY_ALIAS_CATEGORIES: Dict[str, str] = {
        # Image models
        "flux_schnell": "image",
        "flux_dev": "image",
        "flux_pro": "image",
        "sdxl": "image",
        "stable_diffusion": "image",
        # Video models
        "svd": "video",
        "animatediff": "video",
        "kling": "video",
        # Audio models
        "musicgen": "audio",
        "lyria2": "audio",
        "stable_audio": "audio",
        "bark": "audio",
        "whisper": "audio",
    }

    # Shorter TTL for fallback cache to retry API sooner
    FALLBACK_TTL = 60  # 1 minute

    def __init__(self, ttl_seconds: int = DEFAULT_TTL):
        self._cache: Optional[ModelCache] = None
        self._lock = asyncio.Lock()
        self._ttl_seconds = ttl_seconds
        self._http_client: Optional[httpx.AsyncClient] = None

    async def _get_http_client(self) -> httpx.AsyncClient:
        """Get or create HTTP client."""
        if self._http_client is None:
            headers = {"Content-Type": "application/json"}
            api_key = os.getenv("FAL_KEY")
            if api_key:
                headers["Authorization"] = f"Key {api_key}"
            else:
                logger.warning(
                    "FAL_KEY environment variable not set - "
                    "model registry API calls may fail with 401 Unauthorized"
                )
            self._http_client = httpx.AsyncClient(
                base_url=self.FAL_API_BASE, headers=headers, timeout=30.0
            )
        return self._http_client

    async def _fetch_models_page(
        self,
        cursor: Optional[str] = None,
        category: Optional[str] = None,
        limit: int = 100,
    ) -> Dict[str, Any]:
        """Fetch a single page of models from the API."""
        client = await self._get_http_client()
        params: Dict[str, Any] = {"limit": limit}
        if cursor:
            params["cursor"] = cursor
        if category:
            params["category"] = category

        response = await client.get("/models", params=params)
        response.raise_for_status()
        data: Dict[str, Any] = response.json()
        return data

    async def _fetch_all_models(
        self, category: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """Fetch all models with pagination."""
        all_models: List[Dict[str, Any]] = []
        cursor: Optional[str] = None

        while True:
            data = await self._fetch_models_page(cursor=cursor, category=category)
            # API returns "models" array, not "items"
            models = data.get("models", [])

            # Log warning if API returns empty or unexpected response
            if not models and not all_models:
                logger.warning(
                    "API returned empty models list - response keys: %s",
                    list(data.keys()),
                )

            all_models.extend(models)

            cursor = data.get("next_cursor")
            if not cursor or not data.get("has_more", False):
                break

        return all_models

    async def _refresh_cache(self) -> ModelCache:
        """Refresh the model cache from the API."""
        raw_models = await self._fetch_all_models()

        models: Dict[str, FalModel] = {}
        aliases: Dict[str, str] = dict(self.LEGACY_ALIASES)  # Start with legacy aliases
        by_category: Dict[str, List[str]] = {"image": [], "video": [], "audio": []}

        for raw in raw_models:
            model_id = raw.get("endpoint_id", "")
            if not model_id:
                continue

            # Metadata is nested under "metadata" key in API response
            metadata = raw.get("metadata", {})

            # Extract owner from endpoint_id (e.g., "fal-ai/flux/dev" -> "fal-ai")
            owner = model_id.split("/")[0] if "/" in model_id else ""

            # Extract group info (nested under "group" key)
            group = raw.get("group", {})

            model = FalModel(
                id=model_id,
                name=metadata.get("display_name", model_id),
                description=metadata.get("description", ""),
                category=metadata.get("category", ""),
                owner=owner,
                thumbnail_url=metadata.get("thumbnail_url"),
                highlighted=raw.get("highlighted", False),
                group_key=group.get("key") if group else None,
                group_label=group.get("label") if group else None,
                status=raw.get("status", "active"),
                tags=metadata.get("tags"),
            )
            models[model_id] = model

            # Map to our simplified category system
            our_category = self.CATEGORY_MAPPING.get(model.category)
            if our_category and our_category in by_category:
                by_category[our_category].append(model_id)

            # Generate automatic alias from model ID
            auto_alias = self._generate_alias(model_id)
            if auto_alias and auto_alias not in aliases:
                aliases[auto_alias] = model_id

        return ModelCache(
            models=models,
            aliases=aliases,
            by_category=by_category,
            fetched_at=time.time(),
            ttl_seconds=self._ttl_seconds,
        )

    def _generate_alias(self, model_id: str) -> Optional[str]:
        """Generate a friendly alias from a model ID."""
        # Remove "fal-ai/" prefix and convert to snake_case
        if model_id.startswith("fal-ai/"):
            name = model_id[7:]  # Remove "fal-ai/"
            # Replace slashes and dashes with underscores
            alias = name.replace("/", "_").replace("-", "_")
            return alias
        return None

    def _is_cache_valid(self) -> bool:
        """Check if the cache is still valid."""
        if self._cache is None:
            return False
        age = time.time() - self._cache.fetched_at
        return age < self._cache.ttl_seconds

    async def get_cache(self) -> ModelCache:
        """Get the model cache, refreshing if necessary."""
        async with self._lock:
            if not self._is_cache_valid():
                try:
                    self._cache = await self._refresh_cache()
                    logger.info(
                        "Model cache refreshed: %d models, %d aliases",
                        len(self._cache.models),
                        len(self._cache.aliases),
                    )
                except httpx.HTTPStatusError as e:
                    logger.error(
                        "Failed to refresh model cache: HTTP %d - %s",
                        e.response.status_code,
                        str(e),
                    )
                    self._handle_cache_refresh_failure()
                except httpx.TimeoutException as e:
                    logger.error("Timeout refreshing model cache: %s", e)
                    self._handle_cache_refresh_failure()
                except httpx.ConnectError as e:
                    logger.error("Connection error refreshing model cache: %s", e)
                    self._handle_cache_refresh_failure()
                except Exception as e:
                    logger.exception("Unexpected error refreshing model cache: %s", e)
                    self._handle_cache_refresh_failure()
            assert (
                self._cache is not None
            )  # Guaranteed by _handle_cache_refresh_failure
            return self._cache

    def _handle_cache_refresh_failure(self) -> None:
        """Handle cache refresh failure by using stale cache or creating fallback."""
        if self._cache is not None:
            cache_age = time.time() - self._cache.fetched_at
            logger.warning(
                "Using stale cache (age: %.0fs) due to refresh failure", cache_age
            )
        else:
            logger.warning(
                "Creating fallback cache with legacy aliases only - "
                "dynamic model discovery unavailable"
            )
            self._cache = self._create_fallback_cache()

    def _create_fallback_cache(self) -> ModelCache:
        """Create a fallback cache with legacy aliases converted to FalModel objects."""
        models: Dict[str, FalModel] = {}
        by_category: Dict[str, List[str]] = {"image": [], "video": [], "audio": []}

        for alias, model_id in self.LEGACY_ALIASES.items():
            # Skip if we already processed this model_id (handles duplicate aliases)
            if model_id in models:
                continue

            # Get category with warning for missing mappings
            category = self.LEGACY_ALIAS_CATEGORIES.get(alias)
            if category is None:
                logger.warning(
                    "Missing category mapping for legacy alias '%s' (model: %s) "
                    "- defaulting to 'image'",
                    alias,
                    model_id,
                )
                category = "image"

            model = FalModel(
                id=model_id,
                name=alias.replace(
                    "_", " "
                ).title(),  # e.g., "flux_schnell" -> "Flux Schnell"
                description=f"Fal.ai {category} model ({model_id})",
                category=category,
                owner="fal-ai",
            )
            models[model_id] = model
            by_category[category].append(model_id)

        return ModelCache(
            models=models,
            aliases=dict(self.LEGACY_ALIASES),
            by_category=by_category,
            fetched_at=time.time(),
            ttl_seconds=self.FALLBACK_TTL,  # Shorter TTL to retry API sooner
        )

    def is_full_model_id(self, model_input: str) -> bool:
        """Check if input looks like a full model ID vs alias."""
        # Full IDs contain "/" (e.g., "fal-ai/flux-pro/v1.1-ultra")
        return "/" in model_input

    async def resolve_model_id(self, model_input: str) -> str:
        """
        Resolve a model input to a full model ID.

        Accepts either:
        - Full model ID: "fal-ai/flux-pro/v1.1-ultra" -> returns as-is
        - Friendly alias: "flux_pro" -> resolves to "fal-ai/flux-pro"

        Raises ValueError if alias not found.
        """
        # If it looks like a full ID, return as-is
        if self.is_full_model_id(model_input):
            return model_input

        # Otherwise, look up alias
        cache = await self.get_cache()
        if model_input in cache.aliases:
            return cache.aliases[model_input]

        raise ValueError(f"Unknown model alias: {model_input}")

    async def list_models(
        self,
        category: Optional[str] = None,
        search: Optional[str] = None,
        limit: int = 50,
    ) -> List[FalModel]:
        """
        List available models with optional filtering.

        Args:
            category: Filter by category ("image", "video", "audio")
            search: Search query to filter by name/description
            limit: Maximum number of results

        Returns:
            List of FalModel objects
        """
        cache = await self.get_cache()

        if category:
            model_ids = cache.by_category.get(category, [])
            models = [cache.models[mid] for mid in model_ids if mid in cache.models]
        else:
            models = list(cache.models.values())

        if search:
            search_lower = search.lower()
            models = [
                m
                for m in models
                if search_lower in m.name.lower()
                or search_lower in m.description.lower()
                or search_lower in m.id.lower()
            ]

        return models[:limit]

    async def search_models(
        self,
        query: str,
        category: Optional[str] = None,
        limit: int = 10,
    ) -> SearchResult:
        """
        Search models using Fal API's semantic search.

        Uses the API's `q` parameter for free-text search across model name,
        description, and category. Results are sorted with highlighted models first.

        Args:
            query: Free-text search query (e.g., "portrait photo", "anime style")
            category: Optional category filter ("text-to-image", "image-to-video", etc.)
            limit: Maximum number of results

        Returns:
            SearchResult with models and fallback indicator
        """
        client = await self._get_http_client()
        params: Dict[str, Any] = {"q": query, "limit": limit, "status": "active"}
        if category:
            params["category"] = category

        fallback_reason: Optional[str] = None
        try:
            response = await client.get("/models", params=params)
            response.raise_for_status()
            data: Dict[str, Any] = response.json()
        except httpx.HTTPStatusError as e:
            fallback_reason = f"API error (HTTP {e.response.status_code})"
            logger.error(
                "Search API returned HTTP %d for query '%s': %s",
                e.response.status_code,
                query,
                e,
            )
        except httpx.TimeoutException:
            fallback_reason = "API timeout"
            logger.error("Search API timeout for query '%s'", query)
        except httpx.ConnectError as e:
            fallback_reason = "Connection error"
            logger.error("Cannot connect to search API for query '%s': %s", query, e)
        except Exception as e:
            fallback_reason = "Unexpected error"
            logger.exception(
                "Unexpected error searching models for query '%s': %s", query, e
            )

        # If fallback needed, use local cache search
        if fallback_reason:
            fallback_models = await self.list_models(
                category=self._map_to_simple_category(category) if category else None,
                search=query,
                limit=limit,
            )
            return SearchResult(
                models=fallback_models,
                used_fallback=True,
                fallback_reason=fallback_reason,
            )

        raw_models = data.get("models", [])
        models: List[FalModel] = []

        for raw in raw_models:
            model_id = raw.get("endpoint_id", "")
            if not model_id:
                continue

            metadata = raw.get("metadata", {})
            group = raw.get("group", {})
            owner = model_id.split("/")[0] if "/" in model_id else ""

            model = FalModel(
                id=model_id,
                name=metadata.get("display_name", model_id),
                description=metadata.get("description", ""),
                category=metadata.get("category", ""),
                owner=owner,
                thumbnail_url=metadata.get("thumbnail_url"),
                highlighted=raw.get("highlighted", False),
                group_key=group.get("key") if group else None,
                group_label=group.get("label") if group else None,
                status=raw.get("status", "active"),
                tags=metadata.get("tags"),
            )
            models.append(model)

        # Sort: highlighted models first, then by name
        models.sort(key=lambda m: (not m.highlighted, m.name.lower()))

        return SearchResult(models=models[:limit], used_fallback=False)

    def _map_to_simple_category(self, api_category: str) -> Optional[str]:
        """Map API category to our simplified category system."""
        return self.CATEGORY_MAPPING.get(api_category)

    async def recommend_models(
        self,
        task: str,
        category: Optional[str] = None,
        limit: int = 5,
    ) -> RecommendationResult:
        """
        Recommend the best models for a given task.

        Uses semantic search and prioritizes highlighted (featured) models.
        Returns models with relevance reasoning.

        Args:
            task: Description of the task (e.g., "generate professional headshot")
            category: Optional category hint ("image", "video", "audio")
            limit: Maximum number of recommendations

        Returns:
            RecommendationResult with recommendations and fallback indicator
        """
        # Map simplified category to API category if provided
        api_category = None
        if category:
            # Use primary mapping for each simplified category
            category_to_api = {
                "image": "text-to-image",
                "video": "text-to-video",
                "audio": "text-to-audio",
            }
            api_category = category_to_api.get(category)

        # Search using the task as query
        search_result = await self.search_models(
            query=task, category=api_category, limit=limit * 2
        )

        # Build recommendations with reasoning
        recommendations: List[Dict[str, Any]] = []
        for model in search_result.models[:limit]:
            rec = {
                "model_id": model.id,
                "name": model.name,
                "description": model.description,
                "category": model.category,
                "highlighted": model.highlighted,
                "group": model.group_label,
                "reason": self._generate_recommendation_reason(model, task),
            }
            recommendations.append(rec)

        return RecommendationResult(
            recommendations=recommendations,
            used_fallback=search_result.used_fallback,
            fallback_reason=search_result.fallback_reason,
        )

    def _generate_recommendation_reason(self, model: FalModel, task: str) -> str:
        """Generate a human-readable reason for recommending this model."""
        reasons = []

        if model.highlighted:
            reasons.append("Featured model by Fal.ai")

        if model.group_label:
            reasons.append(f"Part of {model.group_label} family")

        category_desc = {
            "text-to-image": "text-to-image generation",
            "image-to-image": "image transformation",
            "image-to-video": "image-to-video animation",
            "text-to-video": "text-to-video generation",
            "text-to-audio": "audio generation",
            "speech-to-text": "speech transcription",
        }
        if model.category in category_desc:
            reasons.append(f"Supports {category_desc[model.category]}")

        if not reasons:
            reasons.append(f"Matches search for '{task}'")

        return "; ".join(reasons)

    async def get_model(self, model_id: str) -> Optional[FalModel]:
        """Get a specific model by ID."""
        cache = await self.get_cache()
        return cache.models.get(model_id)

    async def model_exists(self, model_input: str) -> bool:
        """Check if a model exists (by ID or alias)."""
        try:
            model_id = await self.resolve_model_id(model_input)
            cache = await self.get_cache()
            # For full IDs not in cache, assume they exist (API will validate)
            if self.is_full_model_id(model_input):
                return True
            return model_id in cache.models or model_input in cache.aliases
        except ValueError:
            return False

    async def get_aliases(self) -> Dict[str, str]:
        """Get all available aliases."""
        cache = await self.get_cache()
        return dict(cache.aliases)

    async def get_pricing(self, endpoint_ids: List[str]) -> Dict[str, Any]:
        """
        Fetch pricing information for specified models.

        Args:
            endpoint_ids: List of full endpoint IDs (e.g., ["fal-ai/flux/dev"])

        Returns:
            Dict with "prices" list containing pricing info per model

        Raises:
            httpx.HTTPStatusError: If API request fails
        """
        if not endpoint_ids:
            return {"prices": []}

        client = await self._get_http_client()
        # Build query params with multiple endpoint_id values
        # Type annotation needed for mypy compatibility with httpx
        params: List[Tuple[str, Union[str, int, float, bool, None]]] = [
            ("endpoint_id", eid) for eid in endpoint_ids
        ]
        response = await client.get("/models/pricing", params=params)
        response.raise_for_status()
        result: Dict[str, Any] = response.json()
        return result

    async def get_usage(
        self,
        start: Optional[str] = None,
        end: Optional[str] = None,
        endpoint_ids: Optional[List[str]] = None,
    ) -> Dict[str, Any]:
        """
        Fetch usage and spending history.

        Args:
            start: Start date (YYYY-MM-DD format)
            end: End date (YYYY-MM-DD format)
            endpoint_ids: Optional list of endpoint IDs to filter by

        Returns:
            Dict with "time_series" and "summary" usage data

        Raises:
            httpx.HTTPStatusError: If API request fails (e.g., 401 for non-admin key)
        """
        client = await self._get_http_client()

        # Build query params
        params: Dict[str, Any] = {"expand": "summary"}
        if start:
            params["start"] = start
        if end:
            params["end"] = end

        # Add endpoint_id filters if specified
        if endpoint_ids:
            # For multiple endpoint IDs, we need to make the request with repeated params
            # httpx supports this with a list of tuples
            param_tuples: List[Tuple[str, Union[str, int, float, bool, None]]] = [
                ("expand", "summary")
            ]
            if start:
                param_tuples.append(("start", start))
            if end:
                param_tuples.append(("end", end))
            for eid in endpoint_ids:
                param_tuples.append(("endpoint_id", eid))
            response = await client.get("/models/usage", params=param_tuples)
        else:
            response = await client.get("/models/usage", params=params)

        response.raise_for_status()
        result: Dict[str, Any] = response.json()
        return result

    async def close(self) -> None:
        """Close the HTTP client."""
        if self._http_client:
            await self._http_client.aclose()
            self._http_client = None


# Module-level singleton instance
_registry: Optional[ModelRegistry] = None


async def get_registry() -> ModelRegistry:
    """Get the global ModelRegistry singleton."""
    global _registry
    if _registry is None:
        _registry = ModelRegistry()
    return _registry


def reset_registry() -> None:
    """Reset the global registry (useful for testing)."""
    global _registry
    _registry = None
