"""Model Registry — metadata store for all known models.

Stores rich metadata per the Model Intelligence infobox specification:
- provider, model ID, version, family, type, context window
- capabilities: vision, tool/function calling, structured output, reasoning
- coding, multimodal, local availability, API availability
- license, commercial use, free status, registration required
- rate limits, token limits, discovered, last checked
- status, source, trust, environment, test status
- safety profile, resource requirements
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

logger = logging.getLogger(__name__)


class ModelStatus(str, Enum):
    """Operational status of a model."""
    ACTIVE = "active"
    DEGRADED = "degraded"
    HIDDEN = "hidden"
    FAILED = "failed"
    DEPRECATED = "deprecated"
    UNKNOWN = "unknown"


class ModelTrust(str, Enum):
    """Trust level of a model source."""
    OFFICIAL = "official"
    VERIFIED = "verified"
    COMMUNITY = "community"
    UNOFFICIAL = "unofficial"
    UNKNOWN = "unknown"


class ModelType(str, Enum):
    """Type/category of model."""
    CHAT = "chat"
    COMPLETION = "completion"
    EMBEDDING = "embedding"
    REASONING = "reasoning"
    MULTIMODAL = "multimodal"
    CODE = "code"
    SYSTEM = "system"
    UNKNOWN = "unknown"


class ModelFreeStatus(str, Enum):
    """Free tier status of a model."""
    FREE = "free"
    FREE_FOR_TEST = "free_for_test"
    FREE_FOR_STUDY = "free_for_study"
    PAID = "paid"
    UNKNOWN = "unknown"


@dataclass
class ModelCapabilities:
    """Model capability flags."""

    vision: bool = False
    tool_calling: bool = False
    structured_output: bool = False
    reasoning: bool = False
    coding: bool = False
    multimodal: bool = False
    function_calling: bool = False
    json_mode: bool = False


@dataclass
class ModelRecord:
    """Complete model metadata record."""

    model_id: str
    provider: str
    version: str | None = None
    family: str | None = None
    model_type: ModelType = ModelType.CHAT
    context_window: int | None = None
    capabilities: ModelCapabilities = field(default_factory=ModelCapabilities)

    # Availability
    local_available: bool = False
    api_available: bool = False

    # Licensing
    license: str | None = None
    commercial_use: bool = True
    free_status: ModelFreeStatus = ModelFreeStatus.UNKNOWN
    registration_required: bool = False
    card_required: bool = False

    # Limits
    rate_limit: str | None = None
    max_tokens: int | None = None

    # Discovery
    discovered: float = field(default_factory=time.time)
    last_checked: float = field(default_factory=time.time)

    # Status
    status: ModelStatus = ModelStatus.UNKNOWN
    source: str | None = None
    trust: ModelTrust = ModelTrust.UNKNOWN
    environment: str = "all"
    test_status: str = "not_tested"

    # Safety and resources
    safety_profile: str | None = None
    resource_requirements: dict[str, Any] = field(default_factory=dict)

    # Extra metadata
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def is_free(self) -> bool:
        return self.free_status in (
            ModelFreeStatus.FREE,
            ModelFreeStatus.FREE_FOR_TEST,
            ModelFreeStatus.FREE_FOR_STUDY,
        )

    @property
    def is_local(self) -> bool:
        return self.local_available

    def update(self, **kwargs: Any) -> None:
        """Update fields in-place."""
        for key, value in kwargs.items():
            if hasattr(self, key):
                setattr(self, key, value)
        self.last_checked = time.time()


class ModelRegistry:
    """Registry for model metadata.

    Provides CRUD operations and queries for model discovery,
    validation, and classification. Thread-safe for single-threaded asyncio use.

    Models are keyed by ``(provider, model_id)`` so the same model name can
    exist under multiple providers.
    """

    def __init__(self) -> None:
        self._models: dict[str, ModelRecord] = {}
        self._by_provider: dict[str, set[str]] = {}
        self._by_family: dict[str, set[str]] = {}
        self._by_type: dict[str, set[str]] = {}
        self._by_status: dict[ModelStatus, set[str]] = {
            s: set() for s in ModelStatus
        }
        self._counter = 0

    def _next_key(self, model: ModelRecord) -> str:
        self._counter += 1
        return f"{model.provider}:{model.model_id}#{self._counter}"

    def register(self, model: ModelRecord) -> ModelRecord:
        """Register a model and return its record.

        Each registration creates a new unique record, even if a record with
        the same ``(provider, model_id)`` already exists. This allows multiple
        versions or statuses of the same model name to coexist.
        """
        key = self._next_key(model)
        model.metadata = dict(model.metadata)
        model.metadata.setdefault("_registry_key", key)

        self._models[key] = model
        self._by_provider.setdefault(model.provider, set()).add(key)
        if model.family:
            self._by_family.setdefault(model.family, set()).add(key)
        self._by_type.setdefault(model.model_type.value, set()).add(key)
        self._by_status[model.status].add(key)
        logger.info("Registered model %s (provider=%s)", model.model_id, model.provider)
        return model

    def unregister(self, model_id: str, provider: str | None = None) -> bool:
        """Remove matching models. Returns True if any were removed."""
        keys = [
            k for k, rec in self._models.items()
            if rec.model_id == model_id and (provider is None or rec.provider == provider)
        ]
        if not keys:
            return False
        for key in keys:
            model = self._models.pop(key)
            self._by_provider.get(model.provider, set()).discard(key)
            if model.family:
                self._by_family.get(model.family, set()).discard(key)
            self._by_type.get(model.model_type.value, set()).discard(key)
            self._by_status[model.status].discard(key)
            logger.info("Unregistered model %s", model_id)
        return True

    def get(self, model_id: str, provider: str | None = None) -> ModelRecord | None:
        """Return the first registered record matching ``model_id``.

        When ``provider`` is given, restrict the match to that provider.
        """
        for rec in self._models.values():
            if rec.model_id == model_id and (provider is None or rec.provider == provider):
                return rec
        return None

    def list_by_provider(self, provider: str) -> list[ModelRecord]:
        """List models from a specific provider."""
        return [
            self._models[mid]
            for mid in self._by_provider.get(provider, set())
            if mid in self._models
        ]

    def list_by_family(self, family: str) -> list[ModelRecord]:
        """List models in a family."""
        return [
            self._models[mid]
            for mid in self._by_family.get(family, set())
            if mid in self._models
        ]

    def list_by_type(self, model_type: ModelType | str) -> list[ModelRecord]:
        """List models of a given type."""
        type_key = model_type.value if isinstance(model_type, ModelType) else model_type
        return [
            self._models[mid]
            for mid in self._by_type.get(type_key, set())
            if mid in self._models
        ]

    def list_by_status(self, status: ModelStatus) -> list[ModelRecord]:
        """List models with a given status."""
        return [
            self._models[mid]
            for mid in self._by_status.get(status, set())
            if mid in self._models
        ]

    def list_active(self) -> list[ModelRecord]:
        """List all active models."""
        return [
            self._models[mid]
            for mid in self._by_status.get(ModelStatus.ACTIVE, set())
            if mid in self._models
        ]

    def list_local(self) -> list[ModelRecord]:
        """List all locally available models."""
        return [m for m in self._models.values() if m.local_available]

    def list_free(self) -> list[ModelRecord]:
        """List all free models."""
        return [m for m in self._models.values() if m.is_free]

    def list_all(self) -> list[ModelRecord]:
        """List all registered models."""
        return list(self._models.values())

    def search(
        self,
        *,
        provider: str | None = None,
        family: str | None = None,
        model_type: ModelType | str | None = None,
        status: ModelStatus | None = None,
        local_only: bool | None = None,
        free_only: bool | None = None,
        trust: ModelTrust | None = None,
    ) -> list[ModelRecord]:
        """Search models with multiple filters."""
        results = list(self._models.values())

        if provider:
            results = [m for m in results if m.provider == provider]
        if family:
            results = [m for m in results if m.family == family]
        if model_type:
            type_key = model_type.value if isinstance(model_type, ModelType) else model_type
            results = [m for m in results if m.model_type.value == type_key]
        if status:
            results = [m for m in results if m.status == status]
        if local_only is not None:
            results = [m for m in results if m.local_available == local_only]
        if free_only is not None:
            results = [m for m in results if m.is_free == free_only]
        if trust:
            results = [m for m in results if m.trust == trust]

        return results

    def update(self, model_id: str, provider: str | None = None, **kwargs: Any) -> ModelRecord | None:
        """Update model metadata. Returns updated record or None.

        Looks up the record by ``model_id`` (and optional ``provider``).
        """
        model = self.get(model_id, provider)
        if model is None:
            return None
        key = model.metadata.get("_registry_key")
        if key is None or key not in self._models:
            return None
        old_provider = model.provider
        old_family = model.family
        old_type = model.model_type
        old_status = model.status

        model.update(**kwargs)

        if old_provider != model.provider:
            self._by_provider.get(old_provider, set()).discard(key)
            self._by_provider.setdefault(model.provider, set()).add(key)
        if old_family != model.family:
            if old_family:
                self._by_family.get(old_family, set()).discard(key)
            if model.family:
                self._by_family.setdefault(model.family, set()).add(key)
        if old_type != model.model_type:
            self._by_type.get(old_type.value, set()).discard(key)
            self._by_type.setdefault(model.model_type.value, set()).add(key)
        if old_status != model.status:
            self._by_status[old_status].discard(key)
            self._by_status[model.status].add(key)
        return model

    def count(self) -> int:
        """Total number of registered models."""
        return len(self._models)

    def count_by_provider(self, provider: str) -> int:
        return len(self._by_provider.get(provider, set()))

    def has(self, model_id: str, provider: str | None = None) -> bool:
        if provider is not None:
            return any(
                rec.model_id == model_id and rec.provider == provider
                for rec in self._models.values()
            )
        return any(rec.model_id == model_id for rec in self._models.values())