"""
tests/test_agents/test_sql_agent.py
=====================================
Unit tests for agents/sql_agent.py using stdlib unittest only.

Run standalone:
    python -m unittest tests/test_agents/test_sql_agent.py -v
Run via runner:
    python tests/run_tests.py

Coverage:
- sanitize_sql()       — read-only enforcement, markdown stripping, forbidden keywords
- reduce_sql_result()  — deduplication, row trimming, string truncation
- _extract_json_object() — JSON parsing from LLM text with markdown fences
- process_sql()        — VSSState bridge: success, no-results, error, trace merging
"""

from __future__ import annotations

import sys
import logging
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch

# ── bootstrap ────────────────────────────────────────────────────────────────
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))  # → tests/
import shared  # noqa: E402

logger = logging.getLogger("vss.tests.sql_agent")


# ─────────────────────────────────────────────────────────────────────────────
# 1. sanitize_sql()
# ─────────────────────────────────────────────────────────────────────────────

class TestSanitizeSql(unittest.TestCase):
    """Tests for sanitize_sql() — safety enforcement and cleanup."""

    def setUp(self):
        from agents.sql_agent import sanitize_sql
        self.fn = sanitize_sql

    # ── allowed ──────────────────────────────────────────────────────────────

    def test_plain_select_passes(self):
        result = self.fn("SELECT count() FROM default.tbl")
        self.assertTrue(result.lower().startswith("select"))
        logger.debug("plain SELECT passed: %r", result)

    def test_with_clause_passes(self):
        result = self.fn("WITH cte AS (SELECT 1) SELECT * FROM cte")
        self.assertTrue(result.lower().startswith("with"))

    def test_leading_whitespace_stripped(self):
        result = self.fn("   \n  SELECT 1")
        self.assertIn("1", result)

    # ── markdown cleanup ─────────────────────────────────────────────────────

    def test_sql_fence_stripped(self):
        result = self.fn("```sql\nSELECT 1\n```")
        self.assertNotIn("```", result)
        self.assertTrue(result.lower().startswith("select"))

    def test_generic_fence_stripped(self):
        result = self.fn("```\nSELECT 1\n```")
        self.assertNotIn("```", result)

    def test_backtick_identifiers_removed(self):
        result = self.fn("SELECT `col` FROM `tbl`")
        self.assertNotIn("`", result)

    # ── forbidden keywords ───────────────────────────────────────────────────

    def test_insert_raises(self):
        with self.assertRaises(ValueError):
            self.fn("INSERT INTO tbl VALUES (1)")

    def test_update_raises(self):
        with self.assertRaises(ValueError):
            self.fn("UPDATE tbl SET col=1")

    def test_delete_raises(self):
        with self.assertRaises(ValueError):
            self.fn("DELETE FROM tbl")

    def test_drop_raises(self):
        with self.assertRaises(ValueError):
            self.fn("DROP TABLE tbl")

    def test_create_raises(self):
        with self.assertRaises(ValueError):
            self.fn("CREATE TABLE x (id INT)")

    def test_alter_raises(self):
        with self.assertRaises(ValueError):
            self.fn("ALTER TABLE tbl ADD COLUMN x INT")

    def test_truncate_raises(self):
        with self.assertRaises(ValueError):
            self.fn("TRUNCATE TABLE tbl")

    def test_grant_raises(self):
        with self.assertRaises(ValueError):
            self.fn("GRANT ALL ON tbl TO user")

    def test_non_select_raises(self):
        with self.assertRaises(ValueError):
            self.fn("SHOW TABLES")


# ─────────────────────────────────────────────────────────────────────────────
# 2. reduce_sql_result()
# ─────────────────────────────────────────────────────────────────────────────

