"""
Tests for the data_validation package (handler, helpers, row/cell validators).

Covers pure helper functions, I/O utilities, and handler routing logic.
"""

import csv
import os
import shutil
import tempfile

import pytest
from unittest.mock import Mock, patch

try:
    from data_exchange_agent.tasks.data_validation.helpers import (
        CSV_DELIMITER,
        adapt_sdv_result,
        build_row_table_context,
        resolve_platform,
    )
    from data_exchange_agent.tasks.data_validation.handler import (
        DataValidationHandler,
    )
    from data_exchange_agent.tasks.data_validation.row_validation import (
        RowValidationHandler,
    )

    HANDLER_AVAILABLE = True
except ImportError:
    HANDLER_AVAILABLE = False

try:
    import pandas as pd

    PANDAS_AVAILABLE = True
except ImportError:
    PANDAS_AVAILABLE = False

try:
    from snowflake.snowflake_data_validation.public import Origin, Platform

    SDV_AVAILABLE = True
except ImportError:
    SDV_AVAILABLE = False

try:
    import pyodbc  # noqa: F401

    PYODBC_AVAILABLE = True
except (ImportError, OSError):
    PYODBC_AVAILABLE = False


SKIP_MSG = "data_validation handler or SDV library not available"


# =========================================================================
# Group A: Pure / static functions (no mocking required)
# =========================================================================


class TestDatabaseEngines:
    """Tests for database_engines module functions."""

    def test_is_database_engine_supported(self):
        from data_exchange_agent.data_sources.database_engines import (
            DatabaseEngine,
            is_database_engine_supported,
        )

        assert is_database_engine_supported(DatabaseEngine.ORACLE) is True
        assert is_database_engine_supported(DatabaseEngine.SNOWFLAKE) is True

    def test_get_database_engine_from_string(self):
        from data_exchange_agent.data_sources.database_engines import (
            DatabaseEngine,
            get_database_engine_from_string,
        )

        assert get_database_engine_from_string("oracle") == DatabaseEngine.ORACLE
        assert get_database_engine_from_string("SQLSERVER") == DatabaseEngine.SQLSERVER

    def test_get_database_engine_from_string_invalid(self):
        from data_exchange_agent.data_sources.database_engines import (
            get_database_engine_from_string,
        )

        with pytest.raises(ValueError, match="Unsupported database engine"):
            get_database_engine_from_string("nosql_db")


@pytest.mark.skipif(not HANDLER_AVAILABLE or not SDV_AVAILABLE, reason=SKIP_MSG)
class TestLoadDatatypesMappings:
    """Tests for helpers.load_datatypes_mappings."""

    def test_returns_empty_dict_for_nonexistent_mapping(self):
        from data_exchange_agent.tasks.data_validation.helpers import (
            load_datatypes_mappings,
        )

        with patch(
            "data_exchange_agent.tasks.data_validation.helpers._templates_path",
            return_value="/nonexistent/path",
        ):
            result = load_datatypes_mappings(Platform.REDSHIFT)

        assert result == {}


@pytest.mark.skipif(not HANDLER_AVAILABLE or not SDV_AVAILABLE, reason=SKIP_MSG)
class TestResolvePlatform:
    """Tests for helpers.resolve_platform."""

    def test_known_platforms(self):
        assert resolve_platform("redshift") == Platform.REDSHIFT
        assert resolve_platform("snowflake") == Platform.SNOWFLAKE
        assert resolve_platform("sqlserver") == Platform.SQLSERVER
        assert resolve_platform("sql_server") == Platform.SQLSERVER
        assert resolve_platform("teradata") == Platform.TERADATA

    def test_case_insensitive(self):
        assert resolve_platform("Redshift") == Platform.REDSHIFT
        assert resolve_platform("SNOWFLAKE") == Platform.SNOWFLAKE

    def test_unknown_raises(self):
        with pytest.raises(ValueError, match="Unknown platform"):
            resolve_platform("mysql")


