"""Tests for retrieval learning, reranking, and reindex triggers."""
from __future__ import annotations

import httpx
import pytest

from logos_adk.agent import Agent
from logos_adk.context import RunContext, use_run_context
from logos_adk.learning import FeedbackAPI, RetrievalLearningEngine, RunRecord, SQLiteLearningStore
from logos_adk.providers.mock import MockProvider
from logos_adk.rag.knowledge_base import KnowledgeBase
from logos_adk.rag.vector_store import InMemoryVectorStore


async def _request(app, method: str, path: str, **kwargs):
    transport = httpx.ASGITransport(app=app)
    async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
        return await client.request(method, path, **kwargs)


async def _embed(text: str) -> list[float]:
    normalized = text.lower()
    if "refund" in normalized or "return" in normalized:
        return [1.0, 0.0]
    return [0.0, 1.0]


@pytest.mark.asyncio
async def test_retrieval_learning_reuses_query_rewrite_and_tuned_profile(tmp_path):
    engine = RetrievalLearningEngine(str(tmp_path / "retrieval.db"))
    learning_store = SQLiteLearningStore(str(tmp_path / "learning.db"))
    feedback_api = FeedbackAPI(learning_store)
    kb = KnowledgeBase(
        store=InMemoryVectorStore(),
        embed_fn=_embed,
        retrieval_learning=engine,
    )
    await kb.ingest_text("Refund policy for premium users", metadata={"topic": "refund", "source_key": "refund-doc"})
    await kb.ingest_text("Shipping policy", metadata={"topic": "shipping", "source_key": "shipping-doc"})

    engine.record_query_rewrite(
        original_query="return window",
        rewritten_query="refund policy",
        workspace_id="acme",
        task_type="support",
        user_segment="vip",
        succeeded=True,
    )

    learning_store.create_run(
        RunRecord(
            id="run-1",
            trace_id="trace-1",
            request_id="req-1",
            agent_name="kb-test",
            input_text="return window",
            output_text="refund policy",
            session_id="s-1",
            workspace_id="acme",
            prompt_version="v1",
            model="mock",
            config_fingerprint="cfg-1",
        )
    )
    with use_run_context(
        RunContext(
            trace_id="trace-1",
            session_id="s-1",
            workspace_id="acme",
            task_type="support",
            user_segment="vip",
            prompt_name="kb-test",
            prompt_version="v1",
        )
    ):
        results = await kb.search("return window", limit=2, filter={"topic": "refund"})
    assert results
    feedback_api.record_outcome(trace_id="trace-1", kind="conversion", value=True)
    engine.sync_trace("trace-1", learning_store)

    learning_store.create_run(
        RunRecord(
            id="run-2",
            trace_id="trace-2",
            request_id="req-2",
            agent_name="kb-test",
            input_text="return window",
            output_text="refund policy",
            session_id="s-2",
            workspace_id="acme",
            prompt_version="v1",
            model="mock",
            config_fingerprint="cfg-1",
        )
    )
    with use_run_context(
        RunContext(
            trace_id="trace-2",
            session_id="s-2",
            workspace_id="acme",
            task_type="support",
            user_segment="vip",
            prompt_name="kb-test",
            prompt_version="v1",
        )
    ):
        await kb.search("return window", limit=5)

    event = engine.list_events(trace_id="trace-2")[0]
    assert event.effective_query == "refund policy"
    assert event.limit == 2
    assert event.filter == {"topic": "refund"}
    rewrites = engine.list_query_rewrites(workspace_id="acme", task_type="support", user_segment="vip")
    assert rewrites[0].success_count >= 1


@pytest.mark.asyncio
async def test_retrieval_learning_reranks_sources_and_creates_reindex_triggers(tmp_path):
    from logos_adk.serve import create_app

    engine = RetrievalLearningEngine(str(tmp_path / "retrieval.db"))
    learning_store = SQLiteLearningStore(str(tmp_path / "learning.db"))
    kb = KnowledgeBase(
        store=InMemoryVectorStore(),
        embed_fn=_embed,
        retrieval_learning=engine,
    )
    await kb.ingest_text("Refund policy from old connector", metadata={"source_key": "bad-doc", "connector": "notion"})
    await kb.ingest_text("Accurate refund policy from curated connector", metadata={"source_key": "good-doc", "connector": "confluence"})

    agent = Agent("bot", model=MockProvider(responses=["answer"]))
    agent.knowledge(kb, mode="memory", limit=2)
    app = create_app([agent], learning_store=learning_store)

    run_response = await _request(
        app,
        "POST",
        "/run",
        json={"prompt": "refund", "session_id": "lead-1"},
        headers={
            "X-Workspace-Id": "acme",
            "X-Trace-Id": "trace-run-1",
            "X-Task-Type": "support",
            "X-User-Segment": "vip",
        },
    )
    assert run_response.status_code == 200
    run_payload = run_response.json()

    bad_feedback = await _request(
        app,
        "POST",
        "/feedback",
        json={
            "run_id": run_payload["run_id"],
            "kind": "thumbs_down",
            "failure_types": ["bad_retrieval"],
            "metadata": {"rejected_contexts": ["bad-doc"]},
        },
        headers={"X-Workspace-Id": "acme"},
    )
    assert bad_feedback.status_code == 200

    good_outcome = await _request(
        app,
        "POST",
        "/outcomes",
        json={
            "run_id": run_payload["run_id"],
            "kind": "conversion",
            "value": True,
            "metadata": {"accepted_contexts": ["good-doc"]},
        },
        headers={"X-Workspace-Id": "acme"},
    )
    assert good_outcome.status_code == 200

    scores = {item.source_key: item for item in engine.list_source_scores(workspace_id="acme")}
    assert scores["good-doc"].usefulness_score > scores["bad-doc"].usefulness_score

    triggers = engine.list_reindex_triggers(workspace_id="acme")
    assert triggers
    assert triggers[0].source_key == "bad-doc"
    assert triggers[0].suggested_chunk_size is not None
    assert triggers[0].suggested_overlap is not None

    events = engine.list_events(trace_id="trace-run-1")
    assert events[0].bad_context_detected is True

    with use_run_context(
        RunContext(
            trace_id="trace-run-2",
            session_id="acme:lead-2",
            workspace_id="acme",
            task_type="support",
            user_segment="vip",
            prompt_name="bot",
            prompt_version="unversioned",
        )
    ):
        reranked = await kb.search("refund", limit=2)

    assert reranked[0].metadata["source_key"] == "good-doc"
