"""
Response handler for PostgreSQL query results.

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

from typing import Any, Dict, List, Optional, TypeVar, Generic, Union
from sqlalchemy.orm import Query

T = TypeVar("T")


class PostgresResponse(Generic[T]):
    """
    Wrapper for PostgreSQL/SQLAlchemy query results.

    Attributes:
        query: SQLAlchemy query object
        first: Whether to return first result only
        data: Query results (lazy loaded)
        count: Total count of results

    Properties:
        all: All results as list
        first_item: First result only
    """

    def __init__(
        self,
        query: Query,
        first: bool = False,
        status: bool = True,
        message: str = "",
        error: Optional[str] = None,
    ):
        self._query = query
        self._first = first
        self.status = status
        self.message = message
        self.error = error
        self._data: Optional[Union[List[T], T]] = None
        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:
            results = self._query.all()
            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:
            self._count = self._query.count()
        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,
        }