"""
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, Optional, TypeVar, Generic, Union

from pydantic import BaseModel
from sqlalchemy.orm import Query


T = TypeVar("T")


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

    Attributes:
        metadata: Additional metadata for the query

    Properties:
        count: Total count of results
        query: Get query object
        as_dict: Convert response to dictionary format
    """

    def __init__(
        self,
        pre_query: Query,
        query: Query,
        model,
        is_array: bool = True,
        metadata: Any = None,
    ):
        self._core_class = model
        self._is_list = is_array
        self._query = query
        self._pre_query = pre_query
        self._count: Optional[int] = None
        self.metadata = metadata

    @property
    def core_class(self):
        """Get query object."""
        return self._core_class

    @property
    def data(self) -> Union[T, list[T]]:
        """Get query results."""
        if not self.is_list:
            first_item = self._query.first()
            return first_item if first_item else None
        return self._query.all() if self._query.all() else []

    @property
    def data_as_dict(self) -> Union[Dict[str, Any], list[Dict[str, Any]]] | None:
        """Get query results as dictionary."""
        if not self.count:
            return None
        if self.is_list:
            first_item = self._query.first()
            return first_item.get_dict() if first_item.first() else None
        all_items = self._query.all()
        return [result.get_dict() for result in all_items] if all_items else []

    @property
    def total_count(self) -> int:
        """Lazy load and return total count of results."""
        return self._pre_query.count() if self._pre_query else 0

    @property
    def count(self) -> int:
        """Lazy load and return total count of results."""
        return self._query.count()

    @property
    def core_query(self) -> str:
        """Get query object."""
        return str(self._query)

    @property
    def query(self) -> Query:
        """Get query object."""
        return self._query

    @property
    def is_list(self) -> bool:
        """Check if response is a list."""
        return self._is_list

    @property
    def as_dict(self) -> Dict[str, Any]:
        """Convert response to dictionary format."""
        if isinstance(self.data, list):
            return {
                "metadata": self.metadata,
                "is_list": self._is_list,
                "query": str(self.query),
                "count": self.count,
                "data": [result.get_dict() for result in self.data],
            }
        return {
            "metadata": self.metadata,
            "is_list": self._is_list,
            "query": str(self.query),
            "count": self.count,
            "data": self.data.get_dict() if self.data else {},
        }


class EndpointResponse(BaseModel):
    """Endpoint response model."""

    completed: bool = True
    message: str = "Success"
    pagination_result: Any

    @property
    def response(self):
        """Convert response to dictionary format."""
        resutl_data = getattr(self.pagination_result, "data", None)
        result_pagination = getattr(self.pagination_result, "pagination", None)
        if not result_pagination:
            raise ValueError("Invalid pagination result pagination.")
        pagination_dict = getattr(result_pagination, "as_dict", None)
        if not pagination_dict:
            raise ValueError("Invalid pagination result as_dict.")
        return {
            "completed": self.completed,
            "message": self.message,
            "data": resutl_data,
            "pagination": pagination_dict,
        }