"""
Document mixin for mongoengine read-write router.

Provides Document classes with read-write routing capabilities.
"""

from typing import TypeVar

import mongoengine
from mongoengine import Document
from mongoengine.queryset.manager import QuerySetManager

from .context import RequestContext
from .queryset import RwQuerySet

T = TypeVar("T", bound=Document)


class RwDocumentMixin:
    """
    Mixin for mongoengine Documents with read-write routing.

    This mixin provides:
    - Automatic routing of reads to replicas based on context
    - Writes always go to primary
    - Context tracking for write operations
    - Helper methods for forcing primary reads

    Usage:
        from mongoengine import StringField, IntField
        from mongoengine_rw_router import RwDocumentMixin

        class User(RwDocumentMixin, Document):
            name = StringField()
            age = IntField()

            meta = {
                "abstract": False,
                "queryset_class": RwQuerySet,
            }
    """

    # This will be set by the meta class when used with Document
    _queryset_class: type[RwQuerySet] = RwQuerySet

    @classmethod
    def primary_queryset(cls) -> RwQuerySet:
        """
        Get a QuerySet that forces reads from primary database.

        This is useful when you need to ensure fresh data for critical operations.

        Returns:
            A QuerySet configured to read from primary

        Example:
            user = User.primary_queryset().first()  # Always reads from primary
        """
        qs = cls._get_queryset()
        return qs.force_primary()

    def save(self, *args, **kwargs):
        """
        Save the document to the database.

        This is a write operation, so it:
        1. Always goes to the primary database
        2. Marks the context as having write operations

        Returns:
            The saved document
        """
        RequestContext.mark_write()
        return super().save(*args, **kwargs)

    def delete(self, *args, **kwargs):
        """
        Delete the document from the database.

        This is a write operation, so it:
        1. Always goes to the primary database
        2. Marks the context as having write operations
        """
        RequestContext.mark_write()
        return super().delete(*args, **kwargs)

    def reload(self, *args, **kwargs):
        """
        Reload the document from the database.

        This operation respects the routing configuration:
        - If force_primary context is active, reads from primary
        - If read_after_write_consistency is enabled and writes occurred, reads from primary
        - Otherwise routes based on configured strategy
        """
        from .connections import get_read_db_alias
        from .context import get_user_id

        alias = kwargs.pop("force_primary_alias", None)
        if alias is None:
            alias = get_read_db_alias(user_id=get_user_id())
        with mongoengine.switch_db(type(self), alias):
            return super().reload(*args, **kwargs)

    @classmethod
    def _get_queryset(cls) -> RwQuerySet:
        """Get the QuerySet for this document class."""
        return cls._queryset_class(cls, cls._get_collection())


class RwManager(QuerySetManager):
    """
    Manager for mongoengine Documents with read-write routing.

    This manager uses RwQuerySet for all query operations.
    Use as: objects = RwManager()
    """

    default = RwQuerySet
    queryset_class = RwQuerySet
