"""
Tests for cortex/hippo/tools/expand_query.py — hippo_expand_query tool.

Issues: #1119, #1146
"""

from __future__ import annotations

import json

import pytest
import pytest_asyncio

from cortex.hippo.config import HippoConfig
from cortex.hippo.engine import HippoEngine
from cortex.hippo.retrieval import RetrievalEngine
from cortex.hippo.tools.expand_query import (
    DEFAULT_MAX_ANSWER_TOKENS,
    GatewayDelegator,
    SynthesizeFn,
    _parse_delegated_reply,
    execute_expand_query,
)
from cortex.hippo.tools.recursion_guard import (
    reset_for_tests as reset_recursion_guard,
    stamp_delegated_expansion_context,
)


# ═══════════════════════════════════════════════════════════════════════════
# Fixtures
# ═══════════════════════════════════════════════════════════════════════════


@pytest.fixture
def hippo_config() -> HippoConfig:
    return HippoConfig(enabled=True, db_path=":memory:", fresh_tail_count=4, max_tokens=10_000)


@pytest_asyncio.fixture
async def engine(hippo_config):
    eng = HippoEngine(config=hippo_config)
    await eng.initialize(db_path=":memory:", owner_id="test-owner", conversation_id="test-session")
    yield eng
    await eng.shutdown()


@pytest_asyncio.fixture
async def retrieval(engine):
    return RetrievalEngine(
        conversation_store=engine.conversation_store,
        summary_store=engine.summary_store,
    )


@pytest_asyncio.fixture
async def seeded(engine, retrieval):
    """Engine + retrieval with a 2-level summary DAG for expand-query tests."""
    summary_store = engine.summary_store
    conv_id = engine._conversation.conversation_id

    # Ingest messages for leaf source
    for role, content in [
        ("user", "Let's deploy the app"),
        ("assistant", "Setting up fly.io deploy pipeline"),
        ("user", "What about the database migration?"),
        ("assistant", "Adding owner_id column migration"),
    ]:
        await engine.ingest_message(role=role, content=content)

    # Leaf summaries
    s1 = await summary_store.create_summary(
        conversation_id=conv_id,
        owner_id="test-owner",
        kind="leaf",
        depth=0,
        content="Deployed application to production via fly.io.",
        token_count=12,
    )
    s2 = await summary_store.create_summary(
        conversation_id=conv_id,
        owner_id="test-owner",
        kind="leaf",
        depth=0,
        content="Database schema migration adds owner_id column.",
        token_count=10,
    )

    # Condensed parent
    s3 = await summary_store.create_summary(
        conversation_id=conv_id,
        owner_id="test-owner",
        kind="condensed",
        depth=1,
        content="Deploy and migration work: fly.io deploy + owner_id column.",
        token_count=15,
        descendant_count=2,
        descendant_token_count=22,
    )

    # Link children
    await summary_store.add_summary_parent(s1.summary_id, s3.summary_id)
    await summary_store.add_summary_parent(s2.summary_id, s3.summary_id)

    yield {
        "engine": engine,
        "retrieval": retrieval,
        "conv_id": conv_id,
        "s1": s1,
        "s2": s2,
        "s3": s3,
        "summary_store": summary_store,
        "conversation_store": engine.conversation_store,
    }


@pytest.fixture(autouse=True)
def _reset_guard():
    """Reset recursion guard state between tests."""
    reset_recursion_guard()
    yield
    reset_recursion_guard()


# ═══════════════════════════════════════════════════════════════════════════
# Helpers
# ═══════════════════════════════════════════════════════════════════════════


class MockSynthesizeFn:
    """Mock synthesize callback that returns a predictable answer."""

    def __init__(self, answer: str = "Synthesized answer"):
        self.answer = answer
        self.calls: list[tuple[str, str, int]] = []

    async def __call__(self, context: str, prompt: str, max_tokens: int) -> str:
        self.calls.append((context, prompt, max_tokens))
        return self.answer


class MockGatewayDelegator:
    """Mock gateway delegator that returns a JSON reply."""

    def __init__(self, reply: dict | str | None = None, should_fail: bool = False):
        self._reply = reply
        self._should_fail = should_fail
        self.calls: list[tuple[str, str, float]] = []

    async def delegate(self, system_prompt: str, user_message: str, timeout: float) -> str:
        self.calls.append((system_prompt, user_message, timeout))
        if self._should_fail:
            raise ConnectionError("Gateway unavailable")
        if isinstance(self._reply, dict):
            return json.dumps(self._reply)
        return self._reply or ""


def _extract_details(result: dict) -> dict:
    """Extract the details dict from a tool result."""
    return result.get("details", {})