@pytest.mark.skipif(not HANDLER_AVAILABLE or not PANDAS_AVAILABLE, reason=SKIP_MSG)
class TestAdaptSdvResult:
    """Tests for helpers.adapt_sdv_result."""

    def test_empty_dataframe_returns_expected_columns(self):
        df = adapt_sdv_result(pd.DataFrame(), "my_table")
        assert list(df.columns) == [
            "VALIDATION_TYPE",
            "TABLE_NAME",
            "COLUMN_VALIDATED",
            "EVALUATION_CRITERIA",
            "SOURCE_VALUE",
            "SNOWFLAKE_VALUE",
            "STATUS",
            "COMMENTS",
        ]
        assert len(df) == 0

    def test_table_column_renamed(self):
        raw = pd.DataFrame(
            {
                "VALIDATION_TYPE": ["schema"],
                "TABLE": ["src_table"],
                "COLUMN_VALIDATED": ["col1"],
                "EVALUATION_CRITERIA": ["type_match"],
                "SOURCE_VALUE": ["INT"],
                "SNOWFLAKE_VALUE": ["NUMBER"],
                "STATUS": ["FAILURE"],
                "COMMENTS": [""],
            }
        )
        result = adapt_sdv_result(raw, "src_table")
        assert "TABLE_NAME" in result.columns
        assert "TABLE" not in result.columns
        assert result["TABLE_NAME"].iloc[0] == "src_table"

    def test_missing_columns_filled(self):
        raw = pd.DataFrame({"VALIDATION_TYPE": ["metrics"], "STATUS": ["SUCCESS"]})
        result = adapt_sdv_result(raw, "t")
        assert result["TABLE_NAME"].iloc[0] == "t"
        assert result["COMMENTS"].iloc[0] == ""

    def test_table_name_injected_when_missing(self):
        raw = pd.DataFrame(
            {
                "VALIDATION_TYPE": ["schema"],
                "COLUMN_VALIDATED": ["c"],
                "EVALUATION_CRITERIA": ["e"],
                "SOURCE_VALUE": ["s"],
                "SNOWFLAKE_VALUE": ["t"],
                "STATUS": ["OK"],
                "COMMENTS": [""],
            }
        )
        result = adapt_sdv_result(raw, "injected")
        assert result["TABLE_NAME"].iloc[0] == "injected"


@pytest.mark.skipif(not HANDLER_AVAILABLE or not PANDAS_AVAILABLE, reason=SKIP_MSG)
class TestBuildHashesDataframe:
    """Tests for RowValidationHandler.build_hashes_dataframe."""

    def _make_chunks(self, chunk_ids, hashes):
        return pd.DataFrame({"CHUNK_ID": chunk_ids, "CHUNK_MD5_VALUE": hashes})

    def test_matching_hashes(self):
        src = self._make_chunks([1, 2], ["abc", "def"])
        tgt = self._make_chunks([1, 2], ["abc", "def"])
        df = RowValidationHandler.build_hashes_dataframe(src, tgt, "tbl")
        assert list(df["IS_MATCH"]) == [True, True]
        assert list(df["TABLE_NAME"]) == ["tbl", "tbl"]

    def test_differing_hashes(self):
        src = self._make_chunks([1, 2], ["abc", "def"])
        tgt = self._make_chunks([1, 2], ["abc", "xyz"])
        df = RowValidationHandler.build_hashes_dataframe(src, tgt, "tbl")
        assert df.loc[df["CHUNK_ID"] == 1, "IS_MATCH"].iloc[0] == True  # noqa: E712
        assert df.loc[df["CHUNK_ID"] == 2, "IS_MATCH"].iloc[0] == False  # noqa: E712

    def test_missing_chunk_on_one_side(self):
        src = self._make_chunks([1], ["aaa"])
        tgt = self._make_chunks([1, 2], ["aaa", "bbb"])
        df = RowValidationHandler.build_hashes_dataframe(src, tgt, "tbl")
        assert len(df) == 2
        row_2 = df.loc[df["CHUNK_ID"] == 2].iloc[0]
        assert row_2["SOURCE_HASH"] == ""
        assert row_2["IS_MATCH"] == False  # noqa: E712

    def test_empty_hash_not_considered_match(self):
        src = self._make_chunks([1], [""])
        tgt = self._make_chunks([1], [""])
        df = RowValidationHandler.build_hashes_dataframe(src, tgt, "tbl")
        assert df["IS_MATCH"].iloc[0] == False  # noqa: E712


@pytest.mark.skipif(not HANDLER_AVAILABLE or not SDV_AVAILABLE, reason=SKIP_MSG)
class TestBuildRowTableContext:
    """Tests for helpers.build_row_table_context."""

    def test_three_part_fqn(self):
        ctx = build_row_table_context(
            platform=Platform.SNOWFLAKE,
            origin=Origin.SOURCE,
            table_fqn="db.schema.tbl",
            columns=[],
            index_column_list=["id"],
            chunk_number=2,
            max_failed_rows_number=50,
            row_count=100,
        )
        assert ctx.database_name == "db"
        assert ctx.schema_name == "schema"
        assert ctx.table_name == "tbl"
        assert ctx.row_count == 100
        assert ctx.chunk_number == 2

    def test_two_part_fqn(self):
        ctx = build_row_table_context(
            platform=Platform.REDSHIFT,
            origin=Origin.TARGET,
            table_fqn="schema.tbl",
            columns=[],
            index_column_list=[],
            chunk_number=1,
            max_failed_rows_number=100,
        )
        assert ctx.database_name == ""
        assert ctx.schema_name == "schema"
        assert ctx.table_name == "tbl"

    def test_one_part_fqn(self):
        ctx = build_row_table_context(
            platform=Platform.SNOWFLAKE,
            origin=Origin.SOURCE,
            table_fqn="just_table",
            columns=[],
            index_column_list=[],
            chunk_number=1,
            max_failed_rows_number=100,
        )
        assert ctx.database_name == ""
        assert ctx.schema_name == ""
        assert ctx.table_name == "just_table"


