from datetime import datetime
from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union
from uuid import UUID
from fastapi_pagination.ext.async_sqlmodel import paginate
from fastapi_pagination import Params, Page
from fastapi.encoders import jsonable_encoder
from pydantic import BaseModel
from sqlmodel import SQLModel, select, func
from sqlalchemy.orm import selectinload
from sqlmodel.ext.asyncio.session import AsyncSession
from sqlmodel.sql.expression import Select, SelectOfScalar

ModelType = TypeVar("ModelType", bound=SQLModel)
CreateSchemaType = TypeVar("CreateSchemaType", bound=BaseModel)
UpdateSchemaType = TypeVar("UpdateSchemaType", bound=BaseModel)
SchemaType = TypeVar("SchemaType", bound=BaseModel)
T = TypeVar("T", bound=SQLModel)

class CRUDBase(Generic[ModelType, CreateSchemaType, UpdateSchemaType]):
    def __init__(self, model: Type[ModelType]):
        """
        CRUD object with default methods to Create, Read, Update, Delete (CRUD).
        **Parameters**
        * `model`: A SQLAlchemy model class
        * `schema`: A Pydantic model (schema) class
        """
        self.model = model

    async def get(
        self, db_session: AsyncSession, *, id: Union[UUID, str]
    ) -> Optional[ModelType]:
        response = await db_session.exec(select(self.model).where(self.model.id == id).options(selectinload('*')))
        return response.first()

    async def get_by_ids(self, db_session: AsyncSession, list_ids: List[Union[UUID, str]],) -> Optional[List[ModelType]]:
        response = await db_session.exec(select(self.model).where(self.model.id.in_(list_ids)))
        return response.all()

    async def get_count(
        self, db_session: AsyncSession
    ) -> Optional[ModelType]:
        response = await db_session.exec(select(func.count()).select_from(select(self.model).subquery()))
        return response.one()

    async def get_multi(
        self, db_session: AsyncSession, *, skip: int = 0, limit: int = 100
    ) -> List[ModelType]:
        response = await db_session.exec(
            select(self.model).offset(skip).limit(limit).order_by(self.model.id)
        )
        return response.all()

    async def get_multi_paginated(
        self, db_session: AsyncSession, *, params: Optional[Params] = Params(), query: Optional[Union[T, Select[T], SelectOfScalar[T]]] = None
    ) -> Page[ModelType]:
        if query == None:
            query = self.model
        return await paginate(db_session, query, params)


    async def create(
        self, db_session: AsyncSession, *, obj_in: Union[CreateSchemaType, ModelType], created_by_id: Optional[Union[UUID, str]] = None
    ) -> ModelType:
        db_obj = self.model.from_orm(obj_in)  # type: ignore
        db_obj.created_at = datetime.utcnow()
        db_obj.updated_at = datetime.utcnow()
        if(created_by_id):
            db_obj.created_by_id = created_by_id        

        db_session.add(db_obj)
        await db_session.commit()
        await db_session.refresh(db_obj)
        return db_obj

    async def update(
        self,
        db_session: AsyncSession,
        *,
        obj_current: ModelType,
        obj_new: Union[UpdateSchemaType, Dict[str, Any], ModelType]
    ) -> ModelType:
        obj_data = jsonable_encoder(obj_current)

        if isinstance(obj_new, dict):
            update_data = obj_new
        else:
            update_data = obj_new.dict(exclude_unset=True) #This tells Pydantic to not include the values that were not sent
        for field in obj_data:
            if field in update_data:
                setattr(obj_current, field, update_data[field])
            if field == "updated_at":
                setattr(obj_current, field, datetime.utcnow())

        db_session.add(obj_current)
        await db_session.commit()
        await db_session.refresh(obj_current)
        return obj_current

    async def remove(
        self, db_session: AsyncSession, *, id: Union[UUID, str]
    ) -> ModelType:
        response = await db_session.exec(select(self.model).where(self.model.id == id))
        obj = response.one()
        await db_session.delete(obj)
        await db_session.commit()
        return obj

