"""Evaluation PostgreSQL mixin for Matchbox server."""

from datetime import UTC, datetime, timedelta
from itertools import chain
from typing import TYPE_CHECKING, Any

import numpy as np
import polars as pl
import pyarrow as pa
from pyarrow import Table
from sqlalchemy import BIGINT, func, select
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.orm import Session

from matchbox.common.arrow import (
    SCHEMA_CLUSTER_EXPANSION,
    SCHEMA_EVAL_SAMPLES,
    SCHEMA_JUDGEMENTS,
)
from matchbox.common.datatypes import require
from matchbox.common.db import QueryReturnType, sql_to_df
from matchbox.common.dtos import ResolverStepPath
from matchbox.common.eval import Judgement as CommonJudgement
from matchbox.common.exceptions import (
    MatchboxTooManySamplesRequested,
    MatchboxUserNotFoundError,
)
from matchbox.common.transform import hash_cluster_leaves
from matchbox.server.postgresql.db import MBDB
from matchbox.server.postgresql.orm import (
    Clusters,
    ClusterSourceKey,
    Contains,
    EvalJudgements,
    ResolverClusters,
    SourceConfigs,
    Steps,
    Users,
)
from matchbox.server.postgresql.utils.db import (
    compile_sql,
    ingest_to_temporary_table,
)

if TYPE_CHECKING:
    from pyarrow import Table as ArrowTable

    from matchbox.server.postgresql.adapter.admin import MatchboxPostgresAdminMixin

    # The evaluation mixin has self functions that reference the admin mixin
    # This lets us refernce it in a typesafe way
    _MixinBase = MatchboxPostgresAdminMixin
else:
    ArrowTable = Any
    _MixinBase = object