class TestReduceSqlResult(unittest.TestCase):
    """Tests for reduce_sql_result() — deduplication, trimming, truncation."""

    def setUp(self):
        from agents.sql_agent import reduce_sql_result
        self.fn = reduce_sql_result

    def test_none_returns_none(self):
        self.assertIsNone(self.fn(None))

    def test_string_truncated_at_4000(self):
        result = self.fn("x" * 5000)
        self.assertEqual(len(result), 4000)

    def test_short_string_unchanged(self):
        self.assertEqual(self.fn("hello"), "hello")

    def test_small_list_unchanged(self):
        rows = [{"a": i} for i in range(5)]
        self.assertEqual(self.fn(rows, max_rows=20), rows)

    def test_large_list_trimmed_to_max_rows(self):
        rows = [{"val": i} for i in range(100)]
        result = self.fn(rows, max_rows=20)
        self.assertEqual(len(result), 20)

    def test_deduplication_removes_identical_rows(self):
        rows = [{"val": 1}] * 10
        result = self.fn(rows, max_rows=20, deduplicate=True)
        self.assertEqual(len(result), 1)

    def test_deduplication_disabled_keeps_duplicates(self):
        rows = [{"val": 1}] * 5
        result = self.fn(rows, max_rows=20, deduplicate=False)
        self.assertEqual(len(result), 5)

    def test_head_and_tail_sampling(self):
        rows = [{"val": i} for i in range(50)]
        result = self.fn(rows, max_rows=10)
        self.assertEqual(len(result), 10)
        self.assertEqual(result[0], {"val": 0})
        self.assertEqual(result[-1], {"val": 49})


# ─────────────────────────────────────────────────────────────────────────────
# 3. _extract_json_object()
# ─────────────────────────────────────────────────────────────────────────────

class TestExtractJsonObject(unittest.TestCase):
    """Tests for _extract_json_object() — JSON parsing with markdown fences."""

    def setUp(self):
        from agents.sql_agent import _extract_json_object
        self.fn = _extract_json_object

    def test_plain_json(self):
        result = self.fn('{"key": "value"}')
        self.assertEqual(result, {"key": "value"})

    def test_json_with_markdown_fence(self):
        result = self.fn('```json\n{"key": "value"}\n```')
        self.assertEqual(result, {"key": "value"})

    def test_json_embedded_in_text(self):
        text = 'Here is the answer: {"use_sql_agent": true, "reason": "test"}'
        result = self.fn(text)
        self.assertTrue(result["use_sql_agent"])

    def test_nested_json(self):
        text = '{"selected_tables": ["tbl1", "tbl2"], "reason": "matched"}'
        result = self.fn(text)
        self.assertEqual(result["selected_tables"], ["tbl1", "tbl2"])

    def test_no_json_raises_value_error(self):
        with self.assertRaises(ValueError):
            self.fn("no json here at all")


# ─────────────────────────────────────────────────────────────────────────────
# 4. process_sql() — VSSState bridge
# ─────────────────────────────────────────────────────────────────────────────

