"""In-memory validation using DuckDB.

Compares baseline and actual result sets in three levels:
row counts → per-column summary stats → cell-level diffs.
Results are plain dicts suitable for JSON serialization.

TODO: migrate read-only sproc comparison to use SDV directly
(row/cell validation APIs) instead of the DuckDB row-count and
summary-stats layers here.  Delta comparison for read-write
sprocs already uses SDV validate_cell directly via
``replay.delta_comparator``.
"""

from __future__ import annotations

import logging
import time
from typing import Any

import duckdb
import pandas as pd
import polars as pl
from snowflake.snowflake_data_validation.connector.connector_base import ConnectorBase
from snowflake.snowflake_data_validation.public import (
    ObjectType,
    Origin,
    Platform,
    QueryExecutionResult,
    TableContext,
    ValidationResult,
    validate_cell,
)

from test_runner.common.config import TestRunnerConfig, ValidationConfiguration
from test_runner.common.database_dialect import DatabaseDialect
from test_runner.common.models import (
    MAX_ERROR_LENGTH,
    ColumnTypeMap,
    ComparisonResult,
    RowDict,
    ValidationDifference,
    truncate,
)
from test_runner.validate.transient_table_manager import (
    execute_with_capture,
    fetch_baseline_output_rows,
    fetch_baseline_rows,
    is_numeric_type,
)

LOGGER = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _load_duckdb(
    baseline_rows: list[RowDict],
    actual_rows: list[RowDict],
    col_types: ColumnTypeMap,
) -> tuple[duckdb.DuckDBPyConnection, list[str]]:
    """Register baseline and actual rows as DuckDB tables; return ``(conn, columns)``."""
    columns = list(col_types.keys()) if col_types else (
        list(baseline_rows[0].keys()) if baseline_rows else
        list(actual_rows[0].keys()) if actual_rows else []
    )

    conn = duckdb.connect()

    def _register(name: str, rows: list[RowDict]) -> None:
        if rows:
            df = pd.DataFrame(rows, columns=columns)
        else:
            df = pd.DataFrame(columns=columns)
        # Assign sequential 1-based row numbers for join-based cell diff
        df.insert(0, "_row", range(1, len(df) + 1))
        conn.register(name, df)

    _register("baseline", baseline_rows)
    _register("actual", actual_rows)
    return conn, columns


def _row_count_diff(conn: duckdb.DuckDBPyConnection, columns: list[str]) -> dict[str, Any]:
    """Return a row-counts section with baseline/actual counts."""
    b_count = conn.execute("SELECT COUNT(*) FROM baseline").fetchone()[0]
    a_count = conn.execute("SELECT COUNT(*) FROM actual").fetchone()[0]
    return {
        "baseline": b_count,
        "actual": a_count,
        "match": b_count == a_count,
    }


def _summary_stats_diff(
    conn: duckdb.DuckDBPyConnection,
    columns: list[str],
    col_types: ColumnTypeMap | None = None,
    *,
    tolerance: float,
) -> dict[str, Any]:
    """Return per-column summary stat mismatches for numeric columns.

    Non-numeric columns are skipped (handled by cell-level diff).
    Each mismatch: ``{"column": ..., "metrics": {"sum": {"baseline": x, "actual": y}, ...}}``.
    """
    mismatches: list[dict] = []

    # col_types maps column name -> Snowflake SQL type (e.g. "NUMBER", "VARCHAR").
    # is_numeric_type checks the type string against known numeric type prefixes
    # (_NUMERIC_TYPE_PREFIXES in transient_table_manager) so only numeric columns
    # get aggregate stats; text/date columns are handled by cell-level diff instead.
    for col in columns:
        if col_types and not is_numeric_type(col_types.get(col, "")):
            continue  # skip VARCHAR, DATE, BOOLEAN, etc.
        quoted = f'"{col}"'
        try:
            b_stats = conn.execute(
                f"SELECT COUNT(*) - COUNT({quoted}) AS null_count,"
                f" MIN(TRY_CAST({quoted} AS DOUBLE)) AS min_val,"
                f" MAX(TRY_CAST({quoted} AS DOUBLE)) AS max_val,"
                f" AVG(TRY_CAST({quoted} AS DOUBLE)) AS avg_val,"
                f" SUM(TRY_CAST({quoted} AS DOUBLE)) AS sum_val"
                f" FROM baseline"
            ).fetchone()
            a_stats = conn.execute(
                f"SELECT COUNT(*) - COUNT({quoted}) AS null_count,"
                f" MIN(TRY_CAST({quoted} AS DOUBLE)) AS min_val,"
                f" MAX(TRY_CAST({quoted} AS DOUBLE)) AS max_val,"
                f" AVG(TRY_CAST({quoted} AS DOUBLE)) AS avg_val,"
                f" SUM(TRY_CAST({quoted} AS DOUBLE)) AS sum_val"
                f" FROM actual"
            ).fetchone()
        except Exception as exc:  # noqa: BLE001
            LOGGER.debug("Stats computation skipped for column %r: %s", col, exc)
            continue

        b_null, b_min, b_max, b_avg, b_sum = b_stats
        a_null, a_min, a_max, a_avg, a_sum = a_stats

        # All numeric metrics use tolerance-aware float comparison.
        # Both-None means TRY_CAST failed on both sides (text column) → equal.
        candidates = [
            ("null_count", b_null, a_null),
            ("min", b_min, a_min),
            ("max", b_max, a_max),
            ("avg", b_avg, a_avg),
            ("sum", b_sum, a_sum),
        ]
        metric_diffs = {}
        for metric, b_val, a_val in candidates:
            if metric == "null_count":
                if b_val != a_val:
                    metric_diffs[metric] = {"baseline": b_val, "actual": a_val}
            else:
                if b_val is None and a_val is None:
                    continue  # both non-numeric (text col) — skip
                if b_val is None or a_val is None or abs(b_val - a_val) > tolerance:
                    metric_diffs[metric] = {"baseline": b_val, "actual": a_val}

        if metric_diffs:
            mismatches.append({"column": col, "metrics": metric_diffs})

    return {
        "match": len(mismatches) == 0,
        "mismatches": mismatches,
    }

