import dataclasses
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, Self, TypeVar, overload

if TYPE_CHECKING:
    from .batch import AsyncBatch, Batch
    from .transaction import AsyncTxn, Txn

import datetime

from google.api_core.exceptions import NotFound
from google.cloud.exceptions import Conflict
from google.cloud.firestore import Client
from google.cloud.firestore_v1._helpers import LastUpdateOption
from google.cloud.firestore_v1.base_query import And as FsAnd
from google.cloud.firestore_v1.base_query import FieldFilter as FsFieldFilter
from google.cloud.firestore_v1.base_query import Or as FsOr

from .exceptions import CendryError, DocumentAlreadyExistsError, DocumentNotFoundError
from .filters import And, Or
from .metadata import _clear_metadata, _set_metadata, get_metadata
from .model import FieldFilterResult, Model
from .query import AsyncQuery, Query
from .serialize import (
    deserialize,
    resolve_field_hint,
    resolve_field_path,
    serialize_update_value,
    to_dict,
    validate_required_fields,
)
from .types import TypeRegistry, default_registry

T = TypeVar("T", bound=Model)
BATCH_LIMIT = 500


def _check_batch_limit(n: int) -> None:
    if n > BATCH_LIMIT:
        raise CendryError(f"Batch limit exceeded: {n} items, maximum is {BATCH_LIMIT}")


def _resolve_precondition(
    if_unchanged: bool | datetime.datetime,
    instance: Model | None = None,
) -> LastUpdateOption | None:
    """Resolve if_unchanged into a Firestore LastUpdateOption."""
    if if_unchanged is False:
        return None
    if isinstance(if_unchanged, datetime.datetime):
        return LastUpdateOption(if_unchanged)
    # if_unchanged=True — look up metadata
    if instance is None:  # pragma: no cover
        raise CendryError("if_unchanged=True requires an instance; use a datetime for class+ID")
    meta = get_metadata(instance)
    if meta.update_time is None:  # pragma: no cover
        raise CendryError("No update_time — cannot use if_unchanged")
    return LastUpdateOption(meta.update_time)


class _BaseCendry:
    """Shared logic for sync and async contexts."""

    _client: Any
    type_registry: TypeRegistry

    def _get_collection_ref(self, model_class: type[T], parent: Model | None = None) -> Any:
        """Get a Firestore collection reference, optionally nested under a parent."""
        if parent is not None:
            if parent.id is None:
                raise CendryError("Parent model must have a non-None id")
            parent_ref = self._client.collection(parent.__collection__).document(parent.id)
            return parent_ref.collection(model_class.__collection__)
        return self._client.collection(model_class.__collection__)

    def _build_query(
        self,
        model_class: type[T],
        filters: tuple[Any, ...],
        *,
        order_by: list[Any] | None = None,
        limit: int | None = None,
        start_at: dict[str, Any] | Model | None = None,
        start_after: dict[str, Any] | Model | None = None,
        end_at: dict[str, Any] | Model | None = None,
        end_before: dict[str, Any] | Model | None = None,
        parent: Model | None = None,
        collection_group: bool = False,
    ) -> Any:
        """Build a Firestore query from filters, ordering, and pagination."""
        if collection_group:
            query = self._client.collection_group(model_class.__collection__)
        else:
            query = self._get_collection_ref(model_class, parent)

        for f in filters:
            query = self._apply_filter(query, f)

        if order_by:
            for o in order_by:
                query = query.order_by(o.field, direction=o.direction)

        if limit is not None:
            query = query.limit(limit)

        if start_at is not None:
            query = query.start_at(self._cursor_value(start_at))
        if start_after is not None:
            query = query.start_after(self._cursor_value(start_after))
        if end_at is not None:
            query = query.end_at(self._cursor_value(end_at))
        if end_before is not None:
            query = query.end_before(self._cursor_value(end_before))

        return query

    def _apply_filter(self, query: Any, f: Any) -> Any:
        """Apply a single filter or composite filter to a query."""
        if isinstance(f, FsFieldFilter):
            return query.where(filter=f)
        if isinstance(f, FieldFilterResult):
            return query.where(filter=FsFieldFilter(f.field_name, f.op, f.value))
        if isinstance(f, And):
            fs_filters = [self._to_firestore_filter(sub) for sub in f.filters]
            return query.where(filter=FsAnd(filters=fs_filters))
        if isinstance(f, Or):
            fs_filters = [self._to_firestore_filter(sub) for sub in f.filters]
            return query.where(filter=FsOr(filters=fs_filters))
        raise CendryError(f"Unknown filter type: {type(f)}")

    def _to_firestore_filter(self, f: Any) -> Any:
        """Convert a cendry filter to a Firestore filter."""
        if isinstance(f, FsFieldFilter):
            return f
        if isinstance(f, FieldFilterResult):
            return FsFieldFilter(f.field_name, f.op, f.value)
        if isinstance(f, And):
            return FsAnd(filters=[self._to_firestore_filter(sub) for sub in f.filters])
        if isinstance(f, Or):
            return FsOr(filters=[self._to_firestore_filter(sub) for sub in f.filters])
        raise CendryError(f"Unknown filter type: {type(f)}")  # pragma: no cover

    def _cursor_value(self, cursor: dict[str, Any] | Model) -> dict[str, Any]:
        """Convert a cursor to a dict for Firestore."""
        if isinstance(cursor, Model):
            return to_dict(cursor, by_alias=True, registry=self.type_registry)
        return cursor