# =========================================================================
# Group B: I/O functions (minimal mocking)
# =========================================================================


@pytest.mark.skipif(not HANDLER_AVAILABLE, reason=SKIP_MSG)
class TestCleanupTempDir:
    """Tests for DataValidationHandler._cleanup_temp_dir."""

    def test_removes_directory_and_contents(self):
        d = tempfile.mkdtemp()
        open(os.path.join(d, "file.txt"), "w").close()
        DataValidationHandler._cleanup_temp_dir(d)
        assert not os.path.exists(d)

    def test_nonexistent_directory_does_not_raise(self):
        DataValidationHandler._cleanup_temp_dir("/tmp/nonexistent_dir_xyz_abc")


@pytest.mark.skipif(not HANDLER_AVAILABLE or not PANDAS_AVAILABLE, reason=SKIP_MSG)
class TestWriteL3Csv:
    """Tests for DataValidationHandler._write_l3_csv."""

    def _make_handler(self):
        return DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )

    def test_csv_prepends_workflow_and_metadata_ids(self):
        handler = self._make_handler()
        d = tempfile.mkdtemp()
        try:
            path = handler._write_l3_csv(d, "test.csv", [["a", "b"], ["c", "d"]], 42, 7)
            with open(path, newline="", encoding="utf-8") as f:
                reader = csv.reader(f, delimiter=CSV_DELIMITER)
                rows = list(reader)
            assert len(rows) == 2
            assert rows[0][:2] == ["42", "7"]
            assert rows[0][2:] == ["a", "b"]
            assert rows[1][:2] == ["42", "7"]
        finally:
            shutil.rmtree(d)

    def test_empty_data_produces_empty_file(self):
        handler = self._make_handler()
        d = tempfile.mkdtemp()
        try:
            path = handler._write_l3_csv(d, "empty.csv", [], 1, 1)
            with open(path, encoding="utf-8") as f:
                assert f.read() == ""
        finally:
            shutil.rmtree(d)


@pytest.mark.skipif(not HANDLER_AVAILABLE or not PANDAS_AVAILABLE, reason=SKIP_MSG)
class TestWriteResultCsv:
    """Tests for DataValidationHandler._write_result_csv."""

    def _make_handler(self):
        return DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )

    def test_csv_includes_orchestration_columns(self):
        handler = self._make_handler()
        df = pd.DataFrame(
            {
                "VALIDATION_TYPE": ["schema"],
                "TABLE_NAME": ["t"],
                "STATUS": ["SUCCESS"],
            }
        )
        path = handler._write_result_csv(df, 10, 20)
        try:
            with open(path, newline="", encoding="utf-8") as f:
                reader = csv.reader(f, delimiter=CSV_DELIMITER)
                rows = list(reader)
            assert len(rows) == 1
            assert rows[0][0] == "10"
            assert rows[0][1] == "20"
            assert rows[0][2] == "schema"
        finally:
            shutil.rmtree(os.path.dirname(path))


# =========================================================================
# Group C: Handler methods (require mocking)
# =========================================================================