def _validate_cell_level(
        code_unit_name: str,
        database_dialect: DatabaseDialect,
        fully_qualified_name: str,
        baseline_rows: list[RowDict],
        actual_rows: list[RowDict]
    ) -> ValidationResult:
    source_dataframe = pl.DataFrame(baseline_rows)
    source_query_result = QueryExecutionResult(dataframe=source_dataframe)

    target_dataframe = pl.DataFrame(actual_rows)
    target_query_result = QueryExecutionResult(dataframe=target_dataframe)

    platform = Platform(database_dialect.value)
    table_context = TableContext(
            apply_metric_column_modifier=True,
            chunk_number=1,
            column_mappings={},
            column_selection_list=[],
            columns=[],
            database_name="",
            exclude_metrics=False,
            fully_qualified_name=fully_qualified_name,
            has_where_clause=False,
            id=1,
            internal_fully_qualified_name=fully_qualified_name,
            internal_table_name=code_unit_name,
            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,
            row_count=len(baseline_rows),
            run_id="",
            run_start_time="",
            schema_name="",
            sql_generator=None,
            table_name=code_unit_name,
            templates_loader_manager=None,
            user_index_column_collection=[],
            where_clause="",
        )

    cell_comparison = validate_cell(
            source_query_result=source_query_result,
            target_query_result=target_query_result,
            table_context=table_context,
        )
    
    return cell_comparison

# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def compare_in_memory(
    code_unit_name: str,
    database_dialect: DatabaseDialect,
    fully_qualified_name: str,
    baseline_rows: list[RowDict],
    actual_rows: list[RowDict],
    col_types: ColumnTypeMap,
    tolerance: float = ValidationConfiguration().tolerance,
    max_cell_diffs: int = ValidationConfiguration().max_cell_diffs,
) -> dict[str, Any]:
    """Compare two result sets in memory using DuckDB.

    Returns a dict with ``match``, ``row_counts``, ``summary_stats``, and ``cell_diffs``.
    """
    conn, columns = _load_duckdb(baseline_rows, actual_rows, col_types)

    try:
        rc = _row_count_diff(conn, columns)
        ss = _summary_stats_diff(conn, columns, col_types=col_types, tolerance=tolerance)
        cd: list[dict] = []
        if rc["match"]:
            cell_comparison = _validate_cell_level(
                code_unit_name=code_unit_name,
                database_dialect=database_dialect,
                fully_qualified_name=fully_qualified_name,
                baseline_rows=baseline_rows,
                actual_rows=actual_rows)
            # Rename columns to lowercase format expected by downstream code
            df = cell_comparison.results_dataframe.rename(columns={
                "ROW": "row",
                "COLUMN": "column",
                "SOURCE_VALUE": "baseline",
                "TARGET_VALUE": "actual",
            })
            cd = df.head(max_cell_diffs).to_dict(orient="records")
        else:
            LOGGER.warning(
                "Row count mismatch (baseline=%d, actual=%d); skipping cell-level diff.",
                rc["baseline"], rc["actual"],
            )
    finally:
        conn.close()

    overall_match = rc["match"] and ss["match"] and len(cd) == 0
    return {
        "match": overall_match,
        "row_counts": rc,
        "summary_stats": ss,
        "cell_diffs": cd,
    }


