from __future__ import annotations

import inspect
import json
from pathlib import Path

import pytest

from src.rag.qa_v1_1_73 import api as api_module
from src.rag.qa_v1_1_73.pipeline import QAV1Pipeline
from src.rag.qa_v1_1_73.subquestion_orchestrator import SubquestionRetrievalOrchestrator


def test_pipeline_accepts_optional_progress_callback():
    assert "progress_callback" in inspect.signature(QAV1Pipeline.answer).parameters


def test_subquestion_orchestrator_reports_review_and_bounded_retry():
    events = []
    review_calls = []

    def retrieve(task, *, attempt, refinement):
        return {
            "evidence_top30": [{"evidence_id": f"{task['subquestion_id']}.A{attempt}.E1"}],
            "chains": [],
        }

    def review(_question, groups, *, attempt):
        review_calls.append(attempt)
        return {
            "groups": [
                {
                    "subquestion_id": group["subquestion_id"],
                    "status": "PARTIAL" if attempt == 1 else "SATISFIED",
                    "accepted_evidence_ids": [group["evidence"][0]["evidence_id"]],
                    "accepted_chain_ids": [],
                    "refinement": {"vector_queries": ["refined"]} if attempt == 1 else {},
                }
                for group in groups
            ]
        }

    result = SubquestionRetrievalOrchestrator().run(
        "test question",
        {
            "subquestions": [{"subquestion_id": "Q0", "question": "test question"}],
            "synthesis_order": ["Q0"],
        },
        retrieve=retrieve,
        review=review,
        progress_callback=events.append,
    )

    assert review_calls == [1, 2]
    assert result["audit_groups"][0]["attempt_count"] == 2
    assert [event["stage"] for event in events] == [
        "subquestion_retrieved",
        "agent_review_started",
        "subquestion_retry_started",
        "subquestion_retrieved",
        "agent_review_started",
        "evidence_review_completed",
    ]


def test_api_exposes_sse_chat_stream_without_replacing_chat():
    source = (
        Path(__file__).resolve().parents[1] / "src/rag/qa_v1_1_73/api.py"
    ).read_text(encoding="utf-8")

    assert '@router.post("/chat")' in source
    assert '@router.post("/chat/stream")' in source
    assert "StreamingResponse" in source
    assert "progress_callback" in source


@pytest.mark.asyncio
async def test_sse_chat_stream_emits_progress_before_complete(monkeypatch):
    class FakeLLM:
        available = True

    class FakePipeline:
        llm = FakeLLM()

        def answer(self, question, **kwargs):
            kwargs["progress_callback"]({
                "stage": "entities_identified",
                "details": {"entities": ["WRKY"]},
            })
            return {"success": True, "question": question, "answer": "answer"}

    monkeypatch.setattr(api_module, "get_pipeline", lambda: FakePipeline())
    response = await api_module.chat_stream(api_module.QARequest(question="WRKY?"))
    body = b""
    async for chunk in response.body_iterator:
        body += chunk.encode("utf-8") if isinstance(chunk, str) else chunk
    text = body.decode("utf-8")

    assert text.index("event: progress") < text.index("event: complete")
    complete_data = text.split("event: complete", 1)[1].split("data:", 1)[1].strip()
    assert json.loads(complete_data)["answer"] == "answer"