import datetime as dt
from abc import ABC, abstractmethod
from typing import Callable, Iterator
from collections.abc import Mapping

import requests

from llmdb.schemas import Model, Provider
from llmdb.query import Query, ModelQuery, ProviderQuery


API_URL = "https://models.dev/api.json"


class Collection[T, Q: Query](ABC):
    """Abstract base class for queryable collections.

    Type Parameters:
        T: The type of items in the collection (Model or Provider).
        Q: The query builder type for this collection.
    """

    @abstractmethod
    def get(self, item_id: str) -> T | None:
        """Get an item by its ID.

        Args:
            item_id: The unique identifier of the item.

        Returns:
            The item if found, None otherwise.
        """
        ...

    @abstractmethod
    def filter(self, **kwargs: object) -> Q:
        """Start a filtered query on the collection.

        Args:
            **kwargs: Field-value pairs to filter by. Keys can use dot notation
                for nested fields.

        Returns:
            A query builder with the filter applied.
        """
        ...

    @abstractmethod
    def filter_by(self, predicate: Callable[[T], bool]) -> Q:
        """Start a query with a custom predicate function.

        Args:
            predicate: A callable that takes an item and returns True if it
                should be included.

        Returns:
            A query builder with the predicate applied.
        """
        ...

    @abstractmethod
    def sort_by(self, field: str, desc: bool = False) -> Q:
        """Start a sorted query on the collection.

        Args:
            field: The field name to sort by. Can use dot notation for nested fields.
            desc: If True, sort in descending order.

        Returns:
            A query builder with sorting applied.
        """
        ...

    @abstractmethod
    def all(self) -> list[T]:
        """Return all items in the collection.

        Returns:
            A list of all items.
        """
        ...

    @abstractmethod
    def count(self) -> int:
        """Return the number of items in the collection.

        Returns:
            The total count of items.
        """
        ...

    @abstractmethod
    def __iter__(self) -> Iterator[T]:
        ...

    @abstractmethod
    def __len__(self) -> int:
        ...


class ModelCollection(Collection[Model, ModelQuery]):
    """Collection of models with query and direct lookup capabilities.

    Example:
        ```python
        model = db.models.get("gpt-4o")
        reasoning_models = db.models.filter(reasoning=True).all()
        openai_models = db.models.by_provider("openai").all()
        ```
    """

    def __init__(
        self,
        models: dict[str, Model],
        by_provider: dict[str, list[Model]],
        by_family: dict[str, list[Model]],
    ) -> None:
        self._models = models
        self._by_provider = by_provider
        self._by_family = by_family

    def get(self, item_id: str) -> Model | None:
        """Get a model by its ID.

        Args:
            item_id: The unique identifier of the model.

        Returns:
            The Model if found, None otherwise.
        """
        return self._models.get(item_id)

    def filter(self, **kwargs: object) -> ModelQuery:
        """Start a filtered query on all models.

        Args:
            **kwargs: Field-value pairs to filter by.

        Returns:
            A ModelQuery with the filter applied.

        Example:
            ```python
            db.models.filter(reasoning=True, open_weights=True)
            ```
        """
        return ModelQuery(self._models.values()).filter(**kwargs)

    def filter_by(self, predicate: Callable[[Model], bool]) -> ModelQuery:
        """Start a query with a custom predicate function.

        Args:
            predicate: A callable that takes a Model and returns True if it
                should be included.

        Returns:
            A ModelQuery with the predicate applied.

        Example:
            ```python
            db.models.filter_by(lambda m: m.cost and m.cost.input < 1.0)
            ```
        """
        return ModelQuery(self._models.values()).filter_by(predicate)

    def sort_by(self, field: str, desc: bool = False) -> ModelQuery:
        """Start a sorted query on all models.

        Args:
            field: The field name to sort by.
            desc: If True, sort in descending order.

        Returns:
            A ModelQuery with sorting applied.

        Example:
            ```python
            db.models.sort_by("release_date", desc=True)
            ```
        """
        return ModelQuery(self._models.values()).sort_by(field, desc)

    def by_provider(self, provider_id: str) -> ModelQuery:
        """Get all models for a specific provider.

        Args:
            provider_id: The provider identifier (e.g., "openai", "anthropic").

        Returns:
            A ModelQuery containing the provider's models.

        Example:
            ```python
            openai_models = db.models.by_provider("openai").all()
            ```
        """
        return ModelQuery(self._by_provider.get(provider_id, []))

    def by_family(self, family: str) -> ModelQuery:
        """Get all models in a specific family.

        Args:
            family: The model family name (e.g., "gpt", "claude").

        Returns:
            A ModelQuery containing models in the family.

        Example:
            ```python
            gpt_models = db.models.by_family("gpt").all()
            ```
        """
        return ModelQuery(self._by_family.get(family, []))

    def all(self) -> list[Model]:
        """Return all models as a list.

        Returns:
            A list of all Model objects.
        """
        return list(self._models.values())

    def count(self) -> int:
        """Return the total number of models.

        Returns:
            The count of all models.
        """
        return len(self._models)

    def __iter__(self) -> Iterator[Model]:
        return iter(self._models.values())

    def __len__(self) -> int:
        return len(self._models)


