from typing import Any, Callable, Dict, Generic, List, NewType, Tuple, Type, TypeVar, Union

from pydantic import BaseModel
from tortoise.expressions import Q
from tortoise.models import Model

Total = NewType("Total", int)
ModelType = TypeVar("ModelType", bound=Model)
CreateSchemaType = TypeVar("CreateSchemaType", bound=BaseModel)
UpdateSchemaType = TypeVar("UpdateSchemaType", bound=BaseModel)


class CRUDBase(Generic[ModelType, CreateSchemaType, UpdateSchemaType]):
    def __init__(self, model: Type[ModelType]):
        self.model = model

    async def is_exist(self, **kwargs) -> bool:
        return await self.model.filter(**kwargs).first()
    
    async def all(self, search: Q = Q()) -> List[ModelType]:
        return await self.model.filter(search).all()

    async def get(self, id: int) -> ModelType:
        return await self.model.get(id=id)

    async def list(self, page: int, page_size: int, search: Q = Q(), order: list = []) -> Tuple[Total, List[ModelType]]:
        query = self.model.filter(search)
        return await query.count(), await query.offset((page - 1) * page_size).limit(page_size).order_by(*order)

    async def create(self, obj_in: CreateSchemaType) -> ModelType:
        if isinstance(obj_in, Dict):
            obj_dict = obj_in
        else:
            obj_dict = obj_in.model_dump()
        obj = self.model(**obj_dict)
        await obj.save()
        return obj
    
    async def copy(self, id: int, handler: Callable[[Dict[str, Any]], Dict[str, Any]] = None) -> ModelType:
        obj = await self.get(id=id)
        obj_dict = await obj.to_dict()
        drop_keys = ["id", "created_at", "updated_at"]
        for key in drop_keys:
            obj_dict.pop(key, None)
        if handler:
            obj_dict = handler(obj_dict)
        obj = self.model(**obj_dict)
        await obj.save()
        return obj

    async def update(self, id: int, obj_in: Union[UpdateSchemaType, Dict[str, Any]]) -> ModelType:
        if isinstance(obj_in, Dict):
            obj_dict = obj_in
        else:
            obj_dict = obj_in.model_dump(exclude_unset=True, exclude={"id"})
        obj = await self.get(id=id)
        obj = obj.update_from_dict(obj_dict)
        await obj.save()
        return obj
    
    async def create_or_update(self, obj_in: CreateSchemaType, query_kwargs: Dict[str, Any], update_kwargs: Dict[str, Any] = None) -> ModelType:
        update_kwargs = update_kwargs or {}
        orm = await self.model.filter(**query_kwargs).first()
        if orm: 
            need_update = False if update_kwargs else True
            for key, value in update_kwargs.items():
                if getattr(orm, key) != value:
                    need_update = True
                    break
            if need_update:
                return 'update', await self.update(orm.id, obj_in)
            else: 
                return False, orm
        return 'create', await self.create(obj_in)

    async def remove(self, id: int) -> None:
        obj = await self.get(id=id)
        await obj.delete()