from typing import Callable, Iterator, Any
from collections.abc import Iterable

from llmdb.schemas import Model, Provider


def _get_nested_attr(obj: object, path: str) -> Any:
    """
    Get a nested attribute using dot notation (e.g., 'cost.input').
    """
    for part in path.split("."):
        obj = getattr(obj, part)
    return obj


class Query[T: (Model, Provider)]:
    """Fluent query builder for filtering and sorting collections.

    Provides a chainable API for filtering, sorting, and retrieving items
    from a collection of Model or Provider objects.

    Example:
        ```python
        results = (
            query
            .filter(reasoning=True)
            .sort_by("release_date", desc=True)
            .all()
        )
        ```
    """

    def __init__(self, items: Iterable[T]) -> None:
        self._items = list(items)

    def filter(self, **kwargs: object) -> "Query[T]":
        """Filter items by field equality.

        Supports nested fields using dot notation for accessing attributes
        on nested objects.

        Args:
            **kwargs: Field-value pairs to filter by. Keys can use dot notation
                for nested fields (e.g., `cost.input=0.5`).

        Returns:
            A new Query containing only items matching all conditions.

        Example:
            ```python
            db.models.filter(reasoning=True, open_weights=True)
            db.models.filter(**{"cost.input": 0.5})
            ```
        """
        result = self._items
        for key, value in kwargs.items():
            result = [
                item
                for item in result
                if _get_nested_attr(item, key) == value
            ]
        return Query(result)

    def filter_by(self, predicate: Callable[[T], bool]) -> "Query[T]":
        """Filter items using a custom predicate function.

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

        Returns:
            A new Query containing only items for which predicate returns True.

        Example:
            ```python
            db.models.filter_by(lambda m: m.cost and m.cost.input < 1.0)
            db.providers.filter_by(lambda p: len(p.models) > 10)
            ```
        """
        return Query([item for item in self._items if predicate(item)])

    def sort_by(self, field: str, desc: bool = False) -> "Query[T]":
        """Sort items by a field.

        Supports nested fields using dot notation for accessing attributes
        on nested objects.

        Args:
            field: The field name to sort by. Can use dot notation for nested
                fields (e.g., `cost.input`).
            desc: If True, sort in descending order. Defaults to False.

        Returns:
            A new Query with items sorted by the specified field.

        Example:
            ```python
            db.models.sort_by("release_date", desc=True)
            db.models.sort_by("cost.input")
            ```
        """
        result = sorted(
            self._items,
            key=lambda item: _get_nested_attr(item, field),
            reverse=desc,
        )
        return Query(result)

    def all(self) -> list[T]:
        """Return all matching items as a list.

        Returns:
            A list containing all items in the current query result.

        Example:
            ```python
            models = db.models.filter(reasoning=True).all()
            ```
        """
        return self._items.copy()

    def first(self) -> T | None:
        """Return the first matching item.

        Returns:
            The first item in the query result, or None if empty.

        Example:
            ```python
            model = db.models.filter(id="gpt-4o").first()
            ```
        """
        return self._items[0] if self._items else None

    def count(self) -> int:
        """Return the number of matching items.

        Returns:
            The total count of items in the query result.

        Example:
            ```python
            num_reasoning = db.models.filter(reasoning=True).count()
            ```
        """
        return len(self._items)

    def __iter__(self) -> Iterator[T]:
        return iter(self._items)

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


class ModelQuery(Query[Model]):
    """Query builder for Model collections.

    Extends Query with Model-specific return types for method chaining.

    Example:
        ```python
        models = (
            db.models
            .filter(reasoning=True)
            .sort_by("release_date", desc=True)
            .all()
        )
        ```
    """

    def filter(self, **kwargs: object) -> "ModelQuery":
        """Filter models by field equality.

        Args:
            **kwargs: Field-value pairs to filter by. Keys can use dot notation
                for nested fields (e.g., `cost.input=0.5`).

        Returns:
            A new ModelQuery containing only models matching all conditions.
        """
        result = super().filter(**kwargs)
        return ModelQuery(result._items)

    def filter_by(self, predicate: Callable[[Model], bool]) -> "ModelQuery":
        """Filter models using a custom predicate function.

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

        Returns:
            A new ModelQuery containing only models for which predicate returns True.
        """
        result = super().filter_by(predicate)
        return ModelQuery(result._items)

    def sort_by(self, field: str, desc: bool = False) -> "ModelQuery":
        """Sort models by a field.

        Args:
            field: The field name to sort by. Can use dot notation for nested
                fields (e.g., `cost.input`).
            desc: If True, sort in descending order. Defaults to False.

        Returns:
            A new ModelQuery with models sorted by the specified field.
        """
        result = super().sort_by(field, desc)
        return ModelQuery(result._items)


class ProviderQuery(Query[Provider]):
    """Query builder for Provider collections.

    Extends Query with Provider-specific return types for method chaining.

    Example:
        ```python
        providers = (
            db.providers
            .filter_by(lambda p: len(p.models) > 20)
            .sort_by("name")
            .all()
        )
        ```
    """

    def filter(self, **kwargs: object) -> "ProviderQuery":
        """Filter providers by field equality.

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

        Returns:
            A new ProviderQuery containing only providers matching all conditions.
        """
        result = super().filter(**kwargs)
        return ProviderQuery(result._items)

    def filter_by(self, predicate: Callable[[Provider], bool]) -> "ProviderQuery":
        """Filter providers using a custom predicate function.

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

        Returns:
            A new ProviderQuery containing only providers for which predicate returns True.
        """
        result = super().filter_by(predicate)
        return ProviderQuery(result._items)

    def sort_by(self, field: str, desc: bool = False) -> "ProviderQuery":
        """Sort providers by a field.

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

        Returns:
            A new ProviderQuery with providers sorted by the specified field.
        """
        result = super().sort_by(field, desc)
        return ProviderQuery(result._items)
