"""
Unit Tests for GreatExpectationsLakehouseOperator
Verifies declarative rules (not_null, unique, in_set, row_count),
custom SQL assertions, error handling, and XCom payload formatting.
"""
import os
import sys
import pytest
from unittest.mock import MagicMock, patch

# Ensure airflow/dags directory is on python path
DAGS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "dags"))
if DAGS_DIR not in sys.path:
    sys.path.insert(0, DAGS_DIR)

from helpers.ge_lakehouse_operator import GreatExpectationsLakehouseOperator

try:
    from airflow.exceptions import AirflowException
except ImportError:
    class AirflowException(Exception):
        pass


class TestGreatExpectationsLakehouseOperator:

    @patch("helpers.ge_lakehouse_operator.create_engine")
    def test_declarative_rules_all_pass(self, mock_create_engine):
        mock_conn = MagicMock()
        mock_engine = MagicMock()
        mock_engine.connect.return_value.__enter__.return_value = mock_conn
        mock_create_engine.return_value = mock_engine

        # Mock query outputs:
        # 1. row count = 10 (valid for min 1)
        # 2. null count = 0 (valid for not_null)
        # 3. in_set invalid count = 0 (valid for in_set)
        mock_conn.execute.return_value.scalar.side_effect = [10, 0, 0]

        op = GreatExpectationsLakehouseOperator(
            task_id="test_ge_pass",
            expectation_suite="permits_test_suite",
            table_name="bronze_lakehouse.permits",
            rules=[
                {"expect": "table_row_count_to_be_between", "min_value": 1},
                {"expect": "column_values_to_not_be_null", "column": "permit_id"},
                {"expect": "column_values_to_be_in_set", "column": "status", "allowed_values": ["APPROVED", "REJECTED"]}
            ],
            database_url="sqlite:///:memory:"
        )

        res = op.execute(context={})
        assert res["status"] == "PASSED"
        assert res["overall_success"] is True
        assert res["total_evaluations"] == 3
        assert res["successful_evaluations"] == 3
        assert res["success_percent"] == 100.0

    @patch("helpers.ge_lakehouse_operator.create_engine")
    def test_custom_sql_query_pass(self, mock_create_engine):
        mock_conn = MagicMock()
        mock_engine = MagicMock()
        mock_engine.connect.return_value.__enter__.return_value = mock_conn
        mock_create_engine.return_value = mock_engine

        # Query returns 0 (expected 0)
        mock_conn.execute.return_value.scalar.return_value = 0

        op = GreatExpectationsLakehouseOperator(
            task_id="test_custom_sql",
            expectation_suite="permits_sql_suite",
            queries=[
                {
                    "name": "check_no_negative_amounts",
                    "sql": "SELECT count(*) FROM table WHERE amount < 0",
                    "expected": 0,
                    "error_message": "Found negative amounts!"
                }
            ],
            database_url="sqlite:///:memory:"
        )

        res = op.execute(context={})
        assert res["status"] == "PASSED"
        assert res["overall_success"] is True
        assert len(res["details"]) == 1
        assert res["details"][0]["observed_value"] == 0

    @patch("helpers.ge_lakehouse_operator.create_engine")
    def test_fail_fast_raises_airflow_exception(self, mock_create_engine):
        mock_conn = MagicMock()
        mock_engine = MagicMock()
        mock_engine.connect.return_value.__enter__.return_value = mock_conn
        mock_create_engine.return_value = mock_engine

        # Mock: 3 null values found in permit_id
        mock_conn.execute.return_value.scalar.return_value = 3

        op = GreatExpectationsLakehouseOperator(
            task_id="test_fail_fast",
            expectation_suite="permits_fail_suite",
            table_name="bronze_lakehouse.permits",
            rules=[
                {"expect": "column_values_to_not_be_null", "column": "permit_id"}
            ],
            fail_on_error=True,
            database_url="sqlite:///:memory:"
        )

        with pytest.raises(Exception) as excinfo:
            op.execute(context={})
        assert "Great Expectations Quality Gate GAGAL" in str(excinfo.value)

    @patch("helpers.ge_lakehouse_operator.create_engine")
    def test_fail_on_error_false_does_not_raise(self, mock_create_engine):
        mock_conn = MagicMock()
        mock_engine = MagicMock()
        mock_engine.connect.return_value.__enter__.return_value = mock_conn
        mock_create_engine.return_value = mock_engine

        # Mock: 2 invalid records
        mock_conn.execute.return_value.scalar.return_value = 2

        op = GreatExpectationsLakehouseOperator(
            task_id="test_warn_only",
            expectation_suite="permits_warn_suite",
            table_name="silver.permits",
            rules=[
                {"expect": "column_values_to_be_in_set", "column": "status", "allowed_values": ["APPROVED"]}
            ],
            fail_on_error=False,
            database_url="sqlite:///:memory:"
        )

        res = op.execute(context={})
        assert res["status"] == "FAILED"
        assert res["overall_success"] is False
        assert res["unsuccessful_evaluations"] == 1
        assert res["total_unexpected_records"] == 2