class InMemoryResultComparator:
    """Orchestrates in-memory comparison for a single test case.

    Fetches the baseline from the Snowflake stage, executes the procedure,
    and delegates to :func:`compare_in_memory`.
    """

    def __init__(self, config: TestRunnerConfig, connector: ConnectorBase, code_unit_name: str) -> None:
        self._config = config
        self._connector = connector
        self._code_unit_name = code_unit_name

    def compare(
        self,
        rendered_sql: str,
        params_hash: str,
        dump_baseline_and_actual: bool = False,
        fetch_stmts: list[str] | None = None,
        outputs_sql: str | None = None,
        outputs_only: bool = False,
    ) -> ComparisonResult:
        """Fetch baseline from stage, execute *rendered_sql*, and compare in memory.

        When *dump_baseline_and_actual* is True, attaches ``_baseline_rows`` and
        ``_actual_rows`` to the result for ``write_results_local`` to persist.

        When *fetch_stmts* is provided (e.g. ``SELECT * FROM IDENTIFIER('...')``
        for ``output.tables`` procedures), each statement is executed after the
        CALL and the last result set is used as the actual data.
        """
        database: str = self._config.validation_database or ""

        t0 = time.monotonic()
        if outputs_only:
            baseline_rows, col_types = fetch_baseline_output_rows(
                self._connector, database, params_hash
            )
        else:
            baseline_rows, col_types = fetch_baseline_rows(
                self._connector, database, params_hash
            )
        LOGGER.info(
            "[timing] fetch_baseline_rows (%d rows): %.3fs",
            len(baseline_rows), time.monotonic() - t0,
        )

        # Isolation is handled by the caller via clone databases — the
        # rendered_sql already targets a disposable clone so EXECUTE AS
        # OWNER side effects never touch the original database.
        cursor = self._connector.connection.cursor()
        actual_rows: list[RowDict] = []
        try:
            t1 = time.monotonic()
            execution_sql = outputs_sql if outputs_only and outputs_sql else rendered_sql
            execution_fetch_stmts = None if outputs_only else fetch_stmts
            execute_with_capture(cursor, execution_sql, execution_fetch_stmts)

            col_names = (
                [d[0] for d in cursor.description] if cursor.description else list(col_types.keys())
            )
            raw_rows = cursor.fetchall() or []
            actual_rows = [dict(zip(col_names, r)) for r in raw_rows]
            LOGGER.info(
                "[timing] execute CALL + fetchall (%d rows): %.3fs",
                len(actual_rows), time.monotonic() - t1,
            )
        except Exception as exc:
            LOGGER.warning("In-memory comparison failed during execution: %s", exc)
            return ComparisonResult(
                match=False,
                match_type="error",
                match_level="in_memory",
                differences=[
                    ValidationDifference(type="ERROR", detail=truncate(str(exc), MAX_ERROR_LENGTH)),
                ],
            )
        finally:
            cursor.close()

        vc = self._config.validation_configuration or ValidationConfiguration()
        tolerance: float = vc.tolerance
        max_cell_diffs: int = vc.max_cell_diffs

        t2 = time.monotonic()
        diff = compare_in_memory(
            code_unit_name=self._code_unit_name,
            database_dialect=self._config.source_dialect,
            fully_qualified_name=f"{self._config.validation_database}.{self._code_unit_name}",
            baseline_rows=baseline_rows,
            actual_rows=actual_rows,
            col_types=col_types,
            tolerance=tolerance,
            max_cell_diffs=max_cell_diffs,
        )
        LOGGER.info("[timing] compare_in_memory (duckdb): %.3fs", time.monotonic() - t2)

        human_diffs: list[str] = []
        if not diff["match"]:
            rc = diff["row_counts"]
            if not rc["match"]:
                human_diffs.append(
                    f"row_counts: baseline={rc['baseline']} actual={rc['actual']}"
                )
            for m in diff["summary_stats"]["mismatches"]:
                for metric, vals in m["metrics"].items():
                    human_diffs.append(
                        f"column={m['column']}  metric={metric}"
                        f"  baseline={vals['baseline']}  actual={vals['actual']}"
                    )
            for d in diff["cell_diffs"][:10]:
                human_diffs.append(
                    f"cell diff row={d['row']} col={d['column']} "
                    f"baseline={d['baseline']!r} actual={d['actual']!r}"
                )

        match_type = "data_match" if diff["match"] else (
            "row_count_mismatch" if not diff["row_counts"]["match"] else
            "metrics_mismatch" if diff["summary_stats"]["mismatches"] else
            "data_mismatch"
        )

        result = ComparisonResult(
            match=diff["match"],
            match_type=match_type,
            match_level="in_memory",
            differences=human_diffs,
            in_memory_diff=diff,
        )
        if dump_baseline_and_actual:
            result.baseline_rows = baseline_rows
            result.actual_rows = actual_rows
        return result
