"""A linking methodology based on a deterministic set of conditions."""

import json

import duckdb
import polars as pl
from pydantic import Field, field_validator

from matchbox.client.models import comparison
from matchbox.client.models.linkers.base import Linker, LinkerSettings
from matchbox.common.logging import logger


class DeterministicSettings(LinkerSettings):
    """A data class to enforce the Deterministic linker's settings dictionary shape."""

    comparisons: list[str] | list[list[str]] = Field(
        description="""
            Comparison rules for matching using DuckDB SQL syntax.
            
            Can be specified as:
            - A flat list of strings: All comparisons applied in parallel (OR logic)
            - A nested list of lists: Sequential rounds of matching
            
            Flat list (parallel):
            [
                "left.company_number = right.company_number",
                "left.name = right.name",
            ]
            All comparisons applied to full datasets, results unioned.
            
            Nested list (sequential rounds):
            [
                [
                    "left.company_number = right.company_number",
                    "left.name = right.name",
                ],
                [
                    "left.name_normalised = right.name_normalised",
                    "left.website = right.website",
                ],
            ]
            Each inner list is a "round". Within each round, comparisons use OR 
            logic. After each round, matched records are removed from the pool 
            before the next round.
            
            Use left.field and right.field to refer to columns in the respective 
            sources. Supports all DuckDB SQL operations and functions.
        """,
    )

    @field_validator("comparisons", mode="before")
    @classmethod
    def validate_comparison(
        cls, value: str | list[str] | list[list[str]]
    ) -> list[list[str]]:
        """Normalise to list of lists format."""
        if isinstance(value, str):
            return [[comparison(value, dialect="duckdb")]]
        if not value:
            raise ValueError("comparisons cannot be empty")
        if all(isinstance(v, str) for v in value):
            return [[comparison(v, dialect="duckdb") for v in value]]
        if all(isinstance(v, list) for v in value):
            for round_idx, round_comparisons in enumerate(value):
                if not round_comparisons:
                    raise ValueError(f"Round {round_idx} cannot be empty")
                if not all(isinstance(c, str) for c in round_comparisons):
                    raise ValueError(f"Round {round_idx} must contain only strings")
            return [[comparison(c, dialect="duckdb") for c in r] for r in value]
        raise ValueError(
            "comparisons must be a string, list of strings, or list of lists"
        )


