from __future__ import annotations

from typing import Any, Generic, Mapping, NamedTuple, Optional, Self, Sequence, TypeVar

from sqlalchemy import Result, Select, exists, func, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import DeclarativeBase, contains_eager, load_only

from sera.libs.base_orm import BaseORM
from sera.libs.search_helper import Query, QueryOp
from sera.misc import assert_not_null, to_snake_case
from sera.models import Cardinality, Class, DataProperty, IndexType, ObjectProperty

R = TypeVar("R", bound=BaseORM)
ID = TypeVar("ID")  # ID of a class
SqlResult = TypeVar("SqlResult", bound=Result)


class QueryResult(NamedTuple, Generic[R]):
    records: Sequence[R]
    extra_columns: Mapping[str, Sequence]
    total: Optional[int]


class BaseAsyncSearchService(Generic[ID, R]):
    instance = None

    def __init__(self, cls: Class, orm_classes: Mapping[str, type[DeclarativeBase]]):
        # schema of the class
        self.cls = cls
        self.orm_cls = orm_classes[cls.name]
        self.id_prop = assert_not_null(cls.get_id_property())

        self._cls_id_prop = getattr(self.orm_cls, self.id_prop.name)
        self.is_id_auto_increment = assert_not_null(self.id_prop.db).is_auto_increment

        # mapping from property name to ORM class for object properties
        self.prop2orm: dict[str, type] = {
            prop.name: orm_classes[prop.target.name]
            for prop in cls.properties.values()
            if isinstance(prop, ObjectProperty) and prop.target.db is not None
        }

        # figure out the join clauses so we can join the tables
        # for example, to join between User, UserGroup, and Group
        # the query can look like this:
        # select(User)
        #     .join(UserGroup, UserGroup.user_id == User.id)
        #     .join(Group, Group.id == UserGroup.group_id)
        #     .options(contains_eager(User.group).contains_eager(UserGroup.group))
        self.join_clauses: dict[str, list[dict]] = {}
        for prop in cls.properties.values():
            if (
                isinstance(prop, DataProperty)
                and prop.db is not None
                and prop.db.foreign_key is not None
            ):
                target_tbl = orm_classes[prop.db.foreign_key.name]
                target_cls = prop.db.foreign_key
                source_fk = prop.name
                # the property storing the SQLAlchemy relationship of the foreign key
                source_relprop = prop.name + "_relobj"
                cardinality = Cardinality.ONE_TO_ONE
            elif isinstance(prop, ObjectProperty) and prop.target.db is not None:
                target_tbl = orm_classes[prop.target.name]
                target_cls = prop.target
                source_fk = prop.name + "_id"
                source_relprop = prop.name
                cardinality = prop.cardinality
            else:
                continue

            if cardinality == Cardinality.MANY_TO_MANY:
                # for many-to-many, we need to import the association tables
                assoc_tbl = orm_classes[f"{cls.name}{target_cls.name}"]
                assoc_tbl_source_fk = to_snake_case(cls.name) + "_id"
                assoc_tbl_target_fk = to_snake_case(target_cls.name) + "_id"
                self.join_clauses[prop.name] = [
                    {
                        "class": assoc_tbl,
                        "condition": getattr(assoc_tbl, assoc_tbl_source_fk)
                        == getattr(self.orm_cls, self.id_prop.name),
                        "contains_eager": getattr(self.orm_cls, source_relprop),
                    },
                    {
                        "class": target_tbl,
                        "condition": getattr(assoc_tbl, assoc_tbl_target_fk)
                        == getattr(
                            target_tbl,
                            assert_not_null(target_cls.get_id_property()).name,
                        ),
                        "contains_eager": getattr(
                            assoc_tbl, to_snake_case(target_cls.name)
                        ),
                    },
                ]
            elif cardinality == Cardinality.ONE_TO_MANY:
                # A -> B is 1:N, A.id is stored in B, this does not supported in SERA yet so we do not need
                # to implement it
                raise NotImplementedError()
            else:
                # A -> B is either 1:1 or N:1, we will store the foreign key is in A
                # .join(B, A.<foreign_key> == B.id)
                self.join_clauses[prop.name] = [
                    #     {
                    #         "class": target_tbl,
                    #         "condition": getattr(
                    #             target_tbl,
                    #             assert_not_null(target_cls.get_id_property()).name,
                    #         )
                    #         == getattr(self.orm_cls, source_fk),
                    #         "contains_eager": getattr(self.orm_cls, source_relprop),
                    #     },
                ]

    @classmethod
    def get_instance(cls) -> Self:
        """Get the singleton instance of the service."""
        if cls.instance is None:
            # assume that the subclass overrides the __init__ method
            # so that we don't need to pass the class and orm_cls
            cls.instance = cls()  # type: ignore[call-arg]
        return cls.instance

    async def search(
        self,
        query: Query,
        session: AsyncSession,
    ) -> QueryResult[R]:
        """Retrieving records matched a query.

        Args:
            query: The search query
            session: The database session
        """
        q = self._select()
        extra_cols = []

        if len(query.fields) > 0:
            q = q.options(
                load_only(*[getattr(self.orm_cls, field) for field in query.fields])
            )

        if query.unique:
            q = q.distinct()

        if len(query.sorted_by) > 0:
            q = q.order_by(
                *[
                    (
                        (
                            getattr(self.orm_cls, field.field).desc()
                            if field.order == "desc"
                            else getattr(self.orm_cls, field.field)
                        )
                        if field.prop is None
                        else (
                            getattr(self.prop2orm[field.prop], field.field).desc()
                            if field.order == "desc"
                            else getattr(self.prop2orm[field.prop], field.field)
                        )
                    )
                    for field in query.sorted_by
                ]
            )

        if len(query.group_by) > 0:
            q = q.group_by(
                *[
                    (
                        getattr(self.orm_cls, field.field)
                        if field.prop is None
                        else getattr(self.prop2orm[field.prop], field.field)
                    )
                    for field in query.group_by
                ]
            )

        for clause in query.conditions:
            if clause.op == QueryOp.eq:
                q = q.where(getattr(self.orm_cls, clause.field) == clause.value)
            elif clause.op == QueryOp.ne:
                q = q.where(getattr(self.orm_cls, clause.field) != clause.value)
            elif clause.op == QueryOp.lt:
                q = q.where(getattr(self.orm_cls, clause.field) < clause.value)
            elif clause.op == QueryOp.lte:
                q = q.where(getattr(self.orm_cls, clause.field) <= clause.value)
            elif clause.op == QueryOp.gt:
                q = q.where(getattr(self.orm_cls, clause.field) > clause.value)
            elif clause.op == QueryOp.gte:
                q = q.where(getattr(self.orm_cls, clause.field) >= clause.value)
            elif clause.op == QueryOp.in_:
                q = q.where(getattr(self.orm_cls, clause.field).in_(clause.value))
            elif clause.op == QueryOp.not_in:
                q = q.where(~getattr(self.orm_cls, clause.field).in_(clause.value))
            else:
                assert clause.op == QueryOp.fuzzy
                clause_prop = self.cls.properties[clause.field]
                assert (
                    isinstance(clause_prop, DataProperty) and clause_prop.db is not None
                )
                clause_orm_field = getattr(self.orm_cls, clause.field)
                extra_cols.append(f"{clause.field}_score")

                if clause_prop.db.index_type == IndexType.POSTGRES_FTS_SEVI:
                    # fuzzy search is implemented using Postgres Full-Text Search
                    # sevi is a custom text search configuration that we defined in `configs/postgres-fts.sql`
                    q = q.where(
                        func.to_tsvector("sevi", clause_orm_field).bool_op("@@")(
                            func.plainto_tsquery("sevi", clause.value)
                        )
                    )
                    # TODO: figure out which rank function is better
                    # https://www.postgresql.org/docs/current/textsearch-controls.html#TEXTSEARCH-RANKING
                    q = q.order_by(
                        func.ts_rank_cd(
                            func.to_tsvector("sevi", clause_orm_field),
                            func.plainto_tsquery("sevi", clause.value),
                        ).desc()
                    )
                    q = q.add_columns(
                        func.ts_rank_cd(
                            func.to_tsvector("sevi", clause_orm_field),
                            func.plainto_tsquery("sevi", clause.value),
                        ).label(f"{clause.field}_score")
                    )
                elif clause_prop.db.index_type == IndexType.POSTGRES_TRIGRAM:
                    # fuzzy search is implemented using Postgres trigram index
                    # using a custom function f_unaccent to ignore accents -- see `configs/postgres-fts.sql`
                    q = q.where(
                        func.f_unaccent(clause_orm_field).bool_op("%>")(
                            func.f_unaccent(clause.value)
                        )
                    )
                    q = q.order_by(
                        func.f_unaccent(clause_orm_field).op("<->>")(
                            func.f_unaccent(clause.value)
                        )
                    )
                    q = q.add_columns(
                        (
                            1
                            - func.f_unaccent(clause_orm_field).op("<->>")(
                                func.f_unaccent(clause.value)
                            )
                        ).label(f"{clause.field}_score")
                    )
                else:
                    raise NotImplementedError(
                        f"Fuzzy search is not implemented for index type {clause_prop.db.index_type}"
                    )

        for join_condition in query.join_conditions:
            for join_clause in self.join_clauses[join_condition.prop]:
                q = q.join(
                    join_clause["class"],
                    join_clause["condition"],
                    isouter=join_condition.join_type == "left",
                    full=join_condition.join_type == "full",
                ).options(contains_eager(join_clause["contains_eager"]))

        # Create count query without order_by clause to improve performance
        cq = select(func.count()).select_from(q.order_by(None).subquery())
        rq = q.limit(query.limit).offset(query.offset)

        if len(extra_cols) == 0:
            records = self._process_result(await session.execute(rq)).scalars().all()
            extra_columns = {}
        else:
            records = []
            raw_extra_columns = [[] for col in extra_cols]
            for row in self._process_result(await session.execute(rq)):
                records.append(row[0])
                for i in range(len(extra_cols)):
                    raw_extra_columns[i].append(row[i + 1])
            extra_columns = {
                col: raw_extra_columns[i] for i, col in enumerate(extra_cols)
            }

        if query.return_total:
            total = (await session.execute(cq)).scalar_one()
        else:
            total = None
        return QueryResult(records, extra_columns, total)

    async def get_by_id(
        self, session: AsyncSession, id: ID, **controlled_fields: Any
    ) -> Optional[R]:
        """Retrieving a record by ID."""
        q = self._select().where(self._cls_id_prop == id)
        for field, value in controlled_fields.items():
            if value is not None:
                q = q.where(getattr(self.orm_cls, field) == value)
        result = self._process_result(await session.execute(q)).scalar_one_or_none()
        return result

    async def has_id(
        self, session: AsyncSession, id: ID, **controlled_fields: Any
    ) -> bool:
        """Check whether we have a record with the given ID."""
        q = exists().where(self._cls_id_prop == id)
        for field, value in controlled_fields.items():
            if value is not None:
                q = q.where(getattr(self.orm_cls, field) == value)
        q = q.select()
        result = (await session.execute(q)).scalar()
        return bool(result)

    def _select(self) -> Select:
        """Get the select statement for the class."""
        return select(self.orm_cls)

    def _process_result(self, result: SqlResult) -> SqlResult:
        """Process the result of a query."""
        return result