@pytest.mark.skipif(not HANDLER_AVAILABLE or not SDV_AVAILABLE, reason=SKIP_MSG)
class TestHandleRouting:
    """Tests for DataValidationHandler.handle() dispatch logic."""

    def _make_handler(self):
        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )
        h._row_validator = Mock()
        h._cell_validator = Mock()
        return h

    def test_row_mode_dispatches_to_row_validator(self):
        h = self._make_handler()
        h._row_validator.handle.return_value = {"status": "success"}
        task = {
            "payload": {
                "validation_type": "row",
                "row_validation_mode": "row",
                "table_name": "t",
            }
        }
        result = h.handle(task)
        h._row_validator.handle.assert_called_once()
        assert result["status"] == "success"

    def test_cell_mode_dispatches_to_cell_validator(self):
        h = self._make_handler()
        h._cell_validator.handle.return_value = {"status": "success"}
        task = {
            "payload": {
                "validation_type": "row",
                "row_validation_mode": "cell",
                "table_name": "t",
            }
        }
        result = h.handle(task)
        h._cell_validator.handle.assert_called_once()
        assert result["status"] == "success"

    def test_unknown_type_returns_error(self):
        h = self._make_handler()
        h._run_source_query = Mock(return_value=pd.DataFrame())
        h._run_target_query = Mock(return_value=pd.DataFrame())
        task = {
            "payload": {
                "validation_type": "nonexistent",
                "table_name": "t",
                "source_query": "",
                "target_query": "",
            }
        }
        result = h.handle(task)
        assert result["status"] == "error"
        assert "Unknown validation_type" in result["details"]

    def test_schema_type_calls_validate_schema(self):
        h = self._make_handler()
        h._run_source_query = Mock(return_value=pd.DataFrame({"A": [1]}))
        h._run_target_query = Mock(return_value=pd.DataFrame({"A": [1]}))
        h._upload_to_stage = Mock()
        with patch.object(h, "_validate_schema_sdv") as mock_val:
            mock_val.return_value = pd.DataFrame(
                {
                    "VALIDATION_TYPE": ["schema"],
                    "TABLE_NAME": ["t"],
                    "COLUMN_VALIDATED": ["A"],
                    "EVALUATION_CRITERIA": [""],
                    "SOURCE_VALUE": [""],
                    "SNOWFLAKE_VALUE": [""],
                    "STATUS": ["SUCCESS"],
                    "COMMENTS": [""],
                }
            )
            task = {
                "payload": {
                    "validation_type": "schema",
                    "table_name": "t",
                    "source_query": "SELECT 1",
                    "target_query": "SELECT 1",
                    "source_platform": "redshift",
                }
            }
            result = h.handle(task)
        mock_val.assert_called_once()
        assert result["status"] == "success"

    def test_schema_with_explicit_datatypes_mappings_skips_load(self):
        h = self._make_handler()
        h._run_source_query = Mock(return_value=pd.DataFrame({"A": [1]}))
        h._run_target_query = Mock(return_value=pd.DataFrame({"A": [1]}))
        h._upload_to_stage = Mock()
        with (
            patch.object(h, "_validate_schema_sdv") as mock_val,
            patch(
                "data_exchange_agent.tasks.data_validation.handler.load_datatypes_mappings",
            ) as mock_load,
        ):
            mock_val.return_value = pd.DataFrame(
                {
                    "VALIDATION_TYPE": ["schema"],
                    "TABLE_NAME": ["t"],
                    "COLUMN_VALIDATED": ["A"],
                    "EVALUATION_CRITERIA": [""],
                    "SOURCE_VALUE": [""],
                    "SNOWFLAKE_VALUE": [""],
                    "STATUS": ["SUCCESS"],
                    "COMMENTS": [""],
                }
            )
            task = {
                "payload": {
                    "validation_type": "schema",
                    "table_name": "t",
                    "source_query": "SELECT 1",
                    "target_query": "SELECT 1",
                    "source_platform": "redshift",
                    "datatypes_mappings": {"int": "NUMBER"},
                }
            }
            result = h.handle(task)
        mock_load.assert_not_called()
        assert result["status"] == "success"

    def test_metrics_type_calls_validate_metrics(self):
        h = self._make_handler()
        h._run_source_query = Mock(return_value=pd.DataFrame({"A": [1]}))
        h._run_target_query = Mock(return_value=pd.DataFrame({"A": [1]}))
        h._upload_to_stage = Mock()
        with patch.object(h, "_validate_metrics_sdv") as mock_val:
            mock_val.return_value = pd.DataFrame(
                {
                    "VALIDATION_TYPE": ["metrics"],
                    "TABLE_NAME": ["t"],
                    "COLUMN_VALIDATED": ["A"],
                    "EVALUATION_CRITERIA": [""],
                    "SOURCE_VALUE": [""],
                    "SNOWFLAKE_VALUE": [""],
                    "STATUS": ["SUCCESS"],
                    "COMMENTS": [""],
                }
            )
            task = {
                "payload": {
                    "validation_type": "metrics",
                    "table_name": "t",
                    "source_query": "SELECT 1",
                    "target_query": "SELECT 1",
                    "source_platform": "redshift",
                }
            }
            result = h.handle(task)
        mock_val.assert_called_once()
        assert result["status"] == "success"