fastapi.encoders.jsonable_encoder sqlalchemy.orm.selectinload sqlmodel.select fastapi_pagination.ext.async_sqlmodel.paginate sqlmodel.func.count datetime.datetime.utcnow fastapi_pagination.Params typing.TypeVar
[((518, 554), 'typing.TypeVar', 'TypeVar', (['"""ModelType"""'], {'bound': 'SQLModel'}), "('ModelType', bound=SQLModel)\n", (525, 554), False, 'from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union\n'), ((574, 618), 'typing.TypeVar', 'TypeVar', (['"""CreateSchemaType"""'], {'bound': 'BaseModel'}), "('CreateSchemaType', bound=BaseModel)\n", (581, 618), False, 'from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union\n'), ((638, 682), 'typing.TypeVar', 'TypeVar', (['"""UpdateSchemaType"""'], {'bound': 'BaseModel'}), "('UpdateSchemaType', bound=BaseModel)\n", (645, 682), False, 'from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union\n'), ((696, 734), 'typing.TypeVar', 'TypeVar', (['"""SchemaType"""'], {'bound': 'BaseModel'}), "('SchemaType', bound=BaseModel)\n", (703, 734), False, 'from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union\n'), ((739, 767), 'typing.TypeVar', 'TypeVar', (['"""T"""'], {'bound': 'SQLModel'}), "('T', bound=SQLModel)\n", (746, 767), False, 'from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union\n'), ((2277, 2285), 'fastapi_pagination.Params', 'Params', ([], {}), '()\n', (2283, 2285), False, 'from fastapi_pagination import Params, Page\n'), ((2759, 2776), 'datetime.datetime.utcnow', 'datetime.utcnow', ([], {}), '()\n', (2774, 2776), False, 'from datetime import datetime\n'), ((2805, 2822), 'datetime.datetime.utcnow', 'datetime.utcnow', ([], {}), '()\n', (2820, 2822), False, 'from datetime import datetime\n'), ((3257, 3286), 'fastapi.encoders.jsonable_encoder', 'jsonable_encoder', (['obj_current'], {}), '(obj_current)\n', (3273, 3286), False, 'from fastapi.encoders import jsonable_encoder\n'), ((2454, 2489), 'fastapi_pagination.ext.async_sqlmodel.paginate', 'paginate', (['db_session', 'query', 'params'], {}), '(db_session, query, params)\n', (2462, 2489), False, 'from fastapi_pagination.ext.async_sqlmodel import paginate\n'), ((1349, 1366), 'sqlalchemy.orm.selectinload', 'selectinload', (['"""*"""'], {}), "('*')\n", (1361, 1366), False, 'from sqlalchemy.orm import selectinload\n'), ((3713, 3730), 'datetime.datetime.utcnow', 'datetime.utcnow', ([], {}), '()\n', (3728, 3730), False, 'from datetime import datetime\n'), ((1565, 1583), 'sqlmodel.select', 'select', (['self.model'], {}), '(self.model)\n', (1571, 1583), False, 'from sqlmodel import SQLModel, select, func\n'), ((4024, 4042), 'sqlmodel.select', 'select', (['self.model'], {}), '(self.model)\n', (4030, 4042), False, 'from sqlmodel import SQLModel, select, func\n'), ((1793, 1805), 'sqlmodel.func.count', 'func.count', ([], {}), '()\n', (1803, 1805), False, 'from sqlmodel import SQLModel, select, func\n'), ((1819, 1837), 'sqlmodel.select', 'select', (['self.model'], {}), '(self.model)\n', (1825, 1837), False, 'from sqlmodel import SQLModel, select, func\n'), ((1295, 1313), 'sqlmodel.select', 'select', (['self.model'], {}), '(self.model)\n', (1301, 1313), False, 'from sqlmodel import SQLModel, select, func\n'), ((2062, 2080), 'sqlmodel.select', 'select', (['self.model'], {}), '(self.model)\n', (2068, 2080), False, 'from sqlmodel import SQLModel, select, func\n')]