class MatchboxPostgresEvaluationMixin(_MixinBase):
    """Evaluation mixin for the PostgreSQL adapter for Matchbox."""

    def insert_judgement(self, user_name: str, judgement: CommonJudgement) -> None:  # noqa: D102
        def get_or_create_cluster(leaves: list[int], session: Session) -> int:
            # Compute hash corresponding to set of source clusters (leaves)
            leaf_hashes = [
                session.scalars(
                    select(Clusters.cluster_hash).where(Clusters.cluster_id == leaf_id)
                ).one()
                for leaf_id in leaves
            ]
            cluster_hash = hash_cluster_leaves(leaf_hashes)

            # Upsert cluster
            result = session.execute(
                insert(Clusters)
                .values(cluster_hash=cluster_hash)
                .on_conflict_do_nothing(index_elements=["cluster_hash"])
                .returning(Clusters.cluster_id)
            )

            # Get cluster_id
            newly_created = (cluster_id := result.scalar_one_or_none()) is not None

            if not newly_created:
                cluster_id = session.scalars(
                    select(Clusters.cluster_id).where(
                        Clusters.cluster_hash == cluster_hash
                    )
                ).one()

            # Only insert Contains relationships if we created a new cluster
            if newly_created:
                for leaf_id in leaves:
                    session.execute(
                        insert(Contains)
                        .values(root=cluster_id, leaf=leaf_id)
                        .on_conflict_do_nothing()
                    )

            return require(cluster_id)

        # Check that all referenced leaf IDs exist
        ids = list(chain(*judgement.endorsed)) + judgement.shown
        self.validate_ids(ids)

        with MBDB.get_session() as session:
            # Check that the user exists
            user = session.scalar(select(Users).where(Users.name == user_name))
            if not user:
                raise MatchboxUserNotFoundError(f"User '{user_name}' not found")

            # Get or create shown cluster
            shown_id = get_or_create_cluster(judgement.shown, session)

            # Insert all judgements
            for leaves in judgement.endorsed:
                endorsed_id = get_or_create_cluster(leaves, session)
                session.execute(
                    insert(EvalJudgements).values(
                        user_id=user.user_id,
                        tag=judgement.tag,
                        shown_cluster_id=shown_id,
                        endorsed_cluster_id=endorsed_id,
                        timestamp=datetime.now(UTC),
                    )
                )

            session.commit()

    def get_judgements(self, tag: str | None = None) -> tuple[Table, Table]:  # noqa: D102
        def _cast_tables(
            judgements: pl.DataFrame, cluster_expansion: pl.DataFrame
        ) -> tuple[pa.Table, pa.Table]:
            """Cast judgement tables to conform to data transfer schema."""
            judgements = judgements.cast(pl.Schema(SCHEMA_JUDGEMENTS))
            cluster_expansion = cluster_expansion.cast(
                pl.Schema(SCHEMA_CLUSTER_EXPANSION)
            )

            return (
                judgements.to_arrow().cast(SCHEMA_JUDGEMENTS),
                cluster_expansion.to_arrow().cast(SCHEMA_CLUSTER_EXPANSION),
            )

        judgements_stmt = select(
            Users.name.label("user_name"),
            EvalJudgements.endorsed_cluster_id.label("endorsed"),
            EvalJudgements.shown_cluster_id.label("shown"),
        ).join(Users)

        if tag:
            judgements_stmt = judgements_stmt.where(EvalJudgements.tag == tag)

        with MBDB.get_adbc_connection() as conn:
            judgements = sql_to_df(
                stmt=compile_sql(judgements_stmt),
                connection=conn,
                return_type=QueryReturnType.POLARS,
            )

        if not len(judgements):
            cluster_expansion = pl.DataFrame(schema=pl.Schema(SCHEMA_CLUSTER_EXPANSION))
            return _cast_tables(judgements, cluster_expansion)

        shown_clusters = set(judgements["shown"].to_list())
        endorsed_clusters = set(judgements["endorsed"].to_list())
        referenced_clusters = Table.from_pydict(
            {"root": list(shown_clusters | endorsed_clusters)}
        )

        with ingest_to_temporary_table(
            table_name="judgements",
            schema_name="mb",
            column_types={
                "root": BIGINT(),
            },
            data=referenced_clusters,
        ) as temp_table:
            cluster_expansion_stmt = (
                select(temp_table.c.root, func.array_agg(Contains.leaf).label("leaves"))
                .select_from(temp_table)
                .join(Contains, Contains.root == temp_table.c.root)
                .group_by(temp_table.c.root)
            )

            with MBDB.get_adbc_connection() as conn:
                cluster_expansion = sql_to_df(
                    stmt=compile_sql(cluster_expansion_stmt),
                    connection=conn,
                    return_type=QueryReturnType.POLARS,
                )

        return _cast_tables(judgements, cluster_expansion)

    def sample_for_eval(  # noqa: D102
        self, n: int, path: ResolverStepPath, user_name: str
    ) -> ArrowTable:
        """Sample some clusters from a resolver step."""
        # Not currently checking validity of the user_name
        # If the user ID does not exist, the exclusion by previous judgements breaks
        if n > 100:
            # This reasonable assumption means simple "IS IN" function later is fine
            raise MatchboxTooManySamplesRequested(
                "Can only sample 100 entries at a time."
            )

        with MBDB.get_session() as session:
            # Get user
            user = session.scalar(select(Users).where(Users.name == user_name))
            if not user:
                raise MatchboxUserNotFoundError(f"User '{user_name}' not found")

            # Use ORM to get step metadata
            resolver_step = Steps.from_path(path=path, session=session)
            step_id = resolver_step.step_id

        # Get cluster IDs and features for this resolver and user.
        user_judgements = (
            select(EvalJudgements)
            .where(EvalJudgements.user_id == user.user_id)
            .subquery()
        )
        cluster_features_stmt = (
            select(
                ResolverClusters.cluster_id,
                # Practically our client explicitly avoids showing the same cluster to
                # a user more than oce, so conflicting decisions over time shouldn't be
                # possible. However, the ORM doesn't enforce this, so we take only the
                # latest decision here
                func.max(user_judgements.c.timestamp).label("latest_ts"),
            )
            .join(
                user_judgements,
                ResolverClusters.cluster_id == user_judgements.c.shown_cluster_id,
                isouter=True,
            )
            .where(
                ResolverClusters.step_id == step_id,
            )
            .group_by(ResolverClusters.cluster_id)
        )

        with MBDB.get_adbc_connection() as conn:
            cluster_features = sql_to_df(
                stmt=compile_sql(cluster_features_stmt),
                connection=conn,
                return_type=QueryReturnType.POLARS,
            )

        # Exclude clusters recently judged by this user
        to_sample = cluster_features.filter(
            (pl.col("latest_ts") < datetime.now(UTC) - timedelta(days=365))
            | (pl.col("latest_ts").is_null())
        )

        # Return early if nothing to sample from
        if not len(to_sample):
            return pl.DataFrame(schema=pl.Schema(SCHEMA_EVAL_SAMPLES)).to_arrow()

        # With fewer clusters than requested, return all
        if to_sample.shape[0] <= n:
            sampled_cluster_ids = to_sample.select("cluster_id").to_series().to_list()
        else:
            indices = np.random.choice(to_sample.shape[0], size=n, replace=False)
            sampled_cluster_ids = (
                to_sample[indices].select("cluster_id").to_series().to_list()
            )

        # Get all info we need for the cluster IDs we've sampled, i.e.:
        # source cluster IDs, keys and source steps
        with MBDB.get_adbc_connection() as conn:
            source_clusters = (
                select(Contains.root, Contains.leaf)
                .where(Contains.root.in_(sampled_cluster_ids))
                .subquery()
            )

            # The same leaf can be reused to represent rows across different sources.
            # Only retrieve info for sources upstream of this resolver step.
            source_step_ids = [
                res.step_id for res in resolver_step.ancestors if res.type == "source"
            ]
            source_steps = (
                select(Steps.name, Steps.step_id)
                .where(Steps.step_id.in_(source_step_ids))
                .subquery()
            )

            enrich_stmt = (
                select(
                    source_clusters.c.root,
                    source_clusters.c.leaf,
                    ClusterSourceKey.key,
                    source_steps.c.name.label("source"),
                )
                .select_from(source_clusters)
                .join(
                    ClusterSourceKey,
                    ClusterSourceKey.cluster_id == source_clusters.c.leaf,
                )
                .join(
                    SourceConfigs,
                    SourceConfigs.source_config_id == ClusterSourceKey.source_config_id,
                )
                .join(
                    source_steps,
                    source_steps.c.step_id == SourceConfigs.step_id,
                )
            )

            final_samples = sql_to_df(
                stmt=compile_sql(enrich_stmt),
                connection=conn,
                return_type=QueryReturnType.POLARS,
            )
            return final_samples.cast(pl.Schema(SCHEMA_EVAL_SAMPLES)).to_arrow()
