from pydantic import BaseModel
from sqlalchemy import func

from core.models.base import ORDER_BY_CREATED_AT
from core.utils.table_engine.filter import TableEngineFilter
from core.utils.table_engine.input import TableEngineInput


class TableEngine:
    object: any
    input: TableEngineInput
    is_sql_server: bool
    is_deleted: bool
    model: any
    joins: list = None
    filter: TableEngineFilter
    sort: dict = ORDER_BY_CREATED_AT
    extra_filter: list

    def __init__(self, _object, _input: dict):
        self.object = _object
        _input = _input or {}
        self.input = TableEngineInput(**_input)
        self.input.get_limit()
        self.filter = TableEngineFilter(self.input)
        self.extra_filter = []

    def soft_deleted(self):
        self.is_deleted = True
        self.input.is_deleted = True
        return self

    def sql_server(self):
        self.is_sql_server = True
        return self

    def text_search_cols(self, cols: list):
        self.input.text_search_cols(cols)
        return self

    def set_limit(self, limit: int):
        self.input.p.s = limit
        return self

    def set_model(self, model: BaseModel.__class__):
        self.model = model
        return self

    def join(self, joins):
        self.joins = joins
        return self

    def sort_by(self, sort_by: dict | str):
        self.sort = sort_by
        return self

    def _get_sql_order(self):
        if isinstance(self.sort, dict):
            return self.model.Id.desc()
        return self.sort

    def add_filter(self, filters):
        self.extra_filter.append(filters)
        self.input.extra_filter.append(filters)
        return self

    def get_query(self):
        return self.input.dict()

    async def exec(self):
        if hasattr(self, 'is_sql_server') and self.is_sql_server:
            return await self._exec_sql_server()
        else:
            return await self._exec_mongo()

    async def _exec_mongo(self):
        cursor = self.object.find(
            self.filter.get()).sort(
            **self.sort).limit(
            self.input.get_limit()).skip(self.input.skip())
        return [self.model(**document) async for document in cursor]

    async def _exec_sql_server(self):
        query = self.object \
            .query(self.model)

        if self.joins:
            for j in self.joins:
                if isinstance(j, list):
                    query = query.join(*j)
                else:
                    query = query.join(j)

        query = self.filter.get_sql_server(
            self.model, query, self.extra_filter)

        count = query.count()

        query = query \
            .order_by(self._get_sql_order()) \
            .offset(self.input.skip()) \
            .limit(self.input.get_limit())

        return query.all(), count

    async def count(self):
        if hasattr(self, 'is_sql_server') and self.is_sql_server:
            count = self.object.query(func.count(self.model.Id))
        else:
            count = await self.object.count_documents(
                self.filter.get())
        return count