class Cendry(_BaseCendry):
    """Synchronous Firestore ODM context.

    Args:
        client: Optional Firestore Client. Uses default credentials if not provided.
        type_registry: Optional TypeRegistry override. Uses the global default if not provided.
    """

    def __init__(
        self,
        *,
        client: Client | None = None,
        type_registry: TypeRegistry | None = None,
    ) -> None:
        self._client = client or Client()
        self.type_registry = type_registry or default_registry

    def __enter__(self) -> Self:
        return self

    def __exit__(self, *args: object) -> None:
        self._client.close()

    def get(self, model_class: type[T], document_id: str, *, parent: Model | None = None) -> T:
        """Fetch a single document by ID.

        Args:
            model_class: The Model class to deserialize into.
            document_id: Firestore document ID.
            parent: Parent document for subcollection queries.

        Returns:
            The deserialized model instance.

        Raises:
            DocumentNotFoundError: If the document does not exist.
        """
        col_ref = self._get_collection_ref(model_class, parent)
        doc = col_ref.document(document_id).get()
        if not doc.exists:
            raise DocumentNotFoundError(model_class.__collection__, document_id)
        result = deserialize(model_class, doc.id, doc.to_dict(), registry=self.type_registry)
        _set_metadata(result, update_time=doc.update_time, create_time=doc.create_time)
        return result

    def find(
        self, model_class: type[T], document_id: str, *, parent: Model | None = None
    ) -> T | None:
        """Fetch a single document by ID, returning None if not found.

        Args:
            model_class: The Model class to deserialize into.
            document_id: Firestore document ID.
            parent: Parent document for subcollection queries.

        Returns:
            The deserialized model instance, or None.
        """
        col_ref = self._get_collection_ref(model_class, parent)
        doc = col_ref.document(document_id).get()
        if not doc.exists:
            return None
        result = deserialize(model_class, doc.id, doc.to_dict(), registry=self.type_registry)
        _set_metadata(result, update_time=doc.update_time, create_time=doc.create_time)
        return result

    def get_many(
        self,
        model_class: type[T],
        document_ids: list[str],
        *,
        parent: Model | None = None,
    ) -> list[T]:
        """Batch fetch multiple documents by ID in a single round trip.

        Args:
            model_class: The Model class to deserialize into.
            document_ids: List of Firestore document IDs.
            parent: Parent document for subcollection queries.

        Returns:
            List of deserialized model instances.

        Raises:
            DocumentNotFoundError: If any documents are missing.
        """
        col_ref = self._get_collection_ref(model_class, parent)
        doc_refs = [col_ref.document(doc_id) for doc_id in document_ids]
        results: list[T] = []
        missing: list[str] = []
        for doc in self._client.get_all(doc_refs):
            if not doc.exists:
                missing.append(doc.id)
            else:
                instance = deserialize(
                    model_class, doc.id, doc.to_dict(), registry=self.type_registry
                )
                _set_metadata(instance, update_time=doc.update_time, create_time=doc.create_time)
                results.append(instance)
        if missing:
            raise DocumentNotFoundError(model_class.__collection__, ", ".join(missing))
        return results

    def select(
        self,
        model_class: type[T],
        *filters: Any,
        order_by: list[Any] | None = None,
        limit: int | None = None,
        start_at: dict[str, Any] | Model | None = None,
        start_after: dict[str, Any] | Model | None = None,
        end_at: dict[str, Any] | Model | None = None,
        end_before: dict[str, Any] | Model | None = None,
        parent: Model | None = None,
    ) -> Query[T]:
        """Query documents. Returns a Query with convenience methods."""
        q = self._build_query(
            model_class,
            filters,
            order_by=order_by,
            limit=limit,
            start_at=start_at,
            start_after=start_after,
            end_at=end_at,
            end_before=end_before,
            parent=parent,
        )
        return Query(q, model_class, self._apply_filter, registry=self.type_registry)

    def select_group(
        self,
        model_class: type[T],
        *filters: Any,
        order_by: list[Any] | None = None,
        limit: int | None = None,
        start_at: dict[str, Any] | Model | None = None,
        start_after: dict[str, Any] | Model | None = None,
        end_at: dict[str, Any] | Model | None = None,
        end_before: dict[str, Any] | Model | None = None,
    ) -> Query[T]:
        """Query across all subcollections with the given collection name."""
        q = self._build_query(
            model_class,
            filters,
            order_by=order_by,
            limit=limit,
            start_at=start_at,
            start_after=start_after,
            end_at=end_at,
            end_before=end_before,
            collection_group=True,
        )
        return Query(q, model_class, self._apply_filter, registry=self.type_registry)

    def save(self, instance: T, *, parent: Model | None = None) -> str:
        """Save (upsert) a document. Returns the document ID.

        Args:
            instance: Model instance to save.
            parent: Parent document for subcollection writes.

        Returns:
            The document ID (auto-generated if instance.id was None).
        """
        validate_required_fields(instance)
        col_ref = self._get_collection_ref(type(instance), parent)
        is_new = instance.id is None
        doc_ref = col_ref.document() if is_new else col_ref.document(instance.id)
        write_result = doc_ref.set(to_dict(instance, by_alias=True, registry=self.type_registry))
        if is_new:
            instance.id = doc_ref.id
        _set_metadata(instance, update_time=write_result.update_time)
        return doc_ref.id  # type: ignore[no-any-return]

    def create(self, instance: T, *, parent: Model | None = None) -> str:
        """Create a document. Raises if it already exists. Returns the document ID.

        Args:
            instance: Model instance to create.
            parent: Parent document for subcollection writes.

        Returns:
            The document ID (auto-generated if instance.id was None).

        Raises:
            DocumentAlreadyExistsError: If the document already exists.
        """
        validate_required_fields(instance)
        col_ref = self._get_collection_ref(type(instance), parent)
        is_new = instance.id is None
        doc_ref = col_ref.document() if is_new else col_ref.document(instance.id)
        try:
            write_result = doc_ref.create(
                to_dict(instance, by_alias=True, registry=self.type_registry)
            )
        except Conflict as e:
            raise DocumentAlreadyExistsError(type(instance).__collection__, doc_ref.id) from e
        if is_new:
            instance.id = doc_ref.id
        _set_metadata(instance, update_time=write_result.update_time)
        return doc_ref.id  # type: ignore[no-any-return]

    @overload
    def delete(
        self,
        instance: Model,
        *,
        parent: Model | None = None,
        if_unchanged: bool | datetime.datetime = False,
    ) -> None: ...
    @overload
    def delete(
        self,
        model_class: type[T],
        doc_id: str,
        *,
        parent: Model | None = None,
        must_exist: bool = False,
        if_unchanged: bool | datetime.datetime = False,
    ) -> None: ...

    def delete(  # type: ignore[misc]
        self,
        instance_or_class: Model | type[T],
        doc_id: str | None = None,
        *,
        parent: Model | None = None,
        must_exist: bool = False,
        if_unchanged: bool | datetime.datetime = False,
    ) -> None:
        """Delete a document by instance or by class + ID.

        Args:
            instance_or_class: A Model instance, or a Model class.
            doc_id: Document ID (required when passing a class).
            parent: Parent document for subcollection deletes.
            must_exist: If True, raise DocumentNotFoundError when the document doesn't exist.
            if_unchanged: If True, only delete if unchanged since last read. Accepts datetime.
        """
        if isinstance(instance_or_class, Model):
            if instance_or_class.id is None:
                raise CendryError("Cannot delete a model instance with id=None")
            option = _resolve_precondition(if_unchanged, instance_or_class)
            col_ref = self._get_collection_ref(type(instance_or_class), parent)
            col_ref.document(instance_or_class.id).delete(option=option)
            _clear_metadata(instance_or_class)
        else:
            if doc_id is None:  # pragma: no cover
                raise CendryError("doc_id is required")  # pragma: no cover
            option = _resolve_precondition(if_unchanged)
            col_ref = self._get_collection_ref(instance_or_class, parent)
            if must_exist:
                doc = col_ref.document(doc_id).get()
                if not doc.exists:
                    raise DocumentNotFoundError(instance_or_class.__collection__, doc_id)
            col_ref.document(doc_id).delete(option=option)

    @overload
    def update(
        self,
        instance: Model,
        field_updates: dict[str, Any],
        *,
        parent: Model | None = None,
        if_unchanged: bool | datetime.datetime = False,
    ) -> None: ...
    @overload
    def update(
        self,
        model_class: type[T],
        doc_id: str,
        field_updates: dict[str, Any],
        *,
        parent: Model | None = None,
        if_unchanged: bool | datetime.datetime = False,
    ) -> None: ...

    def update(  # type: ignore[misc]
        self,
        instance_or_class: Model | type[T],
        field_updates_or_doc_id: dict[str, Any] | str,
        field_updates_or_none: dict[str, Any] | None = None,
        *,
        parent: Model | None = None,
        if_unchanged: bool | datetime.datetime = False,
    ) -> None:
        """Partially update a document's fields.

        Args:
            instance_or_class: A Model instance, or a Model class.
            field_updates_or_doc_id: Field updates dict (instance form) or doc ID (class form).
            field_updates_or_none: Field updates dict (class form only).
            parent: Parent document for subcollection updates.
            if_unchanged: If True, only update if unchanged since last read. Accepts datetime.

        Raises:
            DocumentNotFoundError: If the document does not exist.
        """
        if isinstance(instance_or_class, Model):
            if instance_or_class.id is None:
                raise CendryError("Cannot update a model instance with id=None")
            model_class = type(instance_or_class)
            doc_id = instance_or_class.id
            field_updates: dict[str, Any] = field_updates_or_doc_id  # type: ignore[assignment]
        else:
            model_class = instance_or_class
            if not isinstance(field_updates_or_doc_id, str):  # pragma: no cover
                raise CendryError("doc_id must be a string when calling update with a class")
            doc_id = field_updates_or_doc_id
            if field_updates_or_none is None:  # pragma: no cover
                raise CendryError("field_updates is required when calling update with a class")
            field_updates = field_updates_or_none

        resolved = {
            resolve_field_path(model_class, k): serialize_update_value(
                v, hint=resolve_field_hint(model_class, k), registry=self.type_registry
            )
            for k, v in field_updates.items()
        }
        instance_for_meta = instance_or_class if isinstance(instance_or_class, Model) else None
        option = _resolve_precondition(if_unchanged, instance_for_meta)
        col_ref = self._get_collection_ref(model_class, parent)
        try:
            write_result = col_ref.document(doc_id).update(resolved, option=option)
        except NotFound as e:
            raise DocumentNotFoundError(model_class.__collection__, doc_id) from e
        if instance_for_meta is not None:
            _set_metadata(instance_for_meta, update_time=write_result.update_time)

    def refresh(self, instance: T, *, parent: Model | None = None) -> None:
        """Re-fetch a document from Firestore and update the instance in-place.

        Args:
            instance: Model instance to refresh.
            parent: Parent document for subcollection queries.

        Raises:
            DocumentNotFoundError: If the document does not exist.
        """
        if instance.id is None:
            raise CendryError("Cannot refresh a model instance with id=None")
        col_ref = self._get_collection_ref(type(instance), parent)
        doc = col_ref.document(instance.id).get()
        if not doc.exists:
            raise DocumentNotFoundError(type(instance).__collection__, instance.id)
        fresh = deserialize(type(instance), doc.id, doc.to_dict(), registry=self.type_registry)
        for f in dataclasses.fields(instance):
            setattr(instance, f.name, getattr(fresh, f.name))
        _set_metadata(instance, update_time=doc.update_time, create_time=doc.create_time)

    def on_snapshot(
        self,
        model_class: type[T],
        doc_id: str,
        callback: Callable[..., None],
        *,
        parent: Model | None = None,
    ) -> Any:
        """Listen for real-time updates on a single document.

        The callback receives ``(instance, changes, read_time)`` where
        ``instance`` is the deserialized model (or None if deleted).

        Returns a ``Watch`` object with an ``unsubscribe()`` method.

        Note:
            Sync only — async listeners are not supported by the Firestore SDK.
        """
        col_ref = self._get_collection_ref(model_class, parent)
        doc_ref = col_ref.document(doc_id)
        registry = self.type_registry

        def _wrapper(doc_snapshot: Any, changes: Any, read_time: Any) -> None:
            if doc_snapshot and doc_snapshot[0].exists:
                snap = doc_snapshot[0]
                instance = deserialize(model_class, snap.id, snap.to_dict(), registry=registry)
                _set_metadata(instance, update_time=snap.update_time, create_time=snap.create_time)
                callback(instance, changes, read_time)
            else:
                callback(None, changes, read_time)

        return doc_ref.on_snapshot(_wrapper)

    def batch(self) -> "Batch":
        """Create a batch writer for atomic multi-document operations."""
        from .batch import Batch

        return Batch(self._client.batch(), self._get_collection_ref, self.type_registry)

    def save_many(self, instances: list[T], *, parent: Model | None = None) -> None:
        """Save multiple documents atomically. Max 500 items.

        Args:
            instances: List of Model instances to save.
            parent: Parent document for subcollection writes.

        Raises:
            CendryError: If more than 500 items or validation fails.
        """
        _check_batch_limit(len(instances))
        with self.batch() as b:
            for instance in instances:
                b.save(instance, parent=parent)

    @overload
    def delete_many(self, instances: list[Model], *, parent: Model | None = None) -> None: ...
    @overload
    def delete_many(
        self,
        model_class: type[T],
        doc_ids: list[str],
        *,
        parent: Model | None = None,
    ) -> None: ...

    def delete_many(  # type: ignore[misc]
        self,
        instances_or_class: list[Model] | type[T],
        doc_ids: list[str] | None = None,
        *,
        parent: Model | None = None,
    ) -> None:
        """Delete multiple documents atomically. Max 500 items."""
        if isinstance(instances_or_class, list):
            _check_batch_limit(len(instances_or_class))
            with self.batch() as b:
                for instance in instances_or_class:
                    b.delete(instance, parent=parent)
        else:
            if doc_ids is None:  # pragma: no cover
                raise CendryError("doc_ids is required when calling delete_many with a class")
            _check_batch_limit(len(doc_ids))
            with self.batch() as b:
                for doc_id in doc_ids:
                    b.delete(instances_or_class, doc_id, parent=parent)

    @overload
    def transaction(
        self,
        fn: Callable[["Txn"], Any],
        *,
        max_attempts: int = 5,
        read_only: bool = False,
    ) -> Any: ...
    @overload
    def transaction(
        self,
        fn: None = None,
        *,
        max_attempts: int = 5,
        read_only: bool = False,
    ) -> "Txn": ...

    def transaction(
        self,
        fn: Callable[["Txn"], Any] | None = None,
        *,
        max_attempts: int = 5,
        read_only: bool = False,
    ) -> Any:
        """Run a transaction with callback (auto-retry) or context manager (single attempt).

        Args:
            fn: Callback receiving a Txn. If None, returns a context manager.
            max_attempts: Max retry attempts on contention (callback form only).
            read_only: If True, no writes allowed.
        """
        from google.cloud.firestore_v1.transaction import transactional

        from .transaction import Txn

        fs_txn = self._client.transaction(max_attempts=max_attempts, read_only=read_only)
        txn = Txn(fs_txn, self._get_collection_ref, self.type_registry)
        if fn is None:
            return txn

        @transactional
        def _run(transaction: Any) -> Any:  # pragma: no cover
            return fn(txn)

        return _run(fs_txn)  # pragma: no cover


