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

Provides Document classes with read-write routing capabilities.
"""
from typing import TypeVar, Generic

import mongoengine
from mongoengine import Document

from .connections import get_write_db_alias, force_primary
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
        """
        alias = kwargs.pop("force_primary_alias", None)
        if alias:
            with mongoengine.switch_db(type(self), alias):
                return super().reload(*args, **kwargs)
        return super().reload(*args, **kwargs)

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


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

    This manager uses RwQuerySet for all query operations.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._queryset_class = RwQuerySet

    def get_queryset(self) -> RwQuerySet:
        """Get the QuerySet for this manager."""
        return self._queryset_class(self._document)

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

        Returns:
            A QuerySet configured to read from primary
        """
        return self.get_queryset().force_primary()

    def __getattr__(self, name):
        """Proxy attribute access to the QuerySet."""
        return getattr(self.get_queryset(), name)