class TestProcessSql(unittest.IsolatedAsyncioTestCase):
    """
    Tests for process_sql() — the async bridge from VSSState to the
    internal LangGraph SQLAgent pipeline and back.
    External _get_agent() is always mocked so no DB/LLM is needed.
    """

    async def test_success_sets_sql_query(self):
        from agents.sql_agent import process_sql

        state = shared.make_state()
        agent_result = shared.make_sql_result()

        with patch("agents.sql_agent._get_agent") as mock_get:
            mock_get.return_value.run.return_value = agent_result
            result = await process_sql(state)

        self.assertEqual(result["sql_query"], "SELECT count() FROM default.aggregated_analytics_final")
        self.assertIsNone(result["sql_error"])
        logger.info("process_sql success → sql_query set")

    async def test_success_sets_sql_result_formatted(self):
        from agents.sql_agent import process_sql

        state = shared.make_state()
        agent_result = shared.make_sql_result(final_answer="42 events found.")

        with patch("agents.sql_agent._get_agent") as mock_get:
            mock_get.return_value.run.return_value = agent_result
            result = await process_sql(state)

        self.assertEqual(result["sql_result_formatted"], "42 events found.")
        self.assertFalse(result["escalated"])

    async def test_no_execution_result_sets_escalated(self):
        from agents.sql_agent import process_sql

        state = shared.make_state()
        agent_result = shared.make_sql_result(execution_result="", execution_error="")

        with patch("agents.sql_agent._get_agent") as mock_get:
            mock_get.return_value.run.return_value = agent_result
            result = await process_sql(state)

        self.assertTrue(result["escalated"])
        self.assertEqual(result["escalation_reason"], "SQL returned no results")
        logger.info("process_sql no results → escalated=True")

    async def test_execution_error_sets_sql_error(self):
        from agents.sql_agent import process_sql

        state = shared.make_state()
        agent_result = shared.make_sql_result(
            execution_error="DB connection refused",
            execution_result="",
        )

        with patch("agents.sql_agent._get_agent") as mock_get:
            mock_get.return_value.run.return_value = agent_result
            result = await process_sql(state)

        self.assertEqual(result["sql_error"], "DB connection refused")
        self.assertFalse(result["escalated"])
        logger.warning("process_sql DB error → sql_error correctly set")

    async def test_agent_exception_sets_sql_error(self):
        from agents.sql_agent import process_sql

        state = shared.make_state()

        with patch("agents.sql_agent._get_agent") as mock_get:
            mock_get.side_effect = RuntimeError("Agent init failed")
            result = await process_sql(state)

        self.assertIsNotNone(result["sql_error"])
        self.assertIn("Agent init failed", result["sql_error"])
        logger.error("process_sql exception → sql_error=%s", result["sql_error"])

    async def test_internal_trace_merged_into_state(self):
        from agents.sql_agent import process_sql

        state = shared.make_state()
        initial_trace_len = len(state["trace"])
        agent_result = shared.make_sql_result(
            extra_trace=["[00:01] SQL generated", "[00:02] SQL executed"]
        )

        with patch("agents.sql_agent._get_agent") as mock_get:
            mock_get.return_value.run.return_value = agent_result
            result = await process_sql(state)

        self.assertGreater(len(result["trace"]), initial_trace_len + 2)
        combined = " ".join(result["trace"])
        self.assertIn("SQL generated", combined)
        shared.log_state_trace(result, "test_internal_trace_merged_into_state")
        logger.debug("Trace after merge: %d entries", len(result["trace"]))

    async def test_process_sql_adds_own_trace_entries(self):
        """process_sql itself should add at least a 'SQLAgent:' trace entry."""
        from agents.sql_agent import process_sql

        state = shared.make_state()
        agent_result = shared.make_sql_result()

        with patch("agents.sql_agent._get_agent") as mock_get:
            mock_get.return_value.run.return_value = agent_result
            result = await process_sql(state)

        sqlagent_traces = [t for t in result["trace"] if "SQLAgent" in t]
        self.assertGreaterEqual(len(sqlagent_traces), 1)

    async def test_dict_rows_parsed_from_execution_result(self):
        """When execution_result contains dict rows, sql_result_formatted must be set."""
        from agents.sql_agent import process_sql

        state = shared.make_state()
        agent_result = shared.make_sql_result(
            execution_result="[{'count': 42, 'name': 'Fire'}]",
            final_answer="42 fire events.",
        )

        with patch("agents.sql_agent._get_agent") as mock_get:
            mock_get.return_value.run.return_value = agent_result
            result = await process_sql(state)

        self.assertEqual(result["sql_result_formatted"], "42 fire events.")

    async def test_empty_execution_result_string_variants(self):
        """Various 'empty' strings must all trigger escalation."""
        from agents.sql_agent import process_sql

        for empty_val in ("", "[]", "None", "none"):
            with self.subTest(execution_result=empty_val):
                state = shared.make_state()
                agent_result = shared.make_sql_result(
                    execution_result=empty_val,
                    execution_error="",
                )
                with patch("agents.sql_agent._get_agent") as mock_get:
                    mock_get.return_value.run.return_value = agent_result
                    result = await process_sql(state)

                self.assertTrue(result["escalated"],
                    f"execution_result={empty_val!r} should trigger escalation")


if __name__ == "__main__":
    unittest.main(verbosity=2)
