"""Delta comparison using SDV cell-level validation.

Compares baseline deltas (from capture) against actual deltas (from
validate) by sorting both sides and using SDV's ``validate_cell`` for
order-independent set comparison.

Sorting both DataFrames by all columns before the positional cell-level
join in SDV makes the comparison equivalent to set equality — row order
no longer matters, which is correct for insert/delete/update sets.
"""

from __future__ import annotations

import logging
from typing import Any

import polars as pl

from snowflake.snowflake_data_validation.public import (
    ObjectType,
    Origin,
    Platform,
    QueryExecutionResult,
    TableContext,
    validate_cell,
)

from test_runner.common.config import ValidationConfiguration

logger = logging.getLogger(__name__)


def compare_table_deltas(
    baseline_delta: dict[str, Any],
    actual_delta: dict[str, Any],
    table_fqn: str,
    max_cell_diffs: int = ValidationConfiguration().max_cell_diffs,
    ignore_columns: list[str] | None = None,
) -> dict[str, Any]:
    """Compare baseline vs actual delta for a single table.

    Compares inserts-vs-inserts and deletes-vs-deletes as unordered
    sets by sorting both sides before cell-level comparison via SDV.
    Updates are compared by their ``after`` rows.

    When *ignore_columns* is provided, those columns are stripped from
    both baseline and actual row dicts before any comparison.

    Returns a result dict with ``match``, ``table``, ``differences``.
    """
    ignore = {c.upper() for c in ignore_columns} if ignore_columns else set()

    result: dict[str, Any] = {"match": True, "table": table_fqn, "differences": []}

    b_inserts = _strip_columns(baseline_delta.get("inserts", []), ignore)
    a_inserts = _strip_columns(actual_delta.get("inserts", []), ignore)
    b_deletes = _strip_columns(baseline_delta.get("deletes", []), ignore)
    a_deletes = _strip_columns(actual_delta.get("deletes", []), ignore)
    b_updates = baseline_delta.get("updates", [])
    a_updates = actual_delta.get("updates", [])

    if len(b_inserts) != len(a_inserts):
        result["match"] = False
        result["differences"].append({
            "type": "insert_count_mismatch",
            "baseline": len(b_inserts),
            "actual": len(a_inserts),
        })

    if len(b_deletes) != len(a_deletes):
        result["match"] = False
        result["differences"].append({
            "type": "delete_count_mismatch",
            "baseline": len(b_deletes),
            "actual": len(a_deletes),
        })

    if len(b_updates) != len(a_updates):
        result["match"] = False
        result["differences"].append({
            "type": "update_count_mismatch",
            "baseline": len(b_updates),
            "actual": len(a_updates),
        })

    if not result["match"]:
        return result

    for label, b_rows, a_rows in [
        ("insert", b_inserts, a_inserts),
        ("delete", b_deletes, a_deletes),
    ]:
        if not b_rows and not a_rows:
            continue
        diffs = _compare_row_sets(b_rows, a_rows, table_fqn, max_cell_diffs)
        if diffs:
            result["match"] = False
            result["differences"].append({
                "type": f"{label}_content_mismatch",
                "cell_diffs": diffs,
            })

    if b_updates or a_updates:
        b_after = _strip_columns([u.get("after", {}) for u in b_updates], ignore)
        a_after = _strip_columns([u.get("after", {}) for u in a_updates], ignore)
        diffs = _compare_row_sets(b_after, a_after, table_fqn, max_cell_diffs)
        if diffs:
            result["match"] = False
            result["differences"].append({
                "type": "update_content_mismatch",
                "cell_diffs": diffs,
            })

    return result


def _strip_columns(
    rows: list[dict[str, Any]],
    ignore: set[str],
) -> list[dict[str, Any]]:
    """Remove ignored columns (upper-cased) from each row dict."""
    if not ignore or not rows:
        return rows
    return [
        {k: v for k, v in row.items() if k.upper() not in ignore}
        for row in rows
    ]


def _compare_row_sets(
    baseline_rows: list[dict[str, Any]],
    actual_rows: list[dict[str, Any]],
    table_fqn: str,
    max_cell_diffs: int,
) -> list[dict]:
    """Compare two row sets as unordered sets using SDV validate_cell.

    Both sides are sorted by all columns before positional comparison,
    making the comparison order-independent.
    """
    if not baseline_rows and not actual_rows:
        return []

    b_df = pl.DataFrame(baseline_rows)
    a_df = pl.DataFrame(actual_rows)

    b_df = _sortable(b_df).sort(b_df.columns)
    a_df = _sortable(a_df).sort(a_df.columns)

    table_context = _make_table_context(table_fqn, len(baseline_rows))

    validation_result = validate_cell(
        source_query_result=QueryExecutionResult(dataframe=b_df),
        target_query_result=QueryExecutionResult(dataframe=a_df),
        table_context=table_context,
    )

    if validation_result.is_valid:
        return []

    df = validation_result.results_dataframe.rename(columns={
        "ROW": "row",
        "COLUMN": "column",
        "SOURCE_VALUE": "baseline",
        "TARGET_VALUE": "actual",
    })
    return df.head(max_cell_diffs).to_dict(orient="records")


def _sortable(df: pl.DataFrame) -> pl.DataFrame:
    """Cast ``Object`` columns (e.g. Python ``bytes``) to ``Utf8`` so
    Polars can sort the frame without panicking."""
    casts = [
        pl.col(c).cast(pl.Utf8) for c in df.columns
        if df[c].dtype == pl.Object
    ]
    return df.with_columns(casts) if casts else df


def _make_table_context(table_fqn: str, row_count: int) -> TableContext:
    """Build a minimal TableContext for SDV validate_cell."""
    return TableContext(
        apply_metric_column_modifier=True,
        chunk_number=1,
        column_mappings={},
        column_selection_list=[],
        columns=[],
        database_name="",
        exclude_metrics=False,
        fully_qualified_name=table_fqn,
        has_where_clause=False,
        id=1,
        internal_fully_qualified_name=table_fqn,
        internal_table_name=table_fqn,
        is_case_sensitive=False,
        is_exclusion_mode=False,
        is_temporary_table=True,
        max_failed_rows_number=100,
        object_type=ObjectType.TABLE,
        origin=Origin.SOURCE,
        platform=Platform.SNOWFLAKE,
        row_count=row_count,
        run_id="",
        run_start_time="",
        schema_name="",
        sql_generator=None,
        table_name=table_fqn,
        templates_loader_manager=None,
        user_index_column_collection=[],
        where_clause="",
    )