def _extract_text(result: dict) -> str:
    """Extract text content from a tool result."""
    content = result.get("content", [])
    if content and isinstance(content[0], dict):
        return content[0].get("text", "")
    return ""


# ═══════════════════════════════════════════════════════════════════════════
# Tests
# ═══════════════════════════════════════════════════════════════════════════


@pytest.mark.asyncio
async def test_expand_query_by_query(seeded):
    """grep → expand → synthesize pipeline."""
    synth = MockSynthesizeFn("Deploy was done via fly.io with owner_id migration.")
    result = await execute_expand_query(
        params={
            "prompt": "How was the app deployed?",
            "query": "deploy",
            "conversationId": seeded["conv_id"],
        },
        retrieval=seeded["retrieval"],
        owner_id="test-owner",
        conversation_store=seeded["conversation_store"],
        synthesize_fn=synth,
    )

    details = _extract_details(result)
    assert "error" not in details, f"Unexpected error: {details}"
    assert details.get("answer"), "Expected a non-empty answer"
    assert isinstance(details.get("citedIds"), list)
    assert details.get("expandedSummaryCount", 0) > 0
    assert details.get("sourceConversationId") == seeded["conv_id"]

    # Verify synthesize_fn was called
    assert len(synth.calls) == 1
    ctx, prompt, _ = synth.calls[0]
    assert "deploy" in ctx.lower() or "fly.io" in ctx.lower()
    assert prompt == "How was the app deployed?"


@pytest.mark.asyncio
async def test_expand_query_by_summary_ids(seeded):
    """Direct expand → synthesize with explicit summaryIds."""
    synth = MockSynthesizeFn("Summary s3 covers deploy and migration.")
    result = await execute_expand_query(
        params={
            "prompt": "What work was done?",
            "summaryIds": [seeded["s3"].summary_id],
            "conversationId": seeded["conv_id"],
        },
        retrieval=seeded["retrieval"],
        owner_id="test-owner",
        conversation_store=seeded["conversation_store"],
        synthesize_fn=synth,
    )

    details = _extract_details(result)
    assert "error" not in details, f"Unexpected error: {details}"
    assert details.get("answer")
    assert seeded["s3"].summary_id in details.get("citedIds", [])
    assert len(synth.calls) == 1


@pytest.mark.asyncio
async def test_expand_query_prompt_required(seeded):
    """Error if no prompt provided."""
    result = await execute_expand_query(
        params={"query": "deploy"},
        retrieval=seeded["retrieval"],
        owner_id="test-owner",
    )
    details = _extract_details(result)
    assert "error" in details
    assert "prompt" in details["error"].lower()


@pytest.mark.asyncio
async def test_expand_query_needs_query_or_ids(seeded):
    """Error if neither query nor summaryIds provided."""
    result = await execute_expand_query(
        params={"prompt": "Tell me about deploy"},
        retrieval=seeded["retrieval"],
        owner_id="test-owner",
    )
    details = _extract_details(result)
    assert "error" in details
    assert "summaryIds" in details["error"] or "query" in details["error"]


@pytest.mark.asyncio
async def test_expand_query_recursion_blocked(seeded):
    """Guard blocks at depth cap."""
    # Stamp a delegated context at depth=1 (the cap)
    stamp_delegated_expansion_context(
        session_key="test-session-deep",
        request_id="req-deep",
        expansion_depth=1,
        origin_session_key="origin-main",
        stamped_by="test",
    )

    result = await execute_expand_query(
        params={
            "prompt": "What happened?",
            "query": "deploy",
            "conversationId": seeded["conv_id"],
        },
        retrieval=seeded["retrieval"],
        session_key="test-session-deep",
        owner_id="test-owner",
    )

    details = _extract_details(result)
    assert "errorCode" in details or "error" in details
    assert details.get("errorCode") == "EXPANSION_RECURSION_BLOCKED"
    assert "blocked" in details.get("error", "").lower() or "recursion" in details.get("error", "").lower()


@pytest.mark.asyncio
async def test_expand_query_gateway_delegation(seeded):
    """Uses delegator when available."""
    delegator = MockGatewayDelegator(reply={
        "answer": "The app was deployed via fly.io",
        "citedIds": [seeded["s1"].summary_id],
        "expandedSummaryCount": 1,
        "totalSourceTokens": 50,
        "truncated": False,
    })

    result = await execute_expand_query(
        params={
            "prompt": "How was it deployed?",
            "summaryIds": [seeded["s1"].summary_id],
            "conversationId": seeded["conv_id"],
        },
        retrieval=seeded["retrieval"],
        owner_id="test-owner",
        conversation_store=seeded["conversation_store"],
        gateway_delegator=delegator,
    )

    details = _extract_details(result)
    assert "error" not in details, f"Unexpected error: {details}"
    assert details["answer"] == "The app was deployed via fly.io"
    assert seeded["s1"].summary_id in details["citedIds"]
    assert details["expandedSummaryCount"] == 1
    assert details["totalSourceTokens"] == 50

    # Verify delegator was called
    assert len(delegator.calls) == 1
    system_prompt, user_msg, timeout = delegator.calls[0]
    assert "retrieval navigator" in system_prompt.lower()
    assert "How was it deployed?" in user_msg


