from sqlmodel import Session, func, text, SQLModel, select
from typing import Any, List, Literal, Optional, Type, TypeVar, Generic

ReadType = TypeVar("ReadType", bound=SQLModel)
CreateType = TypeVar("CreateType", bound=SQLModel)
UpdateType = TypeVar("UpdateType", bound=SQLModel)


class BaseRepository(Generic[ReadType, CreateType, UpdateType]):
    entity: Type[ReadType]

    def __init__(self, db: Session):
        self.db = db

    def get_entity(self, *args: Any, **kwargs: Any) -> Optional[ReadType]:
        result = self.db.exec(
            select(self.entity)
            .filter(*args)
            .filter_by(**kwargs)
        )

        return result.first()

    def get_entities(self, *args: Any, offset: int = 0, limit: int = 100, order_by: str = 'id', order: Literal['desc', 'asc'] = 'asc', **kwargs: Any) -> List[ReadType]:
        result = self.db.exec(
            select(self.entity)
            .filter(*args)
            .filter_by(**kwargs)
            .offset(offset)
            .limit(limit)
            .order_by(text(f"{order_by} {order}"))
        )

        return result.all()

    def count_entities(self) -> int:
        return self.db.query(