@pytest.mark.skipif(
    not HANDLER_AVAILABLE or not SDV_AVAILABLE or not PANDAS_AVAILABLE,
    reason=SKIP_MSG,
)
class TestGetTargetColumnMetadata:
    """Tests for RowValidationHandler._get_target_column_metadata."""

    def _make_row_handler(self):
        parent = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )
        return RowValidationHandler(parent)

    def test_returns_columns_from_information_schema(self):
        rv = self._make_row_handler()
        rv._parent._run_target_query = Mock(
            return_value=pd.DataFrame(
                {
                    "COLUMN_NAME": ["ID", "NAME"],
                    "DATA_TYPE": ["NUMBER", "VARCHAR"],
                }
            )
        )
        result = rv._get_target_column_metadata("db.schema.tbl", [])
        assert len(result) == 2
        assert result[0].name == "ID"
        assert result[0].data_type == "NUMBER"

    def test_falls_back_on_empty_result(self):
        rv = self._make_row_handler()
        rv._parent._run_target_query = Mock(return_value=pd.DataFrame())
        fallback = [Mock(name="col1")]
        result = rv._get_target_column_metadata("db.schema.tbl", fallback)
        assert result is fallback

    def test_falls_back_on_exception(self):
        rv = self._make_row_handler()
        rv._parent._run_target_query = Mock(side_effect=RuntimeError("boom"))
        fallback = [Mock(name="col1")]
        result = rv._get_target_column_metadata("db.schema.tbl", fallback)
        assert result is fallback

    def test_single_part_fqn_returns_fallback(self):
        rv = self._make_row_handler()
        fallback = [Mock(name="col1")]
        result = rv._get_target_column_metadata("just_table", fallback)
        assert result is fallback


@pytest.mark.skipif(not HANDLER_AVAILABLE or not PANDAS_AVAILABLE, reason=SKIP_MSG)
class TestWriteEmptyRowResults:
    """Tests for RowValidationHandler._write_empty_row_results."""

    def test_uploads_three_csv_files(self):
        parent = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )
        parent._upload_to_stage = Mock()
        rv = RowValidationHandler(parent)

        rv._write_empty_row_results(1, 2, "tbl", "db.sch.tbl", "@stage/path")

        assert parent._upload_to_stage.call_count == 3
        call_args = [c[0][1] for c in parent._upload_to_stage.call_args_list]
        assert "@stage/path/diffs" in call_args
        assert "@stage/path/summary" in call_args
        assert "@stage/path/hashes" in call_args


@pytest.mark.skipif(
    not HANDLER_AVAILABLE or not SDV_AVAILABLE or not PANDAS_AVAILABLE,
    reason=SKIP_MSG,
)
class TestRowValidationHandleEntry:
    """Tests for RowValidationHandler.handle entry-point validation."""

    def _make_row_handler(self):
        parent = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )
        return RowValidationHandler(parent)

    def test_missing_index_column_list_returns_error(self):
        rv = self._make_row_handler()
        result = rv.handle(
            {},
            {
                "table_name": "t",
                "index_column_list": [],
            },
        )
        assert result["status"] == "error"
        assert "index_column_list" in result["details"]


@pytest.mark.skipif(
    not HANDLER_AVAILABLE or not SDV_AVAILABLE or not PANDAS_AVAILABLE,
    reason=SKIP_MSG,
)
class TestCellValidationHandle:
    """Tests for CellValidationHandler.handle."""

    def test_uppercases_target_columns_and_uploads(self):
        from data_exchange_agent.tasks.data_validation.cell_validation import (
            CellValidationHandler,
        )

        parent = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={"redshift": Mock(connection_string="fake")},
        )
        parent._upload_to_stage = Mock()

        source_df = pd.DataFrame({"ID": [1, 2], "NAME": ["a", "b"]})
        target_df = pd.DataFrame({"id": [1, 2], "name": ["a", "b"]})

        parent._run_source_query = Mock(return_value=source_df)
        parent._run_target_query = Mock(return_value=target_df)

        mock_validation = Mock()
        mock_validation.is_valid = True
        mock_validation.results_dataframe = pd.DataFrame()

        cv = CellValidationHandler(parent)
        with patch(
            "data_exchange_agent.tasks.data_validation.cell_validation.validate_cell",
            return_value=mock_validation,
        ):
            result = cv.handle(
                {"payload": {"source_platform": "redshift"}},
                {
                    "table_name": "t",
                    "source_query": "SELECT 1",
                    "target_query": "SELECT 1",
                    "source_platform": "redshift",
                    "index_column_list": ["ID"],
                },
            )

        assert result["status"] == "success"
        assert result["total"] == 0
        parent._upload_to_stage.assert_called_once()

    def test_non_empty_diffs_are_vectorized_to_csv(self):
        from data_exchange_agent.tasks.data_validation.cell_validation import (
            CellValidationHandler,
        )

        parent = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={"redshift": Mock(connection_string="fake")},
        )
        source_df = pd.DataFrame({"ID": [1, 2], "NAME": ["a", "b"]})
        target_df = pd.DataFrame({"id": [1, 2], "name": ["a", "x"]})
        parent._run_source_query = Mock(return_value=source_df)
        parent._run_target_query = Mock(return_value=target_df)
        parent._upload_to_stage = Mock()
        parent._write_l3_csv = Mock(return_value="/tmp/diffs.csv")
        parent._cleanup_temp_dir = Mock()

        mock_validation = Mock()
        mock_validation.is_valid = False
        mock_validation.results_dataframe = pd.DataFrame(
            {
                "ROW": [1, 2],
                "COLUMN": ["col_a", "col_b"],
                "SOURCE_VALUE": ["src1", "src2"],
                "TARGET_VALUE": ["tgt1", "tgt2"],
            }
        )

        cv = CellValidationHandler(parent)
        with patch(
            "data_exchange_agent.tasks.data_validation.cell_validation.validate_cell",
            return_value=mock_validation,
        ):
            result = cv.handle(
                {"payload": {"source_platform": "redshift"}},
                {
                    "table_name": "t",
                    "source_query": "SELECT 1",
                    "target_query": "SELECT 1",
                    "source_platform": "redshift",
                    "index_column_list": ["ID"],
                },
            )

        assert result["status"] == "success"
        assert result["total"] == 2

        call_args = parent._write_l3_csv.call_args
        diffs_data = call_args[0][2]
        assert len(diffs_data) == 2
        assert diffs_data[0][0] == "t"
        assert diffs_data[0][2] == "col_a"


