"""Tests for test_runner.validate.in_memory_comparator."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

from snowflake.snowflake_data_validation.public import Platform, ValidationResult

from test_runner.common.config import ValidationConfiguration
from test_runner.common.database_dialect import DatabaseDialect
from test_runner.validate.in_memory_comparator import (
    _validate_cell_level,
    compare_in_memory,
    InMemoryResultComparator,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

COL_TYPES = {"ProductName": "VARCHAR", "UnitsSold": "NUMBER"}

# Default values for the new required parameters
_DEFAULT_PROC_NAME = "test_proc"
_DEFAULT_PLATFORM = DatabaseDialect.SNOWFLAKE
_DEFAULT_FQN = "db.test_proc"


def _rows(*tuples) -> list[dict]:
    return [{"ProductName": t[0], "UnitsSold": t[1]} for t in tuples]


def _compare(baseline, actual, col_types, **kwargs):
    """Helper wrapper for compare_in_memory with default procedure metadata."""
    return compare_in_memory(
        code_unit_name=_DEFAULT_PROC_NAME,
        database_dialect=_DEFAULT_PLATFORM,
        fully_qualified_name=_DEFAULT_FQN,
        baseline_rows=baseline,
        actual_rows=actual,
        col_types=col_types,
        **kwargs,
    )


# ---------------------------------------------------------------------------
# TestValidateCellLevel – results_dataframe attribute
# ---------------------------------------------------------------------------


class TestValidateCellLevel:
    """Verify _validate_cell_level returns a ValidationResult with results_dataframe."""

    def test_returns_validation_result(self):
        baseline = _rows(("A", 1))
        actual = _rows(("A", 1))
        result = _validate_cell_level(
            code_unit_name=_DEFAULT_PROC_NAME,
            database_dialect=_DEFAULT_PLATFORM,
            fully_qualified_name=_DEFAULT_FQN,
            baseline_rows=baseline,
            actual_rows=actual,
        )
        assert isinstance(result, ValidationResult)

    def test_has_results_dataframe_attribute(self):
        baseline = _rows(("A", 1))
        actual = _rows(("A", 1))
        result = _validate_cell_level(
            code_unit_name=_DEFAULT_PROC_NAME,
            database_dialect=_DEFAULT_PLATFORM,
            fully_qualified_name=_DEFAULT_FQN,
            baseline_rows=baseline,
            actual_rows=actual,
        )
        assert hasattr(result, "results_dataframe"), (
            "ValidationResult must expose results_dataframe"
        )

    def test_results_dataframe_empty_when_identical(self):
        rows = _rows(("A", 1), ("B", 2))
        result = _validate_cell_level(
            code_unit_name=_DEFAULT_PROC_NAME,
            database_dialect=_DEFAULT_PLATFORM,
            fully_qualified_name=_DEFAULT_FQN,
            baseline_rows=rows,
            actual_rows=rows,
        )
        assert result.is_valid is True
        assert result.results_dataframe.empty

    def test_results_dataframe_populated_when_different(self):
        baseline = _rows(("A", 1), ("B", 2))
        actual = _rows(("A", 1), ("B", 99))
        result = _validate_cell_level(
            code_unit_name=_DEFAULT_PROC_NAME,
            database_dialect=_DEFAULT_PLATFORM,
            fully_qualified_name=_DEFAULT_FQN,
            baseline_rows=baseline,
            actual_rows=actual,
        )
        assert result.is_valid is False
        assert not result.results_dataframe.empty

    def test_results_dataframe_has_expected_columns(self):
        """The upstream dataframe must have ROW/COLUMN/SOURCE_VALUE/TARGET_VALUE
        so compare_in_memory can rename them to lowercase."""
        baseline = _rows(("A", 1))
        actual = _rows(("A", 2))
        result = _validate_cell_level(
            code_unit_name=_DEFAULT_PROC_NAME,
            database_dialect=_DEFAULT_PLATFORM,
            fully_qualified_name=_DEFAULT_FQN,
            baseline_rows=baseline,
            actual_rows=actual,
        )
        expected_cols = {"ROW", "COLUMN", "SOURCE_VALUE", "TARGET_VALUE"}
        assert expected_cols <= set(result.results_dataframe.columns)


class TestCellDiffsUsesResultsDataframe:
    """Ensure compare_in_memory reads results_dataframe."""

    def test_cell_diffs_populated(self):
        """End-to-end: compare_in_memory must produce cell_diffs."""
        baseline = _rows(("A", 10))
        actual = _rows(("A", 20))
        result = _compare(baseline, actual, COL_TYPES)
        assert len(result["cell_diffs"]) >= 1
        diff = result["cell_diffs"][0]
        assert set(diff.keys()) >= {"row", "column", "baseline", "actual"}

    @patch("test_runner.validate.in_memory_comparator._validate_cell_level")
    def test_accesses_results_dataframe_attribute(self, mock_validate):
        """compare_in_memory must read .results_dataframe on the ValidationResult."""
        mock_df = MagicMock()
        mock_df.rename.return_value.head.return_value.to_dict.return_value = []
        mock_result = MagicMock(spec=ValidationResult)
        mock_result.results_dataframe = mock_df
        mock_validate.return_value = mock_result

        rows = _rows(("A", 1))
        _compare(rows, rows, COL_TYPES)

        mock_df.rename.assert_called_once()


# ---------------------------------------------------------------------------
# TestCompareInMemory – row_counts level
# ---------------------------------------------------------------------------

class TestRowCounts:
    def test_identical_rows_match(self):
        rows = _rows(("Widget", 10), ("Gadget", 5))
        result = _compare(rows, rows, COL_TYPES)
        assert result["row_counts"]["match"] is True
        assert result["row_counts"]["baseline"] == 2
        assert result["row_counts"]["actual"] == 2

    def test_different_counts_fail(self):
        baseline = _rows(("Widget", 10), ("Gadget", 5))
        actual = _rows(("Widget", 10))
        result = _compare(baseline, actual, COL_TYPES)
        assert result["row_counts"]["match"] is False
        assert result["match"] is False

    def test_empty_vs_empty_is_match(self):
        result = _compare([], [], COL_TYPES)
        assert result["row_counts"]["match"] is True

    def test_empty_baseline_nonempty_actual_fails(self):
        result = _compare([], _rows(("X", 1)), COL_TYPES)
        assert result["row_counts"]["match"] is False


# ---------------------------------------------------------------------------
# TestCompareInMemory – summary_stats level
# ---------------------------------------------------------------------------

class TestSummaryStats:
    def test_identical_stats_match(self):
        rows = _rows(("A", 10), ("B", 20))
        result = _compare(rows, rows, COL_TYPES)
        assert result["summary_stats"]["match"] is True
        assert result["summary_stats"]["mismatches"] == []

    def test_numeric_difference_detected(self):
        baseline = _rows(("A", 10), ("B", 20))
        actual   = _rows(("A", 10), ("B", 99))  # B's units differ
        result = _compare(baseline, actual, COL_TYPES)
        assert result["summary_stats"]["match"] is False
        cols = {m["column"] for m in result["summary_stats"]["mismatches"]}
        assert "UnitsSold" in cols
        # Every entry has column + metrics dict
        for m in result["summary_stats"]["mismatches"]:
            assert {"column", "metrics"} <= set(m.keys())
            assert isinstance(m["metrics"], dict)

    def test_only_differing_metrics_reported(self):
        baseline = _rows(("A", 10), ("B", 20))
        actual   = _rows(("A", 10), ("B", 99))
        result = _compare(baseline, actual, COL_TYPES)
        units_entry = next(m for m in result["summary_stats"]["mismatches"]
                           if m["column"] == "UnitsSold")
        # max and sum differ; min does NOT (both have min=10)
        assert "max" in units_entry["metrics"]
        assert "sum" in units_entry["metrics"]
        assert "min" not in units_entry["metrics"]
        # Each metric value has baseline and actual keys
        for vals in units_entry["metrics"].values():
            assert "baseline" in vals and "actual" in vals

    def test_numeric_strings_compared_as_numbers(self):
        """Baseline numeric strings are coerced before loading — no lexicographic traps."""
        from test_runner.validate.transient_table_manager import coerce_numeric
        # Simulate what fetch_baseline_rows now does before handing rows to DuckDB
        col_types = {"TOTALVALUE": "NUMBER"}
        baseline_raw = [{"TOTALVALUE": "145683238.00"}, {"TOTALVALUE": "17853738.00"}]
        actual_raw   = [{"TOTALVALUE": 145683238.00},   {"TOTALVALUE": 17853738.00}]
        baseline = [{k: coerce_numeric(v) for k, v in r.items()} for r in baseline_raw]
        result = _compare(baseline, actual_raw, col_types)
        assert result["summary_stats"]["match"] is True
        assert result["match"] is True

    def test_varchar_columns_skipped_in_summary_stats(self):
        """VARCHAR columns must not appear in summary_stats mismatches."""
        col_types = {"NAME": "VARCHAR", "AMOUNT": "NUMBER"}
        baseline = [{"NAME": "Alice", "AMOUNT": 100}, {"NAME": "Bob", "AMOUNT": 200}]
        actual   = [{"NAME": "Alice", "AMOUNT": 100}, {"NAME": "Charlie", "AMOUNT": 200}]
        result = _compare(baseline, actual, col_types)
        # NAME differs but must NOT appear in summary_stats — only in cell_diffs
        cols_in_stats = {m["column"] for m in result["summary_stats"]["mismatches"]}
        assert "NAME" not in cols_in_stats
        # Cell diff should catch the NAME change
        assert any(d["column"] == "NAME" for d in result["cell_diffs"])

    def test_null_count_difference_detected(self):
        baseline = [{"ProductName": "A", "UnitsSold": None}, {"ProductName": "B", "UnitsSold": 5}]
        actual   = [{"ProductName": "A", "UnitsSold": 10},   {"ProductName": "B", "UnitsSold": 5}]
        result = _compare(baseline, actual, COL_TYPES)
        assert result["summary_stats"]["match"] is False
        units_entry = next(m for m in result["summary_stats"]["mismatches"]
                           if m["column"] == "UnitsSold")
        assert "null_count" in units_entry["metrics"]


# ---------------------------------------------------------------------------
# TestCompareInMemory – cell_diffs level
# ---------------------------------------------------------------------------

class TestCellDiffs:
    def test_no_diffs_when_identical(self):
        rows = _rows(("A", 1), ("B", 2))
        result = _compare(rows, rows, COL_TYPES)
        assert result["cell_diffs"] == []
        assert result["match"] is True

    def test_detects_single_cell_change(self):
        baseline = _rows(("A", 1), ("B", 2))
        actual   = _rows(("A", 1), ("B", 99))
        result = _compare(baseline, actual, COL_TYPES)
        assert len(result["cell_diffs"]) == 1
        diff = result["cell_diffs"][0]
        assert diff["column"] == "UnitsSold"
        assert diff["row"] == 1  # 0-based index from validate_cell
        assert str(diff["baseline"]) == "2"
        assert str(diff["actual"]) == "99"

    def test_cell_diff_includes_row_col_baseline_actual_keys(self):
        baseline = _rows(("X", 10))
        actual   = _rows(("Y", 10))
        result = _compare(baseline, actual, COL_TYPES)
        assert any(d["column"] == "ProductName" for d in result["cell_diffs"])
        diff = next(d for d in result["cell_diffs"] if d["column"] == "ProductName")
        assert set(diff.keys()) >= {"row", "column", "baseline", "actual"}

    def test_cell_diffs_skipped_on_row_count_mismatch(self):
        baseline = _rows(("A", 1), ("B", 2))
        actual   = _rows(("A", 1))
        result = _compare(baseline, actual, COL_TYPES)
        assert result["cell_diffs"] == []

    def test_cell_diffs_capped_at_max(self):
        # Generate more rows than the default cap (ValidationConfiguration.max_cell_diffs = 1000)
        default_cap = ValidationConfiguration().max_cell_diffs
        n = default_cap + 5
        baseline = [{"Col": f"base-{i}"} for i in range(n)]
        actual   = [{"Col": f"actual-{i}"} for i in range(n)]
        result = _compare(baseline, actual, {"Col": "VARCHAR"})
        assert len(result["cell_diffs"]) <= default_cap

    def test_cell_diffs_capped_at_custom_max(self):
        """max_cell_diffs parameter controls the ceiling."""
        cap = 5
        n = cap + 3
        baseline = [{"Col": f"base-{i}"} for i in range(n)]
        actual   = [{"Col": f"actual-{i}"} for i in range(n)]
        result = _compare(baseline, actual, {"Col": "VARCHAR"}, max_cell_diffs=cap)
        assert len(result["cell_diffs"]) <= cap


# ---------------------------------------------------------------------------
# TestCompareInMemory – overall match key
# ---------------------------------------------------------------------------

class TestOverallMatch:
    def test_perfect_match(self):
        rows = _rows(("A", 1))
        assert _compare(rows, rows, COL_TYPES)["match"] is True

    def test_any_diff_means_no_match(self):
        baseline = _rows(("A", 1))
        actual   = _rows(("A", 2))
        assert _compare(baseline, actual, COL_TYPES)["match"] is False


# ---------------------------------------------------------------------------
# TestInMemoryResultComparator
# ---------------------------------------------------------------------------

class TestSummaryStatsTolerance:
    """Epsilon / tolerance behaviour in _summary_stats_diff."""

    def test_difference_within_tolerance_is_ignored(self):
        """A diff smaller than the tolerance must NOT appear in mismatches."""
        baseline = [{"V": 100.0}, {"V": 200.0}]
        actual   = [{"V": 100.0}, {"V": 200.0005}]  # avg/sum differ by < 0.001
        col_types = {"V": "NUMBER"}
        result = _compare(baseline, actual, col_types, tolerance=0.001)
        assert result["summary_stats"]["match"] is True, result["summary_stats"]

    def test_difference_exceeding_tolerance_is_reported(self):
        """A diff larger than the tolerance MUST appear in mismatches."""
        baseline = [{"V": 100.0}, {"V": 200.0}]
        actual   = [{"V": 100.0}, {"V": 200.002}]  # avg/sum differ by > 0.001
        col_types = {"V": "NUMBER"}
        result = _compare(baseline, actual, col_types, tolerance=0.001)
        assert result["summary_stats"]["match"] is False, result["summary_stats"]

    def test_default_tolerance_applied(self):
        """compare_in_memory default tolerance comes from ValidationConfiguration."""
        assert ValidationConfiguration().tolerance == 0.001
        # A tiny diff that falls within the default tolerance should be ignored.
        baseline = [{"V": 1000.0}]
        actual   = [{"V": 1000.0009}]
        col_types = {"V": "NUMBER"}
        result = _compare(baseline, actual, col_types)
        assert result["summary_stats"]["match"] is True


class TestInMemoryResultComparator:
    def _make_config(
        self,
        database: str = "MYDB",
        tolerance: float | None = None,
        max_cell_diffs: int | None = None,
    ) -> MagicMock:
        cfg = MagicMock()
        cfg.validation_database = database
        cfg.source_dialect = DatabaseDialect.SNOWFLAKE
        vc_kwargs: dict = {}
        if tolerance is not None:
            vc_kwargs["tolerance"] = tolerance
        if max_cell_diffs is not None:
            vc_kwargs["max_cell_diffs"] = max_cell_diffs
        cfg.validation_configuration = ValidationConfiguration(**vc_kwargs)
        return cfg

    def _make_connector(self, actual_tuples: list[tuple], col_names: list[str]) -> MagicMock:
        cursor = MagicMock()
        cursor.description = [(c,) for c in col_names]
        cursor.fetchall.return_value = actual_tuples
        connector = MagicMock()
        connector.connection.cursor.return_value = cursor
        return connector

    @patch("test_runner.validate.in_memory_comparator.fetch_baseline_rows")
    def test_pass_when_rows_match(self, mock_fetch):
        mock_fetch.return_value = (
            [{"ProductName": "A", "UnitsSold": 10}],
            {"ProductName": "VARCHAR", "UnitsSold": "NUMBER"},
        )
        connector = self._make_connector([("A", 10)], ["ProductName", "UnitsSold"])
        comp = InMemoryResultComparator(self._make_config(), connector, "dbo.Proc")
        result = comp.compare("CALL dbo.Proc()", "hash1")
        assert result.match is True
        assert result.match_type == "data_match"
        assert result.in_memory_diff["match"] is True

    @patch("test_runner.validate.in_memory_comparator.fetch_baseline_rows")
    def test_fail_when_rows_differ(self, mock_fetch):
        mock_fetch.return_value = (
            [{"ProductName": "A", "UnitsSold": 10}],
            {"ProductName": "VARCHAR", "UnitsSold": "NUMBER"},
        )
        connector = self._make_connector([("A", 99)], ["ProductName", "UnitsSold"])
        comp = InMemoryResultComparator(self._make_config(), connector, "dbo.Proc")
        result = comp.compare("CALL dbo.Proc()", "hash1")
        assert result.match is False
        assert result.match_type != "data_match"
        assert result.in_memory_diff["match"] is False

    @patch("test_runner.validate.in_memory_comparator.fetch_baseline_rows")
    def test_uses_database_from_config(self, mock_fetch):
        mock_fetch.return_value = ([], {})
        connector = self._make_connector([], [])
        comp = InMemoryResultComparator(self._make_config("PRODDB"), connector, "dbo.P")
        comp.compare("CALL dbo.P()", "h")
        mock_fetch.assert_called_once_with(connector, "PRODDB", "h")

    @patch("test_runner.validate.in_memory_comparator.fetch_baseline_rows")
    def test_differences_list_populated_on_fail(self, mock_fetch):
        mock_fetch.return_value = (
            [{"X": 1}],
            {"X": "NUMBER"},
        )
        connector = self._make_connector([(2,)], ["X"])
        comp = InMemoryResultComparator(self._make_config(), connector, "P")
        result = comp.compare("CALL P()", "h")
        assert result.differences  # non-empty
        assert result.match is False

    @patch("test_runner.validate.in_memory_comparator.fetch_baseline_rows")
    def test_result_has_runner_compatible_attrs(self, mock_fetch):
        mock_fetch.return_value = ([], {})
        connector = self._make_connector([], [])
        comp = InMemoryResultComparator(self._make_config(), connector, "P")
        result = comp.compare("CALL P()", "h")
        # Must have all attributes that runner.py accesses
        assert hasattr(result, "match")
        assert hasattr(result, "match_type")
        assert hasattr(result, "match_level")
        assert hasattr(result, "differences")
        assert hasattr(result, "in_memory_diff")

    @patch("test_runner.validate.in_memory_comparator.fetch_baseline_rows")
    def test_cursor_closed_after_execution(self, mock_fetch):
        mock_fetch.return_value = ([], {})
        connector = self._make_connector([], [])
        comp = InMemoryResultComparator(self._make_config(), connector, "P")
        comp.compare("CALL P()", "h")
        connector.connection.cursor.return_value.close.assert_called()

    @patch("test_runner.validate.in_memory_comparator.fetch_baseline_rows")
    def test_tolerance_from_config_applies_to_summary_stats(self, mock_fetch):
        """Tolerance applies to summary_stats but cell diffs use exact comparison."""
        # baseline sum=300.0, actual sum=300.0009 — within tolerance=0.001
        mock_fetch.return_value = (
            [{"V": 100.0}, {"V": 200.0}],
            {"V": "NUMBER"},
        )
        connector = self._make_connector([(100.0,), (200.0009,)], ["V"])
        comp = InMemoryResultComparator(self._make_config(tolerance=0.001), connector, "P")
        result = comp.compare("CALL P()", "h")
        # Summary stats match (within tolerance), but cell diff detects exact difference
        assert result.in_memory_diff["summary_stats"]["match"] is True
        assert result.match is False  # Cell diff causes overall mismatch

    @patch("test_runner.validate.in_memory_comparator.fetch_baseline_rows")
    def test_tolerance_from_config_catches_large_float_diff(self, mock_fetch):
        """Comparator must not suppress diffs that exceed the config tolerance."""
        # baseline sum=300.0, actual sum=300.002 — exceeds tolerance=0.001
        mock_fetch.return_value = (
            [{"V": 100.0}, {"V": 200.0}],
            {"V": "NUMBER"},
        )
        connector = self._make_connector([(100.0,), (200.002,)], ["V"])
        comp = InMemoryResultComparator(self._make_config(tolerance=0.001), connector, "P")
        result = comp.compare("CALL P()", "h")
        assert result.match is False

    @patch("test_runner.validate.in_memory_comparator.fetch_baseline_rows")
    def test_max_cell_diffs_from_config_limits_output(self, mock_fetch):
        """Comparator must read max_cell_diffs from config and cap cell diffs."""
        n = 10
        mock_fetch.return_value = (
            [{"C": f"b{i}"} for i in range(n)],
            {"C": "VARCHAR"},
        )
        connector = self._make_connector([(f"a{i}",) for i in range(n)], ["C"])
        comp = InMemoryResultComparator(self._make_config(max_cell_diffs=3), connector, "P")
        result = comp.compare("CALL P()", "h")
        assert len(result.in_memory_diff["cell_diffs"]) <= 3

    @patch("test_runner.validate.in_memory_comparator.fetch_baseline_output_rows")
    def test_outputs_only_uses_output_baseline_rows(self, mock_fetch_outputs):
        mock_fetch_outputs.return_value = (
            [{"OUT_COL": 21}],
            {"OUT_COL": "NUMBER"},
        )
        connector = self._make_connector([(21,)], ["OUT_COL"])
        comp = InMemoryResultComparator(self._make_config("MYDB"), connector, "dbo.P")
        result = comp.compare(
            "CALL dbo.P()", "h", outputs_only=True, outputs_sql="BEGIN SELECT 21 AS OUT_COL; END;"
        )
        assert result.match is True
        mock_fetch_outputs.assert_called_once_with(connector, "MYDB", "h")

    @patch("test_runner.validate.in_memory_comparator.fetch_baseline_output_rows")
    @patch("test_runner.validate.in_memory_comparator.execute_with_capture")
    def test_outputs_only_executes_outputs_sql_without_fetch_stmts(
        self, mock_execute_with_capture, mock_fetch_outputs,
    ):
        mock_fetch_outputs.return_value = (
            [{"OUT_COL": 1}],
            {"OUT_COL": "NUMBER"},
        )
        connector = self._make_connector([(1,)], ["OUT_COL"])
        comp = InMemoryResultComparator(self._make_config(), connector, "dbo.P")
        comp.compare(
            "CALL dbo.P()", "h", outputs_only=True, outputs_sql="BEGIN SELECT 1 AS OUT_COL; END;"
        )
        call_args = mock_execute_with_capture.call_args
        assert call_args.args[1] == "BEGIN SELECT 1 AS OUT_COL; END;"
        assert call_args.args[2] is None

    @patch("test_runner.validate.in_memory_comparator.fetch_baseline_rows")
    def test_dump_baseline_and_actual_attaches_rows(self, mock_fetch):
        baseline_rows = [{"ProductName": "A", "UnitsSold": 10}]
        mock_fetch.return_value = (
            baseline_rows,
            {"ProductName": "VARCHAR", "UnitsSold": "NUMBER"},
        )
        connector = self._make_connector([("A", 10)], ["ProductName", "UnitsSold"])
        comp = InMemoryResultComparator(self._make_config(), connector, "dbo.Proc")
        result = comp.compare("CALL dbo.Proc()", "hash1", dump_baseline_and_actual=True)
        assert result.baseline_rows == baseline_rows
        assert result.actual_rows == [{"ProductName": "A", "UnitsSold": 10}]
