from typing import Any, List, Optional, TypeVar, Union

from sqlalchemy.orm import Query
from sqlalchemy.orm.session import Session

T = TypeVar("T")


class QueryResponse:
    """Handler for SQLAlchemy query results with error handling."""

    def __init__(self, db: Session, query: Query, first: bool = False):
        self.db = db
        self.first = first
        self.__query = query

    def get(self, index: int) -> Optional[T]:
        """Get item at specific index if it exists."""
        count = self.count
        if count and not index > count:
            return self.data[index - 1]
        return None

    @property
    def data(self) -> Union[Optional[T], List[T]]:
        """Get query results with error handling."""
        try:
            if self.first:
                return self.__query.first()
            return self.__query.all()
        except Exception as e:
            # Handle any database errors by rolling back
            self.db.rollback()
            return None if self.first else []

    @property
    def count(self) -> int:
        """Get total count of query results."""
        return self.__query.count()

    @property
    def query(self) -> Query:
        """Get the underlying SQLAlchemy query."""
        return self.__query