class DeterministicLinker(Linker):
    """A deterministic linker that links based on a set of boolean conditions.

    Uses DuckDB as the SQL backend, enabling rich SQL operations while maintaining
    a Polars DataFrame interface. Supports both parallel matching (single round)
    and sequential matching (multiple rounds where matched records are removed
    after each round).
    """

    settings: DeterministicSettings

    def prepare(self, left: pl.DataFrame, right: pl.DataFrame) -> None:
        """Prepare the linker for linking."""
        pass

    def link(self, left: pl.DataFrame, right: pl.DataFrame) -> pl.DataFrame:
        """Link the left and right dataframes.

        If comparisons is a flat list, applies all comparisons in parallel.
        If comparisons is a nested list, applies each round sequentially,
        removing matched records from the pool after each round.
        """
        con: duckdb.DuckDBPyConnection = duckdb.connect(":memory:")
        try:
            all_matches: list[pl.DataFrame] = []
            remaining_left, remaining_right = left, right

            for round_num, round_comparisons in enumerate(
                self.settings.comparisons, start=1
            ):
                if remaining_left.is_empty() or remaining_right.is_empty():
                    logger.info(f"Round {round_num}: Skipping - no records remaining")
                    break

                logger.info(
                    f"Round {round_num}: {len(remaining_left):,} left × "
                    f"{len(remaining_right):,} right"
                )

                matches = self._link_round(
                    con, remaining_left, remaining_right, round_comparisons, round_num
                )

                logger.info(f"Round {round_num}: Found {len(matches):,} matches")

                if not matches.is_empty():
                    all_matches.append(matches)
                    matched_left = matches.select("left_id").unique()
                    matched_right = matches.select("right_id").unique()
                    remaining_left = remaining_left.join(
                        matched_left,
                        left_on=self.settings.left_id,
                        right_on="left_id",
                        how="anti",
                    )
                    remaining_right = remaining_right.join(
                        matched_right,
                        left_on=self.settings.right_id,
                        right_on="right_id",
                        how="anti",
                    )

            return self._finalise_results(all_matches)
        finally:
            con.close()

    def _link_round(
        self,
        con: duckdb.DuckDBPyConnection,
        left: pl.DataFrame,
        right: pl.DataFrame,
        comparisons: list[str],
        round_num: int,
    ) -> pl.DataFrame:
        """Apply all comparisons in a round using OR logic via DuckDB."""
        con.register("left_df", left)
        con.register("right_df", right)

        subqueries: list[str] = []
        for condition in comparisons:
            subquery: str = f"""
                SELECT
                    l.{self.settings.left_id} AS left_id,
                    r.{self.settings.right_id} AS right_id,
                    1.0 AS probability
                FROM left_df l
                INNER JOIN right_df r
                    ON {condition}
            """
            subqueries.append(subquery)

        query: str = f"""
            SELECT DISTINCT *
            FROM ({" UNION ALL ".join(subqueries)})
        """

        max_est: int = self._get_max_cardinality(con, query)
        logger.info(f"Round {round_num}: Estimated max cardinality: {max_est:,}")

        return con.execute(query).pl()

    def _get_max_cardinality(self, con: duckdb.DuckDBPyConnection, query: str) -> int:
        """Get max cardinality estimate from DuckDB plan, or -1 if unavailable."""
        explain = con.execute(
            f"PRAGMA explain_output = 'all'; EXPLAIN (FORMAT json) {query}"
        ).fetchall()
        plans = {k: json.loads(v) for k, v in dict(explain).items()}

        estimates: list[dict] = []
        for plan_name, plan_type in [
            ("physical_plan", "physical"),
            ("logical_opt", "optimised"),
        ]:
            if plan := plans.get(plan_name):
                estimates.extend(self._traverse_plan(plan[0], plan_type))

        # Debug: log full tree breakdown
        logger.debug("Plan breakdown:")
        for plan_type in ["physical", "optimised"]:
            plan_ests = [e for e in estimates if e["plan_type"] == plan_type]
            if plan_ests:
                logger.debug(f"  {plan_type.capitalize()} plan:")
                for est in plan_ests:
                    indent = "    " + "  " * est["depth"]
                    logger.debug(f"{indent}{est['node_name']}: {est['cardinality']:,}")

        # Return max for info logging
        return max(
            (e["cardinality"] for e in estimates if e["cardinality"] > 0),
            default=-1,
        )

    def _traverse_plan(self, node: dict, plan_type: str, depth: int = 0) -> list[dict]:
        """Recursively collect cardinality estimates with metadata from plan tree."""
        estimates: list[dict] = []
        cardinality = node.get("extra_info", {}).get("Estimated Cardinality")
        if cardinality is not None:
            estimates.append(
                {
                    "cardinality": int(cardinality),
                    "node_name": node.get("name", "UNKNOWN"),
                    "depth": depth,
                    "plan_type": plan_type,
                }
            )
        for child in node.get("children", []):
            estimates.extend(self._traverse_plan(child, plan_type, depth + 1))
        return estimates

    def _finalise_results(self, all_matches: list[pl.DataFrame]) -> pl.DataFrame:
        """Combine matches from all rounds and ensure correct schema."""
        if all_matches:
            return pl.concat(all_matches).with_columns(
                pl.col("probability").cast(pl.Float32)
            )
        return pl.DataFrame(
            {"left_id": [], "right_id": [], "probability": []}
        ).with_columns(pl.col("probability").cast(pl.Float32))