class ProviderCollection(Collection[Provider, ProviderQuery]):
    """Collection of providers with query and direct lookup capabilities.

    Example:
        ```python
        provider = db.providers.get("openai")
        all_providers = db.providers.all()
        ```
    """

    def __init__(self, providers: dict[str, Provider]) -> None:
        self._providers = providers

    def get(self, item_id: str) -> Provider | None:
        """Get a provider by its ID.

        Args:
            item_id: The unique identifier of the provider.

        Returns:
            The Provider if found, None otherwise.
        """
        return self._providers.get(item_id)

    def filter(self, **kwargs: object) -> ProviderQuery:
        """Start a filtered query on all providers.

        Args:
            **kwargs: Field-value pairs to filter by.

        Returns:
            A ProviderQuery with the filter applied.
        """
        return ProviderQuery(self._providers.values()).filter(**kwargs)

    def filter_by(self, predicate: Callable[[Provider], bool]) -> ProviderQuery:
        """Start a query with a custom predicate function.

        Args:
            predicate: A callable that takes a Provider and returns True if it
                should be included.

        Returns:
            A ProviderQuery with the predicate applied.

        Example:
            ```python
            db.providers.filter_by(lambda p: len(p.models) > 20)
            ```
        """
        return ProviderQuery(self._providers.values()).filter_by(predicate)

    def sort_by(self, field: str, desc: bool = False) -> ProviderQuery:
        """Start a sorted query on all providers.

        Args:
            field: The field name to sort by.
            desc: If True, sort in descending order.

        Returns:
            A ProviderQuery with sorting applied.
        """
        return ProviderQuery(self._providers.values()).sort_by(field, desc)

    def all(self) -> list[Provider]:
        """Return all providers as a list.

        Returns:
            A list of all Provider objects.
        """
        return list(self._providers.values())

    def count(self) -> int:
        """Return the total number of providers.

        Returns:
            The count of all providers.
        """
        return len(self._providers)

    def __iter__(self) -> Iterator[Provider]:
        return iter(self._providers.values())

    def __len__(self) -> int:
        return len(self._providers)


class Database:
    """Wrapper around the models.dev API data.

    Example:
        ```python
        from llmdb import db

        # Access triggers lazy loading
        model = db.models.get("gpt-4o")

        # Force refresh from API
        db.refresh()
        ```
    """

    def __init__(self) -> None:
        self._providers: dict[str, Provider] = {}
        self._models: dict[str, Model] = {}
        self._models_by_provider: dict[str, list[Model]] = {}
        self._models_by_family: dict[str, list[Model]] = {}
        self._fetched_at: dt.datetime | None = None
        self._loaded = False

    def _ensure_loaded(self) -> None:
        if not self._loaded:
            self.refresh()

    def refresh(self) -> None:
        """Fetch fresh data from the API and rebuild all indexes.

        Clears existing data and repopulates from the API. Updates the
        `fetched_at` timestamp on success.

        Raises:
            requests.HTTPError: If the API request fails.
        """
        response = requests.get(API_URL, timeout=30.0)
        response.raise_for_status()
        data: Mapping[str, dict] = response.json()

        self._providers.clear()
        self._models.clear()
        self._models_by_provider.clear()
        self._models_by_family.clear()

        for provider_id, provider_data in data.items():
            models_data = provider_data.pop("models", {})
            provider = Provider.model_validate({**provider_data, "models": {}})
            self._providers[provider_id] = provider
            self._models_by_provider[provider_id] = []

            for model_id, model_data in models_data.items():
                model = Model.model_validate({**model_data, "provider_id": provider_id})
                provider.models[model_id] = model
                self._models[model_id] = model
                self._models_by_provider[provider_id].append(model)

                if model.family is not None:
                    if model.family not in self._models_by_family:
                        self._models_by_family[model.family] = []
                    self._models_by_family[model.family].append(model)

        self._fetched_at = dt.datetime.now(dt.UTC)
        self._loaded = True

    @property
    def models(self) -> ModelCollection:
        """Access the models collection for queries and lookups.

        Returns:
            A ModelCollection providing access to all models.
        """
        self._ensure_loaded()
        return ModelCollection(
            self._models,
            self._models_by_provider,
            self._models_by_family,
        )

    @property
    def providers(self) -> ProviderCollection:
        """Access the providers collection for queries and lookups.

        Returns:
            A ProviderCollection providing access to all providers.
        """
        self._ensure_loaded()
        return ProviderCollection(self._providers)

    @property
    def fetched_at(self) -> dt.datetime | None:
        """Timestamp of when data was last fetched from the API.

        Returns:
            The datetime of the last fetch, or None if not yet loaded.
        """
        return self._fetched_at