@pytest.mark.skipif(not HANDLER_AVAILABLE or not SDV_AVAILABLE, reason=SKIP_MSG)
class TestValidateSchemaSDV:
    """Tests for DataValidationHandler._validate_schema_sdv."""

    def test_delegates_to_sdv_and_adapts_result(self):
        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )
        source_df = pd.DataFrame({"ID": [1], "NAME": ["a"]})
        target_df = pd.DataFrame({"ID": [1], "NAME": ["a"]})

        mock_validation = Mock()
        mock_validation.is_valid = True
        mock_validation.results_dataframe = pd.DataFrame(
            {
                "VALIDATION_TYPE": ["schema"],
                "TABLE_NAME": ["t"],
                "COLUMN_VALIDATED": ["ID"],
                "EVALUATION_CRITERIA": ["DATA_TYPE"],
                "SOURCE_VALUE": ["int"],
                "SNOWFLAKE_VALUE": ["NUMBER"],
                "STATUS": ["SUCCESS"],
            }
        )
        table_ctx = Mock()

        with patch(
            "data_exchange_agent.tasks.data_validation.handler.validate_schema",
            return_value=mock_validation,
        ) as mock_vs:
            result = h._validate_schema_sdv(source_df, target_df, table_ctx, {}, {}, "t")
            mock_vs.assert_called_once()

        assert isinstance(result, pd.DataFrame)
        assert len(result) == 1


@pytest.mark.skipif(not HANDLER_AVAILABLE or not SDV_AVAILABLE, reason=SKIP_MSG)
class TestValidateMetricsSDV:
    """Tests for DataValidationHandler._validate_metrics_sdv."""

    def test_delegates_to_sdv_and_adapts_result(self):
        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )
        source_df = pd.DataFrame({"ID": [1, 2, 3]})
        target_df = pd.DataFrame({"ID": [1, 2, 3]})

        mock_validation = Mock()
        mock_validation.is_valid = True
        mock_validation.results_dataframe = pd.DataFrame(
            {
                "VALIDATION_TYPE": ["metrics"],
                "TABLE_NAME": ["t"],
                "COLUMN_VALIDATED": ["ID"],
                "EVALUATION_CRITERIA": ["SUM"],
                "SOURCE_VALUE": ["6"],
                "SNOWFLAKE_VALUE": ["6"],
                "STATUS": ["SUCCESS"],
            }
        )
        table_ctx = Mock()

        with patch(
            "data_exchange_agent.tasks.data_validation.handler.validate_metrics",
            return_value=mock_validation,
        ) as mock_vm:
            result = h._validate_metrics_sdv(source_df, target_df, table_ctx, {}, 0.001, "t")
            mock_vm.assert_called_once()

        assert isinstance(result, pd.DataFrame)
        assert len(result) == 1

    def test_uppercases_column_validated_when_present(self):
        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )
        source_df = pd.DataFrame({"COLUMN_VALIDATED": ["id", "name"], "VALUE": [6, 3]})
        target_df = pd.DataFrame({"COLUMN_VALIDATED": ["id", "name"], "VALUE": [6, 3]})

        mock_validation = Mock()
        mock_validation.is_valid = True
        mock_validation.results_dataframe = pd.DataFrame()
        table_ctx = Mock()

        with patch(
            "data_exchange_agent.tasks.data_validation.handler.validate_metrics",
            return_value=mock_validation,
        ):
            h._validate_metrics_sdv(source_df, target_df, table_ctx, {}, 0.001, "t")


