"""Tests for retrieval agent — plan parsing, broad query detection."""

from __future__ import annotations

import json
from unittest.mock import AsyncMock

from conftest import FakeLLMProvider, run_async

from arandu.read.retrieval_agent import (
    _parse_plan,
    plan_retrieval,
)

# ===========================================================================
# Test: plan parsing
# ===========================================================================


class TestParsePlan:
    def test_parse_similarity_plan(self) -> None:
        raw = json.dumps(
            {
                "strategy": "similarity",
                "pattern_queries": [],
                "similarity_query": "city where user lives",
                "max_facts": 20,
                "reason": "specific search",
            }
        )
        plan = _parse_plan(raw, "where do I live?")

        assert plan is not None
        assert plan.strategy == "multi_signal"  # V5: mapped from similarity
        assert plan.similarity_query == "city where user lives"
        assert plan.max_facts == 20

    def test_parse_skip_plan(self) -> None:
        raw = json.dumps(
            {
                "strategy": "skip",
                "pattern_queries": [],
                "similarity_query": None,
                "max_facts": 0,
                "reason": "greeting",
            }
        )
        plan = _parse_plan(raw, "hi")

        assert plan is not None
        assert plan.strategy == "skip"

    def test_parse_pattern_with_entities(self) -> None:
        raw = json.dumps(
            {
                "strategy": "pattern",
                "pattern_queries": [{"entity_pattern": "social:%"}],
                "similarity_query": None,
                "max_facts": 50,
                "reason": "aggregation",
                "entities": ["social:kevin"],
            }
        )
        plan = _parse_plan(raw, "who are my friends?")

        assert plan is not None
        assert plan.strategy == "multi_signal"
        assert "social:kevin" in plan.entities
        assert len(plan.pattern_queries) == 1
        assert plan.pattern_queries[0].entity_pattern == "social:%"

    def test_parse_time_travel(self) -> None:
        raw = json.dumps(
            {
                "strategy": "similarity",
                "similarity_query": "salary",
                "max_facts": 30,
                "reason": "temporal",
                "as_of_range": ["2025-03-01T00:00:00-03:00", "2025-03-31T23:59:59-03:00"],
            }
        )
        plan = _parse_plan(raw, "how much did I earn in March?")

        assert plan is not None
        assert plan.as_of_range is not None
        start, end = plan.as_of_range
        assert start.year == 2025
        assert start.month == 3
        assert end.month == 3

    def test_parse_invalid_json_returns_none(self) -> None:
        assert _parse_plan("not json at all", "query") is None

    def test_parse_empty_returns_none(self) -> None:
        assert _parse_plan("", "query") is None
        assert _parse_plan(None, "query") is None

    def test_parse_markdown_fenced_json(self) -> None:
        raw = '```json\n{"strategy":"skip","reason":"greeting"}\n```'
        plan = _parse_plan(raw, "hi")

        assert plan is not None
        assert plan.strategy == "skip"

    def test_parse_max_facts_clamped(self) -> None:
        raw = json.dumps({"strategy": "similarity", "max_facts": 9999})
        plan = _parse_plan(raw, "query")
        assert plan is not None
        assert plan.max_facts == 200  # clamped to max

    def test_parse_entities_from_pattern(self) -> None:
        """Entities are derived from pattern_queries when no explicit 'entities' field."""
        raw = json.dumps(
            {
                "strategy": "pattern",
                "pattern_queries": [{"entity_pattern": "person:kevin:%"}],
            }
        )
        plan = _parse_plan(raw, "query")
        assert plan is not None
        assert "person:kevin" in plan.entities


# ===========================================================================
# Test: plan_retrieval via LLM
# ===========================================================================


def _mock_session_with_schema(pairs: list[tuple[str, str]]) -> AsyncMock:
    """Create a mock AsyncSession that returns schema pairs from execute()."""
    from unittest.mock import MagicMock

    mock_result = MagicMock()
    mock_result.all.return_value = pairs

    session = AsyncMock()
    session.execute = AsyncMock(return_value=mock_result)
    return session


class TestPlanRetrieval:
    def test_empty_query_returns_multi_signal(self) -> None:
        session = AsyncMock()
        llm = FakeLLMProvider()

        plan = run_async(plan_retrieval(session, "user_123", "", llm))

        assert plan.strategy == "multi_signal"
        assert plan.reason == "empty_query"

    def test_greeting_handled_by_llm(self) -> None:
        """Greetings are now handled by the LLM (no regex shortcut)."""
        session = _mock_session_with_schema([])

        llm_response = json.dumps(
            {
                "strategy": "skip",
                "similarity_query": None,
                "max_facts": 0,
                "reason": "greeting",
            }
        )
        llm = FakeLLMProvider([llm_response])

        from arandu.read.retrieval_agent import _schema_cache

        _schema_cache.clear()

        plan = run_async(plan_retrieval(session, "user_123", "oi", llm))

        assert plan.strategy == "skip"

    def test_llm_called_for_real_query(self) -> None:
        session = _mock_session_with_schema([("person:kevin", "preference")])

        llm_response = json.dumps(
            {
                "strategy": "similarity",
                "similarity_query": "Kevin preferences",
                "max_facts": 20,
                "reason": "specific search",
            }
        )
        llm = FakeLLMProvider([llm_response])

        from arandu.read.retrieval_agent import _schema_cache

        _schema_cache.clear()

        plan = run_async(plan_retrieval(session, "user_123", "what does Kevin like?", llm))

        assert plan.strategy == "multi_signal"
        assert plan.similarity_query == "Kevin preferences"

    def test_llm_failure_returns_fallback(self) -> None:
        session = _mock_session_with_schema([("person:kevin", "preference")])

        llm = FakeLLMProvider(["this is not json"])

        from arandu.read.retrieval_agent import _schema_cache

        _schema_cache.clear()

        plan = run_async(plan_retrieval(session, "user_123", "where do I live?", llm))

        assert plan.strategy == "multi_signal"
        assert plan.reason == "agent_fallback"

    def test_broad_query_detected_from_reason(self) -> None:
        """Broad query detection now uses LLM reason field, not regex."""
        session = _mock_session_with_schema([])

        llm_response = json.dumps(
            {
                "strategy": "similarity",
                "similarity_query": "everything about user",
                "max_facts": 100,
                "reason": "broad overview",
            }
        )
        llm = FakeLLMProvider([llm_response])

        from arandu.read.retrieval_agent import _schema_cache

        _schema_cache.clear()

        plan = run_async(plan_retrieval(session, "user_123", "tell me everything about me", llm))

        assert plan.broad_query is True
