"""
QuerySet wrapper for mongoengine read-write router.

Wraps mongoengine's QuerySet to implement read routing based on context.
"""

import mongoengine
from mongoengine import QuerySet as BaseQuerySet
from pymongo import read_preferences

from .connections import get_read_db_alias, get_write_db_alias
from .context import RequestContext, get_user_id


class RwQuerySet(BaseQuerySet):
    """
    QuerySet wrapper that routes reads to replicas based on context.

    This QuerySet automatically routes read operations to the appropriate
    database (primary or replica) based on the current routing strategy
    and request context.
    """

    def __init__(self, *args, force_primary: bool = False, **kwargs):
        """
        Initialize the RwQuerySet.

        Args:
            *args: Arguments to pass to BaseQuerySet
            force_primary: If True, force all reads to go to primary
            **kwargs: Keyword arguments to pass to BaseQuerySet
        """
        super().__init__(*args, **kwargs)
        self._force_primary = force_primary

    def clone(self, *args, **kwargs):
        """
        Clone the QuerySet, preserving the force_primary flag.

        Returns:
            A new RwQuerySet instance with the same configuration
        """
        force_primary = kwargs.pop("force_primary", self._force_primary)
        new_qs = self.__class__(
            self._document, self._collection_obj, force_primary=force_primary
        )
        return self._clone_into(new_qs)

    def _get_db_alias(self) -> str:
        """
        Get the appropriate database alias for this query.

        Returns:
            The database alias to use for the query
        """
        if self._force_primary:
            return get_write_db_alias()

        # Get user_id from context for consistent hashing
        user_id = get_user_id()
        return get_read_db_alias(user_id=user_id)

    def _cursor_args(self):
        """
        Get cursor arguments, setting appropriate read preference.

        This ensures that when we're querying a replica, we use
        the appropriate read preference.
        """
        args = super()._cursor_args()

        # Get the target alias to determine read preference
        alias = self._get_db_alias()
        # If targeting a replica, set secondary preferred read preference
        if alias != "default":
            args["read_preference"] = read_preferences.ReadPreference.SECONDARY_PREFERRED

        return args

    # Override methods that execute database reads
    # We need to execute these with the appropriate database alias

    def __iter__(self):
        """Iterate over the queryset results."""
        alias = self._get_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().__iter__()

    def count(self, *args, **kwargs):
        """Count the number of matching documents."""
        alias = self._get_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().count(*args, **kwargs)

    def first(self, *args, **kwargs):
        """Get the first matching document."""
        alias = self._get_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().first(*args, **kwargs)

    def last(self, *args, **kwargs):
        """Get the last matching document."""
        alias = self._get_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().last(*args, **kwargs)

    def get(self, *args, **kwargs):
        """Get a single matching document, raising DoesNotExist if not found."""
        alias = self._get_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().get(*args, **kwargs)

    def exists(self, *args, **kwargs):
        """Check if any matching documents exist."""
        alias = self._get_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().exists(*args, **kwargs)

    def limit(self, *args, **kwargs):
        """Limit the number of results."""
        alias = self._get_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().limit(*args, **kwargs)

    def skip(self, *args, **kwargs):
        """Skip a number of results."""
        alias = self._get_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().skip(*args, **kwargs)

    def distinct(self, *args, **kwargs):
        """Get distinct values for a field."""
        alias = self._get_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().distinct(*args, **kwargs)

    def only(self, *fields, **kwargs):
        """Select only specific fields."""
        alias = self._get_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().only(*fields, **kwargs)

    def exclude(self, *fields, **kwargs):
        """Exclude specific fields."""
        alias = self._get_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().exclude(*fields, **kwargs)

    def aggregate(self, *pipeline, **kwargs):
        """Perform aggregation."""
        alias = self._get_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().aggregate(*pipeline, **kwargs)

    def item_frequencies(self, *args, **kwargs):
        """Get item frequencies."""
        alias = self._get_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().item_frequencies(*args, **kwargs)

    def __len__(self):
        """Get the length of the queryset (count)."""
        return self.count()

    def __getitem__(self, key):
        """Support list-like indexing."""
        alias = self._get_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().__getitem__(key)

    def delete(self, *args, **kwargs):
        """
        Delete matching documents.

        This is a write operation, so it always goes to primary.
        """
        RequestContext.mark_write()
        alias = get_write_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().delete(*args, **kwargs)

    def update(self, *args, **kwargs):
        """
        Update matching documents.

        This is a write operation, so it always goes to primary.
        """
        RequestContext.mark_write()
        alias = get_write_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().update(*args, **kwargs)

    def update_one(self, *args, **kwargs):
        """
        Update a single matching document.

        This is a write operation, so it always goes to primary.
        """
        RequestContext.mark_write()
        alias = get_write_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().update_one(*args, **kwargs)

    def upsert_one(self, *args, **kwargs):
        """
        Upsert a single document.

        This is a write operation, so it always goes to primary.
        """
        RequestContext.mark_write()
        alias = get_write_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().upsert_one(*args, **kwargs)

    def modify(self, *args, **kwargs):
        """
        Modify and return a single document.

        This is a write operation, so it always goes to primary.
        """
        RequestContext.mark_write()
        alias = get_write_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().modify(*args, **kwargs)

    def create(self, *args, **kwargs):
        """
        Create a new document.

        This is a write operation, so it always goes to primary.
        """
        RequestContext.mark_write()
        alias = get_write_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().create(*args, **kwargs)

    def bulk_create(self, *args, **kwargs):
        """
        Create multiple documents.

        This is a write operation, so it always goes to primary.
        """
        RequestContext.mark_write()
        alias = get_write_db_alias()
        with mongoengine.switch_db(self._document, alias):
            return super().bulk_create(*args, **kwargs)

    # Method to force primary reads
    def force_primary(self) -> "RwQuerySet":
        """
        Return a new QuerySet that forces reads from primary.

        Returns:
            A new RwQuerySet with force_primary enabled
        """
        return self.clone(force_primary=True)