@pytest.mark.skipif(not HANDLER_AVAILABLE or not PANDAS_AVAILABLE, reason=SKIP_MSG)
class TestRunTargetQuery:
    """Tests for DataValidationHandler._run_target_query."""

    def test_returns_empty_dataframe_for_empty_query(self):
        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )
        result = h._run_target_query("")
        assert isinstance(result, pd.DataFrame)
        assert result.empty

    def test_executes_query_and_returns_dataframe(self):
        mock_cursor = Mock()
        mock_cursor.description = [("COL_A",), ("COL_B",)]
        mock_cursor.fetchall.return_value = [(1, "x"), (2, "y")]

        mock_ds = Mock()
        mock_ds.get_cursor.return_value.__enter__ = Mock(return_value=mock_cursor)
        mock_ds.get_cursor.return_value.__exit__ = Mock(return_value=False)

        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=mock_ds,
            source_connections_config={},
        )
        result = h._run_target_query("SELECT 1")
        assert len(result) == 2
        assert list(result.columns) == ["COL_A", "COL_B"]


@pytest.mark.skipif(not HANDLER_AVAILABLE or not PANDAS_AVAILABLE, reason=SKIP_MSG)
class TestUploadToStage:
    """Tests for DataValidationHandler._upload_to_stage."""

    def test_executes_put_command(self):
        mock_cursor = Mock()
        mock_ds = Mock()
        mock_ds.get_cursor.return_value.__enter__ = Mock(return_value=mock_cursor)
        mock_ds.get_cursor.return_value.__exit__ = Mock(return_value=False)

        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=mock_ds,
            source_connections_config={},
        )
        h._upload_to_stage("/tmp/file.csv", "stage/path")
        mock_cursor.execute.assert_called_once()
        assert "PUT" in mock_cursor.execute.call_args[0][0]


@pytest.mark.skipif(not HANDLER_AVAILABLE or not PANDAS_AVAILABLE, reason=SKIP_MSG)
class TestRunSourceQuery:
    """Tests for DataValidationHandler._run_source_query."""

    def test_returns_empty_dataframe_for_empty_query(self):
        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )
        result = h._run_source_query("", {})
        assert isinstance(result, pd.DataFrame)
        assert result.empty

    def test_raises_for_unknown_engine(self):
        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={"redshift": Mock()},
        )
        task = {"payload": {"source_platform": "oracle"}}
        with pytest.raises(ValueError, match="Engine 'oracle' not found"):
            h._run_source_query("SELECT 1", task)

    @pytest.mark.skipif(not PYODBC_AVAILABLE, reason="pyodbc not loadable")
    @patch("pyodbc.connect")
    @patch("data_exchange_agent.tasks.data_validation.handler.register_output_converters")
    def test_executes_query_via_odbc(self, mock_register, mock_connect):
        mock_cursor = Mock()
        mock_cursor.description = [("COL_A",), ("COL_B",)]
        mock_cursor.fetchall.return_value = [("v1", "v2")]
        mock_conn = Mock()
        mock_conn.cursor.return_value = mock_cursor
        mock_connect.return_value = mock_conn

        mock_cfg = Mock()
        mock_cfg.connection_string = "DSN=test"

        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={"sqlserver": mock_cfg},
        )
        task = {"payload": {"source_platform": "sqlserver"}}
        result = h._run_source_query("SELECT 1", task)

        assert list(result.columns) == ["COL_A", "COL_B"]
        assert len(result) == 1
        mock_conn.close.assert_called_once()


@pytest.mark.skipif(not HANDLER_AVAILABLE or not PANDAS_AVAILABLE or not PYODBC_AVAILABLE, reason=SKIP_MSG)
class TestOpenSourceConnection:
    """Tests for DataValidationHandler._open_source_connection."""

    @patch("pyodbc.connect")
    @patch("data_exchange_agent.tasks.data_validation.handler.register_output_converters")
    def test_returns_connection(self, mock_register, mock_connect):
        mock_conn = Mock()
        mock_connect.return_value = mock_conn
        mock_cfg = Mock()
        mock_cfg.connection_string = "DSN=test"

        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={"sqlserver": mock_cfg},
        )
        task = {"payload": {"source_platform": "sqlserver"}}
        conn = h._open_source_connection(task)

        assert conn is mock_conn
        mock_register.assert_called_once_with(mock_conn)

    @patch("pyodbc.connect")
    @patch("data_exchange_agent.tasks.data_validation.handler.register_output_converters")
    def test_executes_use_database_command(self, mock_register, mock_connect):
        mock_cursor = Mock()
        mock_conn = Mock()
        mock_conn.cursor.return_value = mock_cursor
        mock_connect.return_value = mock_conn
        mock_cfg = Mock()
        mock_cfg.connection_string = "DSN=test"

        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={"sqlserver": mock_cfg},
        )
        task = {"payload": {"source_platform": "sqlserver", "use_database_command": "USE mydb"}}
        h._open_source_connection(task)

        mock_cursor.execute.assert_called_once_with("USE mydb")
        mock_cursor.close.assert_called_once()


