"""
Response handler for MongoDB query results.

This module provides a wrapper class for MongoDB query results,
adding convenience methods for accessing data and managing query state.
"""

from typing import Any, Dict, List, Optional, TypeVar, Generic, Union
from pymongo.cursor import Cursor

T = TypeVar("T")


class MongoResponse(Generic[T]):
    """
    Wrapper for MongoDB query results.

    Attributes:
        cursor: MongoDB cursor object
        first: Whether to return first result only
        data: Query results (lazy loaded)
        count: Total count of results
    """

    def __init__(
        self,
        cursor: Optional[Cursor] = None,
        first: bool = False,
        status: bool = True,
        message: str = "",
        error: Optional[str] = None,
        data: Optional[Union[List[T], T]] = None,
    ):
        self._cursor = cursor
        self._first = first
        self.status = status
        self.message = message
        self.error = error
        self._data: Optional[Union[List[T], T]] = data
        self._count: Optional[int] = None

    @property
    def data(self) -> Union[List[T], T, None]:
        """
        Lazy load and return query results.
        Returns first item if first=True, otherwise returns all results.
        """
        if self._data is None and self._cursor is not None:
            results = list(self._cursor)
            self._data = results[0] if self._first and results else results
        return self._data

    @property
    def count(self) -> int:
        """Lazy load and return total count of results."""
        if self._count is None:
            if self._cursor is not None:
                self._count = self._cursor.count()
            else:
                self._count = len(self.all)
        return self._count

    @property
    def all(self) -> List[T]:
        """Get all results as list."""
        return (
            self.data
            if isinstance(self.data, list)
            else [self.data] if self.data else []
        )

    @property
    def first(self) -> Optional[T]:
        """Get first result only."""
        return self.data if self._first else (self.data[0] if self.data else None)

    def as_dict(self) -> Dict[str, Any]:
        """Convert response to dictionary format."""
        return {
            "status": self.status,
            "message": self.message,
            "data": self.data,
            "count": self.count,
            "error": self.error,
        }