from typing import Optional, Union, List

from pydantic import BaseModel
from sqlalchemy import desc, asc
from sqlalchemy.exc import MultipleResultsFound, NoResultFound
from sqlalchemy.orm import InstrumentedAttribute
from sqlmodel import SQLModel

from common.database.exceptions import DatabaseRowUnavailable
from common.database.session_manager import GlobalSessionManager as SessionManager
from common.database.utils import find_primary_key


class BaseSQLModel(SQLModel):
    @classmethod
    async def create(
        cls, model_create: Union["ModelCreate", BaseModel, None] = None, **kwargs
    ) -> "BaseSQLModel":
        if model_create is not None:
            create_dict = (
                model_create.create_dict()
                if isinstance(model_create, ModelCreate)
                else model_create.model_dump()
            )
            create_kwargs = {**create_dict, **kwargs}
        else:
            create_kwargs = dict(kwargs)
        # Validate the create_kwargs
        validated_kwargs = cls.model_validate(create_kwargs).model_dump()
        # Create the model
        model = await SessionManager.async_create(cls, **validated_kwargs)
        return model

    @classmethod
    async def get(
        cls,
        id: str = None,
        raise_error: bool = False,
        mode: str = "one",
        offset: Optional[int] = None,
        limit: Optional[int] = None,
        select_columns: Optional[List[str | InstrumentedAttribute]] = None,
        **kwargs,
    ) -> Union["BaseSQLModel", list["BaseSQLModel"], dict, list[dict], None]:
        try:
            kwargs = kwargs or {}

            if id is not None:
                pk_id = find_primary_key(cls)
                assert (
                    pk_id not in kwargs
                ), f"Primary key {pk_id} is already in kwargs, but id=... arg is specified."
                kwargs[pk_id] = id

            # Add offset and limit directly to kwargs
            if offset is not None:
                kwargs["offset"] = offset
            if limit is not None:
                kwargs["limit"] = limit

            if select_columns is not None:
                assert (
                    len(select_columns) >= 1
                ), "select_columns must have at least one column."
                kwargs["select_columns"] = select_columns
                kwargs["get_scalars"] = False
                kwargs["as_dict"] = True

            return await SessionManager.async_get(cls, mode=mode, **kwargs)
        except (MultipleResultsFound, NoResultFound) as e:
            error_messages = {
                MultipleResultsFound: "Multiple records found for the given criteria.",
                NoResultFound: "No record found for the given criteria.",
            }
            if raise_error:
                raise DatabaseRowUnavailable(error_messages[type(e)])
            return None

    @classmethod
    async def get_or_create(cls, **kwargs) -> "BaseSQLModel":
        model = await cls.get(raise_error=False, **kwargs)
        if model is None:
            model = await cls.create(model_create=cls.model_validate(kwargs))  # noqa
        return model

    async def update(
        self, model_update: Union["ModelUpdate", BaseModel, dict, None] = None, **kwargs
    ) -> "BaseSQLModel":
        if model_update is not None:
            if isinstance(model_update, ModelUpdate):
                update_dict = model_update.update_dict()
            elif isinstance(model_update, BaseModel):
                update_dict = model_update.model_dump(exclude_unset=True)
            elif isinstance(model_update, dict):
                update_dict = model_update
            else:
                raise TypeError(
                    "model_update must be of type ModelUpdate, BaseModel, or dict"
                )
            update_kwargs = {**update_dict, **kwargs}
        else:
            update_kwargs = dict(kwargs)

        model_validate = self.model_validate(self.model_dump(), update=update_kwargs)
        update_kwargs = model_validate.model_dump(exclude_unset=True)

        self.sqlmodel_update(update_kwargs)

        return await self.commit(do_refresh=True)

    @classmethod
    async def batch_update(
        cls,
        get_kwargs: Optional[dict] = None,
        join_models: Optional[list] = None,
        custom_predicates: Optional[list] = None,
        do_commit: bool = True,
        **update_values,
    ) -> int:
        return await SessionManager.async_batch_update(
            cls,
            get_kwargs=get_kwargs,
            join_models=join_models,
            custom_predicates=custom_predicates,
            do_commit=do_commit,
            **update_values,
        )

    async def delete(self) -> None:
        await SessionManager.async_delete(self)

    async def refresh(self) -> "BaseSQLModel":
        return await SessionManager.async_refresh(self)

    async def commit(self, do_refresh: bool = True) -> "BaseSQLModel":
        await SessionManager.async_add_and_commit(self)
        if do_refresh:
            await self.refresh()
        return self

    @classmethod
    async def count(cls, **kwargs) -> int:
        return await SessionManager.async_count(cls, **kwargs)

    @classmethod
    async def sort(
        cls,
        offset: Optional[int] = None,
        select_columns: Optional[List[str | InstrumentedAttribute]] = None,
        **kwargs,
    ) -> Union[List["BaseSQLModel"], List[dict]]:
        if offset is not None:
            kwargs["offset"] = offset
        if select_columns is not None:
            kwargs["select_columns"] = select_columns
            kwargs["get_scalars"] = False
            kwargs["as_dict"] = True
        return await SessionManager.async_sort(cls, **kwargs)

    @classmethod
    async def simple_sort(
        cls,
        *,
        field: str,
        descending: bool = True,
        offset: Optional[int] = None,
        limit: Optional[int] = None,
        select_columns: Optional[List[str | InstrumentedAttribute]] = None,
        **kwargs,
    ) -> Union[List["BaseSQLModel"], List[dict]]:
        """Simple sorting interface for common use cases.

        Args:
            field: Field name to sort by
            descending: Sort in descending order if True, ascending if False
            offset: Number of items to skip
            limit: Maximum number of items to return
            select_columns: List of column names to select (returns Row objects if specified)
            **kwargs: Additional filter parameters

        Example:
            # Sort by creation date, newest first
            items = await Model.simple_sort(
                field="created_at",
                descending=True,
                limit=10,
                status="active"  # additional filter
            )

            # Get only specific columns
            rows = await Model.simple_sort(
                field="created_at",
                descending=True,
                select_columns=["id", "name", "created_at"],
                limit=10
            )
        """
        # If field is already in kwargs as a filter predicate, we can't pass
        # it again as a sort kwarg (would overwrite the filter).
        # Use custom_sort_orders instead.
        if field in kwargs:
            sort_fn = desc if descending else asc
            custom_sort_orders = [
                (sort_fn(SessionManager.get_model_attr(cls, field)), limit)
            ]
            return await cls.sort(
                offset=offset,
                select_columns=select_columns,
                custom_sort_orders=custom_sort_orders,
                **kwargs,
            )

        # Normal case: field is not used as a filter, safe to use as sort kwarg
        sort_order = (
            SessionManager.descending() if descending else SessionManager.ascending()
        )
        if limit is not None:
            sort_order.limit = limit

        return await cls.sort(
            offset=offset,
            select_columns=select_columns,
            **kwargs,
            **{field: sort_order},
        )

    @classmethod
    async def exists(cls, **kwargs) -> bool:
        count = await cls.count(**kwargs)
        return count > 0


class ModelRead(BaseSQLModel):
    pass


class ModelCreate(BaseSQLModel):
    def create_dict(self) -> dict:
        return self.model_dump()


class ModelUpdate(BaseSQLModel):
    def update_dict(self):
        return self.model_dump(exclude_unset=True)