@pytest.mark.skipif(not HANDLER_AVAILABLE or not PANDAS_AVAILABLE, reason=SKIP_MSG)
class TestExecuteTargetStatement:
    """Tests for DataValidationHandler._execute_target_statement."""

    def test_executes_statement_via_cursor(self):
        mock_cursor = Mock()
        mock_ds = Mock()
        mock_ds.get_cursor.return_value.__enter__ = Mock(return_value=mock_cursor)
        mock_ds.get_cursor.return_value.__exit__ = Mock(return_value=False)

        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=mock_ds,
            source_connections_config={},
        )
        h._execute_target_statement("CREATE TABLE t (id INT)")
        mock_cursor.execute.assert_called_once_with("CREATE TABLE t (id INT)")


@pytest.mark.skipif(not HANDLER_AVAILABLE or not PANDAS_AVAILABLE, reason=SKIP_MSG)
class TestGetTargetRowCount:
    """Tests for DataValidationHandler._get_target_row_count."""

    def test_returns_row_count(self):
        mock_cursor = Mock()
        mock_cursor.fetchone.return_value = (42,)
        mock_ds = Mock()
        mock_ds.get_cursor.return_value.__enter__ = Mock(return_value=mock_cursor)
        mock_ds.get_cursor.return_value.__exit__ = Mock(return_value=False)

        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=mock_ds,
            source_connections_config={},
        )
        assert h._get_target_row_count("db.schema.table") == 42

    def test_returns_zero_for_none_result(self):
        mock_cursor = Mock()
        mock_cursor.fetchone.return_value = None
        mock_ds = Mock()
        mock_ds.get_cursor.return_value.__enter__ = Mock(return_value=mock_cursor)
        mock_ds.get_cursor.return_value.__exit__ = Mock(return_value=False)

        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=mock_ds,
            source_connections_config={},
        )
        assert h._get_target_row_count("db.schema.table") == 0


@pytest.mark.skipif(not HANDLER_AVAILABLE or not PANDAS_AVAILABLE, reason=SKIP_MSG)
class TestExecuteSourceStatementOnConn:
    """Tests for DataValidationHandler._execute_source_statement_on_conn."""

    def test_executes_and_closes_cursor(self):
        mock_conn = Mock()
        mock_cursor = Mock()
        mock_conn.cursor.return_value = mock_cursor

        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )
        h._execute_source_statement_on_conn(mock_conn, "DROP TABLE tmp")
        mock_cursor.execute.assert_called_once_with("DROP TABLE tmp")
        mock_cursor.close.assert_called_once()


@pytest.mark.skipif(not HANDLER_AVAILABLE or not PANDAS_AVAILABLE, reason=SKIP_MSG)
class TestRunSourceQueryOnConn:
    """Tests for DataValidationHandler._run_source_query_on_conn."""

    def test_returns_dataframe_with_uppercased_columns(self):
        mock_conn = Mock()
        mock_cursor = Mock()
        mock_cursor.description = [("col_a",), ("col_b",)]
        mock_cursor.fetchall.return_value = [(1, "x"), (2, "y")]
        mock_conn.cursor.return_value = mock_cursor

        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )
        result = h._run_source_query_on_conn(mock_conn, "SELECT 1")
        assert list(result.columns) == ["COL_A", "COL_B"]
        assert len(result) == 2
        mock_cursor.close.assert_called_once()


@pytest.mark.skipif(not HANDLER_AVAILABLE or not PANDAS_AVAILABLE, reason=SKIP_MSG)
class TestGetRowCountOnConn:
    """Tests for DataValidationHandler._get_row_count_on_conn."""

    def test_returns_count(self):
        mock_conn = Mock()
        mock_cursor = Mock()
        mock_cursor.fetchone.return_value = (100,)
        mock_conn.cursor.return_value = mock_cursor

        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )
        assert h._get_row_count_on_conn(mock_conn, "db.t") == 100
        mock_cursor.close.assert_called_once()

    def test_returns_zero_when_no_row(self):
        mock_conn = Mock()
        mock_cursor = Mock()
        mock_cursor.fetchone.return_value = None
        mock_conn.cursor.return_value = mock_cursor

        h = DataValidationHandler(
            logger=Mock(),
            snowflake_datasource=Mock(),
            source_connections_config={},
        )
        assert h._get_row_count_on_conn(mock_conn, "db.t") == 0