@pytest.mark.asyncio
async def test_expand_query_gateway_fallback(seeded):
    """Falls back to local synthesis when delegator fails."""
    failing_delegator = MockGatewayDelegator(should_fail=True)
    synth = MockSynthesizeFn("Local fallback answer")

    result = await execute_expand_query(
        params={
            "prompt": "What work was done?",
            "summaryIds": [seeded["s3"].summary_id],
            "conversationId": seeded["conv_id"],
        },
        retrieval=seeded["retrieval"],
        owner_id="test-owner",
        conversation_store=seeded["conversation_store"],
        gateway_delegator=failing_delegator,
        synthesize_fn=synth,
    )

    details = _extract_details(result)
    assert "error" not in details, f"Unexpected error: {details}"
    # Delegator was tried, then local synthesis happened
    assert len(failing_delegator.calls) == 1
    assert len(synth.calls) == 1
    assert details["answer"] == "Local fallback answer"


@pytest.mark.asyncio
async def test_expand_query_max_tokens(seeded):
    """Respects maxTokens param passed to synthesize_fn."""
    synth = MockSynthesizeFn("Short answer")
    result = await execute_expand_query(
        params={
            "prompt": "Summarize briefly",
            "summaryIds": [seeded["s1"].summary_id],
            "conversationId": seeded["conv_id"],
            "maxTokens": 500,
        },
        retrieval=seeded["retrieval"],
        owner_id="test-owner",
        conversation_store=seeded["conversation_store"],
        synthesize_fn=synth,
    )

    details = _extract_details(result)
    assert "error" not in details
    # Verify max_tokens was forwarded to synthesize_fn
    assert len(synth.calls) == 1
    _, _, max_tokens = synth.calls[0]
    assert max_tokens == 500


@pytest.mark.asyncio
async def test_expand_query_token_cap(seeded):
    """Respects tokenCap for expansion budget."""
    synth = MockSynthesizeFn("Capped answer")
    result = await execute_expand_query(
        params={
            "prompt": "What happened?",
            "summaryIds": [seeded["s1"].summary_id, seeded["s2"].summary_id],
            "conversationId": seeded["conv_id"],
            "tokenCap": 5,  # Very small cap
        },
        retrieval=seeded["retrieval"],
        owner_id="test-owner",
        conversation_store=seeded["conversation_store"],
        synthesize_fn=synth,
    )

    details = _extract_details(result)
    assert "error" not in details
    # With tokenCap=5, expansion should be truncated or very limited
    assert isinstance(details.get("totalSourceTokens", 0), int)


# ═══════════════════════════════════════════════════════════════════════════
# Unit tests for reply parsing
# ═══════════════════════════════════════════════════════════════════════════


class TestParseDelegatedReply:
    def test_valid_json(self):
        reply = json.dumps({
            "answer": "The deploy used fly.io",
            "citedIds": ["sum_abc", "sum_def"],
            "expandedSummaryCount": 2,
            "totalSourceTokens": 100,
            "truncated": False,
        })
        parsed = _parse_delegated_reply(reply, 3)
        assert parsed["answer"] == "The deploy used fly.io"
        assert parsed["citedIds"] == ["sum_abc", "sum_def"]
        assert parsed["expandedSummaryCount"] == 2
        assert parsed["totalSourceTokens"] == 100
        assert parsed["truncated"] is False

    def test_fenced_json(self):
        reply = 'Here is the result:\n```json\n{"answer": "fenced", "citedIds": ["sum_x"]}\n```'
        parsed = _parse_delegated_reply(reply, 1)
        assert parsed["answer"] == "fenced"
        assert parsed["citedIds"] == ["sum_x"]

    def test_plain_text_fallback(self):
        reply = "Just a plain text answer with no JSON"
        parsed = _parse_delegated_reply(reply, 5)
        assert parsed["answer"] == reply
        assert parsed["citedIds"] == []
        assert parsed["expandedSummaryCount"] == 5

    def test_empty_reply(self):
        parsed = _parse_delegated_reply(None, 0)
        assert parsed["answer"] == ""
        assert parsed["citedIds"] == []

    def test_invalid_json(self):
        reply = '{"answer": "broken'
        parsed = _parse_delegated_reply(reply, 2)
        assert parsed["answer"] == reply
        assert parsed["expandedSummaryCount"] == 2