class AsyncCendry(_BaseCendry):
    """Asynchronous Firestore ODM context.

    Works with anyio (asyncio + trio).

    Args:
        client: Optional async Firestore Client.
        type_registry: Optional TypeRegistry override.

    Note:
        `select()` and `select_group()` are regular `def` methods — they return
        `AsyncQuery` synchronously. Only `get`, `find`, `get_many` are `async def`.
    """

    def __init__(
        self,
        *,
        client: Any = None,
        type_registry: TypeRegistry | None = None,
    ) -> None:
        if client is None:
            from google.cloud.firestore import AsyncClient

            client = AsyncClient()
        self._client = client
        self.type_registry = type_registry or default_registry

    async def __aenter__(self) -> Self:
        return self

    async def __aexit__(self, *args: object) -> None:
        await self._client.close()

    async def get(
        self, model_class: type[T], document_id: str, *, parent: Model | None = None
    ) -> T:
        """Get a document by ID. Raises DocumentNotFoundError if it doesn't exist."""
        col_ref = self._get_collection_ref(model_class, parent)
        doc = await col_ref.document(document_id).get()
        if not doc.exists:
            raise DocumentNotFoundError(model_class.__collection__, document_id)
        result = deserialize(model_class, doc.id, doc.to_dict(), registry=self.type_registry)
        _set_metadata(result, update_time=doc.update_time, create_time=doc.create_time)
        return result

    async def find(
        self, model_class: type[T], document_id: str, *, parent: Model | None = None
    ) -> T | None:
        """Get a document by ID. Returns None if it doesn't exist."""
        col_ref = self._get_collection_ref(model_class, parent)
        doc = await col_ref.document(document_id).get()
        if not doc.exists:
            return None
        result = deserialize(model_class, doc.id, doc.to_dict(), registry=self.type_registry)
        _set_metadata(result, update_time=doc.update_time, create_time=doc.create_time)
        return result

    async def get_many(
        self,
        model_class: type[T],
        document_ids: list[str],
        *,
        parent: Model | None = None,
    ) -> list[T]:
        """Batch fetch multiple documents by ID in a single round trip.

        Args:
            model_class: The Model class to deserialize into.
            document_ids: List of Firestore document IDs.
            parent: Parent document for subcollection queries.

        Returns:
            List of deserialized model instances.

        Raises:
            DocumentNotFoundError: If any documents are missing.
        """
        col_ref = self._get_collection_ref(model_class, parent)
        doc_refs = [col_ref.document(doc_id) for doc_id in document_ids]
        results: list[T] = []
        missing: list[str] = []
        async for doc in self._client.get_all(doc_refs):
            if not doc.exists:
                missing.append(doc.id)
            else:
                instance = deserialize(
                    model_class, doc.id, doc.to_dict(), registry=self.type_registry
                )
                _set_metadata(instance, update_time=doc.update_time, create_time=doc.create_time)
                results.append(instance)
        if missing:
            raise DocumentNotFoundError(model_class.__collection__, ", ".join(missing))
        return results

    def select(
        self,
        model_class: type[T],
        *filters: Any,
        order_by: list[Any] | None = None,
        limit: int | None = None,
        start_at: dict[str, Any] | Model | None = None,
        start_after: dict[str, Any] | Model | None = None,
        end_at: dict[str, Any] | Model | None = None,
        end_before: dict[str, Any] | Model | None = None,
        parent: Model | None = None,
    ) -> AsyncQuery[T]:
        """Query documents. Returns an AsyncQuery with convenience methods."""
        q = self._build_query(
            model_class,
            filters,
            order_by=order_by,
            limit=limit,
            start_at=start_at,
            start_after=start_after,
            end_at=end_at,
            end_before=end_before,
            parent=parent,
        )
        return AsyncQuery(q, model_class, self._apply_filter, registry=self.type_registry)

    def select_group(
        self,
        model_class: type[T],
        *filters: Any,
        order_by: list[Any] | None = None,
        limit: int | None = None,
        start_at: dict[str, Any] | Model | None = None,
        start_after: dict[str, Any] | Model | None = None,
        end_at: dict[str, Any] | Model | None = None,
        end_before: dict[str, Any] | Model | None = None,
    ) -> AsyncQuery[T]:
        """Query across all subcollections with the given collection name."""
        q = self._build_query(
            model_class,
            filters,
            order_by=order_by,
            limit=limit,
            start_at=start_at,
            start_after=start_after,
            end_at=end_at,
            end_before=end_before,
            collection_group=True,
        )
        return AsyncQuery(q, model_class, self._apply_filter, registry=self.type_registry)

    async def save(self, instance: T, *, parent: Model | None = None) -> str:
        """Save (upsert) a document. Returns the document ID."""
        validate_required_fields(instance)
        col_ref = self._get_collection_ref(type(instance), parent)
        is_new = instance.id is None
        doc_ref = col_ref.document() if is_new else col_ref.document(instance.id)
        write_result = await doc_ref.set(
            to_dict(instance, by_alias=True, registry=self.type_registry)
        )
        if is_new:
            instance.id = doc_ref.id
        _set_metadata(instance, update_time=write_result.update_time)
        return doc_ref.id  # type: ignore[no-any-return]

    async def create(self, instance: T, *, parent: Model | None = None) -> str:
        """Create a document. Raises if it already exists. Returns the document ID."""
        validate_required_fields(instance)
        col_ref = self._get_collection_ref(type(instance), parent)
        is_new = instance.id is None
        doc_ref = col_ref.document() if is_new else col_ref.document(instance.id)
        try:
            write_result = await doc_ref.create(
                to_dict(instance, by_alias=True, registry=self.type_registry)
            )
        except Conflict as e:
            raise DocumentAlreadyExistsError(type(instance).__collection__, doc_ref.id) from e
        if is_new:
            instance.id = doc_ref.id
        _set_metadata(instance, update_time=write_result.update_time)
        return doc_ref.id  # type: ignore[no-any-return]

    @overload
    async def delete(
        self,
        instance: Model,
        *,
        parent: Model | None = None,
        if_unchanged: bool | datetime.datetime = False,
    ) -> None: ...
    @overload
    async def delete(
        self,
        model_class: type[T],
        doc_id: str,
        *,
        parent: Model | None = None,
        must_exist: bool = False,
        if_unchanged: bool | datetime.datetime = False,
    ) -> None: ...

    async def delete(  # type: ignore[misc]
        self,
        instance_or_class: Model | type[T],
        doc_id: str | None = None,
        *,
        parent: Model | None = None,
        must_exist: bool = False,
        if_unchanged: bool | datetime.datetime = False,
    ) -> None:
        """Delete a document by instance or by class + ID."""
        if isinstance(instance_or_class, Model):
            if instance_or_class.id is None:
                raise CendryError("Cannot delete a model instance with id=None")
            option = _resolve_precondition(if_unchanged, instance_or_class)
            col_ref = self._get_collection_ref(type(instance_or_class), parent)
            await col_ref.document(instance_or_class.id).delete(option=option)
            _clear_metadata(instance_or_class)
        else:
            if doc_id is None:  # pragma: no cover
                raise CendryError("doc_id is required")  # pragma: no cover
            option = _resolve_precondition(if_unchanged)
            col_ref = self._get_collection_ref(instance_or_class, parent)
            if must_exist:
                doc = await col_ref.document(doc_id).get()
                if not doc.exists:
                    raise DocumentNotFoundError(instance_or_class.__collection__, doc_id)
            await col_ref.document(doc_id).delete(option=option)

    @overload
    async def update(
        self,
        instance: Model,
        field_updates: dict[str, Any],
        *,
        parent: Model | None = None,
        if_unchanged: bool | datetime.datetime = False,
    ) -> None: ...
    @overload
    async def update(
        self,
        model_class: type[T],
        doc_id: str,
        field_updates: dict[str, Any],
        *,
        parent: Model | None = None,
        if_unchanged: bool | datetime.datetime = False,
    ) -> None: ...

    async def update(  # type: ignore[misc]
        self,
        instance_or_class: Model | type[T],
        field_updates_or_doc_id: dict[str, Any] | str,
        field_updates_or_none: dict[str, Any] | None = None,
        *,
        parent: Model | None = None,
        if_unchanged: bool | datetime.datetime = False,
    ) -> None:
        """Partially update a document's fields."""
        if isinstance(instance_or_class, Model):
            if instance_or_class.id is None:
                raise CendryError("Cannot update a model instance with id=None")
            model_class = type(instance_or_class)
            doc_id = instance_or_class.id
            field_updates: dict[str, Any] = field_updates_or_doc_id  # type: ignore[assignment]
        else:
            model_class = instance_or_class
            if not isinstance(field_updates_or_doc_id, str):  # pragma: no cover
                raise CendryError("doc_id must be a string when calling update with a class")
            doc_id = field_updates_or_doc_id
            if field_updates_or_none is None:  # pragma: no cover
                raise CendryError("field_updates is required when calling update with a class")
            field_updates = field_updates_or_none

        resolved = {
            resolve_field_path(model_class, k): serialize_update_value(
                v, hint=resolve_field_hint(model_class, k), registry=self.type_registry
            )
            for k, v in field_updates.items()
        }
        instance_for_meta = instance_or_class if isinstance(instance_or_class, Model) else None
        option = _resolve_precondition(if_unchanged, instance_for_meta)
        col_ref = self._get_collection_ref(model_class, parent)
        try:
            write_result = await col_ref.document(doc_id).update(resolved, option=option)
        except NotFound as e:
            raise DocumentNotFoundError(model_class.__collection__, doc_id) from e
        if instance_for_meta is not None:
            _set_metadata(instance_for_meta, update_time=write_result.update_time)

    async def refresh(self, instance: T, *, parent: Model | None = None) -> None:
        """Re-fetch a document from Firestore and update the instance in-place."""
        if instance.id is None:
            raise CendryError("Cannot refresh a model instance with id=None")
        col_ref = self._get_collection_ref(type(instance), parent)
        doc = await col_ref.document(instance.id).get()
        if not doc.exists:
            raise DocumentNotFoundError(type(instance).__collection__, instance.id)
        fresh = deserialize(type(instance), doc.id, doc.to_dict(), registry=self.type_registry)
        for f in dataclasses.fields(instance):
            setattr(instance, f.name, getattr(fresh, f.name))
        _set_metadata(instance, update_time=doc.update_time, create_time=doc.create_time)

    def batch(self) -> "AsyncBatch":
        """Create an async batch writer for atomic multi-document operations."""
        from .batch import AsyncBatch

        return AsyncBatch(self._client.batch(), self._get_collection_ref, self.type_registry)

    async def save_many(self, instances: list[T], *, parent: Model | None = None) -> None:
        """Save multiple documents atomically. Max 500 items."""
        _check_batch_limit(len(instances))
        async with self.batch() as b:
            for instance in instances:
                b.save(instance, parent=parent)

    @overload
    async def delete_many(
        self, instances: list[Model], *, parent: Model | None = None
    ) -> None: ...
    @overload
    async def delete_many(
        self,
        model_class: type[T],
        doc_ids: list[str],
        *,
        parent: Model | None = None,
    ) -> None: ...

    async def delete_many(  # type: ignore[misc]
        self,
        instances_or_class: list[Model] | type[T],
        doc_ids: list[str] | None = None,
        *,
        parent: Model | None = None,
    ) -> None:
        """Delete multiple documents atomically. Max 500 items."""
        if isinstance(instances_or_class, list):
            _check_batch_limit(len(instances_or_class))
            async with self.batch() as b:
                for instance in instances_or_class:
                    b.delete(instance, parent=parent)
        else:
            if doc_ids is None:  # pragma: no cover
                raise CendryError("doc_ids is required when calling delete_many with a class")
            _check_batch_limit(len(doc_ids))
            async with self.batch() as b:
                for doc_id in doc_ids:
                    b.delete(instances_or_class, doc_id, parent=parent)

    @overload
    async def transaction(
        self,
        fn: Callable[["AsyncTxn"], Any],
        *,
        max_attempts: int = 5,
        read_only: bool = False,
    ) -> Any: ...
    @overload
    def transaction(
        self,
        fn: None = None,
        *,
        max_attempts: int = 5,
        read_only: bool = False,
    ) -> "AsyncTxn": ...

    def transaction(
        self,
        fn: Callable[["AsyncTxn"], Any] | None = None,
        *,
        max_attempts: int = 5,
        read_only: bool = False,
    ) -> Any:
        """Run a transaction with callback (auto-retry) or context manager (single attempt).

        Args:
            fn: Async callback receiving an AsyncTxn. If None, returns a context manager.
            max_attempts: Max retry attempts on contention (callback form only).
            read_only: If True, no writes allowed.
        """
        from .transaction import AsyncTxn

        fs_txn = self._client.transaction(max_attempts=max_attempts, read_only=read_only)
        txn = AsyncTxn(fs_txn, self._get_collection_ref, self.type_registry)
        if fn is None:
            return txn

        from google.cloud.firestore_v1.async_transaction import (
            async_transactional,  # pragma: no cover
        )

        @async_transactional  # pragma: no cover
        async def _run(transaction: Any) -> Any:  # pragma: no cover
            return await fn(txn)  # pragma: no cover

        async def _execute() -> Any:  # pragma: no cover
            return await _run(fs_txn)  # pragma: no cover

        return _execute()  # pragma: no cover
