"""tests/test_verification_sql_validator.py — Unit tests for SQLSemanticValidator and SQLExecutionOracle."""
import pytest
import sqlite3

pytestmark = pytest.mark.unit

SCHEMA = "CREATE TABLE orders (id INT, amount FLOAT, status TEXT);"
SIMPLE_SQL = "SELECT id, amount FROM orders WHERE status = 'completed'"


class TestImport:
    def test_import_validator(self):
        from llm_guard.verification.sql_validator import SQLSemanticValidator
        assert SQLSemanticValidator is not None

    def test_import_oracle(self):
        from llm_guard.verification.sql_validator import SQLExecutionOracle
        assert SQLExecutionOracle is not None

    def test_import_result(self):
        from llm_guard.verification.sql_validator import SQLValidationResult
        assert SQLValidationResult is not None

    def test_import_oracle_result(self):
        from llm_guard.verification.sql_validator import SQLOracleResult
        assert SQLOracleResult is not None


class TestSQLSemanticValidatorInstantiation:
    def test_no_key_construction(self):
        from llm_guard.verification.sql_validator import SQLSemanticValidator
        v = SQLSemanticValidator(api_key="", schema=SCHEMA)
        assert v._schema == SCHEMA

    def test_dialect_stored(self):
        from llm_guard.verification.sql_validator import SQLSemanticValidator
        v = SQLSemanticValidator(api_key="", dialect="postgres")
        assert v._dialect == "postgres"


class TestSQLHeuristicValidation:
    """Tests for the zero-cost heuristic fallback (no API key → _heuristic_sql_validate)."""

    def _validate(self, question, sql, schema=SCHEMA, dialect="ansi"):
        from llm_guard.verification.sql_validator import SQLSemanticValidator
        v = SQLSemanticValidator(api_key="", schema=schema, dialect=dialect)
        return v.validate(question, sql)

    def test_valid_simple_select(self):
        result = self._validate("Get completed orders", SIMPLE_SQL)
        assert result.verdict in ("VALID", "UNCERTAIN")
        assert 0.0 <= result.risk_score <= 1.0

    def test_missing_group_by_detected(self):
        sql = "SELECT status, SUM(amount) FROM orders"
        result = self._validate("Total by status", sql)
        # Should detect missing GROUP BY
        assert result.verdict in ("UNCERTAIN", "INVALID")

    def test_join_without_on_flagged(self):
        sql = "SELECT a.id, b.amount FROM orders a JOIN orders b"
        result = self._validate("Join orders", sql, schema="CREATE TABLE orders (id INT, amount FLOAT);")
        assert len(result.issues) > 0 or result.verdict in ("UNCERTAIN", "INVALID")

    def test_no_schema_warns(self):
        import warnings
        from llm_guard.verification.sql_validator import SQLSemanticValidator
        v = SQLSemanticValidator(api_key="", schema="")
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            result = v.validate("query", SIMPLE_SQL)
            assert any("schema" in str(x.message).lower() for x in w)

    def test_dialect_databricks_not_strict_groupby(self):
        sql = "SELECT status, SUM(amount) FROM orders"
        result = self._validate("Total by status", sql, dialect="databricks")
        # Databricks allows loose GROUP BY
        assert result.verdict in ("VALID", "UNCERTAIN")

    def test_result_has_judge_model_heuristic(self):
        result = self._validate("query", SIMPLE_SQL)
        assert result.judge_model == "heuristic"

    def test_validate_batch(self):
        from llm_guard.verification.sql_validator import SQLSemanticValidator
        v = SQLSemanticValidator(api_key="", schema=SCHEMA)
        pairs = [
            {"question": "Q1", "sql": SIMPLE_SQL},
            {"question": "Q2", "sql": "SELECT COUNT(*) FROM orders"},
        ]
        results = v.validate_batch(pairs)
        assert len(results) == 2


class TestSQLExecutionOracle:
    def _sqlite_factory(self, tmp_path):
        db = str(tmp_path / "test.db")
        conn = sqlite3.connect(db)
        conn.execute("CREATE TABLE orders (id INT, amount FLOAT, status TEXT)")
        conn.execute("INSERT INTO orders VALUES (1, 100.0, 'completed')")
        conn.commit()
        conn.close()
        return lambda: sqlite3.connect(db)

    def test_successful_query(self, tmp_path):
        from llm_guard.verification.sql_validator import SQLExecutionOracle
        oracle = SQLExecutionOracle(self._sqlite_factory(tmp_path))
        result = oracle.check("SELECT id, amount FROM orders")
        assert result.passed is True
        assert result.shape == (1, 2)
        assert result.execution_risk == 0.0

    def test_error_query(self, tmp_path):
        from llm_guard.verification.sql_validator import SQLExecutionOracle
        oracle = SQLExecutionOracle(self._sqlite_factory(tmp_path))
        result = oracle.check("SELECT nonexistent_column FROM orders")
        assert result.passed is False
        assert result.execution_risk == 0.90

    def test_empty_result(self, tmp_path):
        from llm_guard.verification.sql_validator import SQLExecutionOracle
        oracle = SQLExecutionOracle(self._sqlite_factory(tmp_path))
        result = oracle.check("SELECT * FROM orders WHERE status = 'missing'")
        assert result.passed is True
        assert result.execution_risk == 0.50

    def test_blend_with_behavioral(self, tmp_path):
        from llm_guard.verification.sql_validator import SQLExecutionOracle
        oracle = SQLExecutionOracle(self._sqlite_factory(tmp_path))
        oracle_result = oracle.check("SELECT * FROM orders")
        blended = oracle.blend_with_behavioral(oracle_result, behavioral_risk=0.5)
        assert 0.0 <= blended <= 1.0

    def test_check_raw_with_list_result(self, tmp_path):
        from llm_guard.verification.sql_validator import SQLExecutionOracle
        oracle = SQLExecutionOracle(lambda: sqlite3.connect(":memory:"))
        result = oracle.check_raw("SELECT 1", result=[(1, 2), (3, 4)])
        assert result.passed is True
        assert result.execution_risk == 0.0

    def test_check_raw_with_error(self):
        from llm_guard.verification.sql_validator import SQLExecutionOracle
        oracle = SQLExecutionOracle(lambda: sqlite3.connect(":memory:"))
        result = oracle.check_raw("SELECT 1", result=None, error="syntax error")
        assert result.passed is False
        assert result.execution_risk == 0.90
