from __future__ import annotations

from datetime import datetime
from difflib import get_close_matches
from time import perf_counter

from app.modules.rag.intent_router_v2 import ConversationState, IntentRouterV2
from tests.pipeline_setup.suite_02_pipeline.pipeline_intent_rag.helpers.diagnostics import (
    apply_retrieval_report,
    assign_repo_scope,
    init_diagnostics,
    validate_constraints,
)
from tests.pipeline_setup.suite_02_pipeline.pipeline_intent_rag.helpers.models import PhraseCase, PipelineResult
from tests.unit_tests.rag.intent_router_testkit import repo_context


class RouterOnlyRunner:
    def __init__(self, started_at: datetime, router: IntentRouterV2 | None = None) -> None:
        self._started_at = started_at
        self._router = router or IntentRouterV2()

    def run_case(self, case: PhraseCase) -> PipelineResult:
        router_started = perf_counter()
        route_result = self._router.route(case.text, ConversationState(), repo_context())
        router_ms = _ms(router_started)
        query = route_result.query_plan.normalized or route_result.query_plan.raw or case.text
        diagnostics = init_diagnostics(route_result, router_ms=router_ms, include_retrieval=False)
        return PipelineResult(
            case=case,
            mode="router_only",
            run_started_at=self._started_at,
            rag_session_id=case.rag_session_id,
            intent=route_result.intent,
            graph_id=route_result.graph_id,
            conversation_mode=route_result.conversation_mode,
            query=query,
            symbol_resolution=route_result.symbol_resolution.model_dump(),
            diagnostics=diagnostics,
        )


class IntentRouterRagPipelineRunner:
    def __init__(
        self,
        started_at: datetime,
        rag_adapter,
        session_resolver,
        router: IntentRouterV2 | None = None,
    ) -> None:
        self._started_at = started_at
        self._rag_adapter = rag_adapter
        self._session_resolver = session_resolver
        self._router = router or IntentRouterV2()

    def run_case(self, case: PhraseCase) -> PipelineResult:
        router_started = perf_counter()
        route_result = self._router.route(case.text, ConversationState(), repo_context())
        router_ms = _ms(router_started)
        diagnostics = init_diagnostics(route_result, router_ms=router_ms, include_retrieval=True)
        rag_query = route_result.query_plan.normalized or route_result.query_plan.raw or case.text
        rag_session_id = self._session_resolver.resolve(case)
        retrieval_started = perf_counter()
        rag_rows = self._retrieve_rows(
            rag_session_id=rag_session_id,
            rag_query=rag_query,
            route_result=route_result,
        )
        diagnostics["timings_ms"]["retrieval_total"] = _ms(retrieval_started)
        consume_report = getattr(self._rag_adapter, "consume_retrieval_report", None)
        retrieval_report = consume_report() if callable(consume_report) else None
        apply_retrieval_report(diagnostics, retrieval_report)
        symbol_started = perf_counter()
        symbol_resolution = self._resolve_symbol(route_result.symbol_resolution.model_dump(), rag_rows)
        diagnostics["timings_ms"]["symbol_resolution"] = _ms(symbol_started)
        if (
            route_result.query_plan.sub_intent == "EXPLAIN"
            and list(route_result.query_plan.symbol_candidates or [])
            and str(symbol_resolution.get("status") or "") != "resolved"
        ):
            rag_rows = []
            if isinstance(diagnostics.get("retrieval"), dict):
                diagnostics["retrieval"]["fallback"] = {"used": False, "reason": None}
        merge_started = perf_counter()
        rag_rows = self._rag_adapter.hydrate_resolved_symbol_sources(
            rag_session_id=rag_session_id,
            base_query=rag_query,
            rag_rows=rag_rows,
            symbol_resolution=symbol_resolution,
            retrieval_spec=route_result.retrieval_spec,
            retrieval_constraints=route_result.retrieval_constraints,
        )
        rag_rows = self._rag_adapter.force_symbol_context_c0(
            rag_session_id=rag_session_id,
            rag_rows=rag_rows,
            symbol_resolution=symbol_resolution,
            limit=20,
        )
        diagnostics["timings_ms"]["merge_rank"] = _ms(merge_started)
        expected_repo_id = str(getattr(route_result.retrieval_spec.filters, "repo_id", "") or "").strip() or None
        assign_repo_scope(diagnostics, rag_rows, expected_repo_id=expected_repo_id)
        diagnostics["constraint_violations"] = validate_constraints(
            route_result,
            rag_rows,
            symbol_resolution,
            expected_repo_id=expected_repo_id,
        )
        return PipelineResult(
            case=case,
            mode="router_rag",
            run_started_at=self._started_at,
            rag_session_id=rag_session_id,
            intent=route_result.intent,
            graph_id=route_result.graph_id,
            conversation_mode=route_result.conversation_mode,
            query=rag_query,
            rag_rows=rag_rows,
            symbol_resolution=symbol_resolution,
            diagnostics=diagnostics,
        )

    def _resolve_symbol(self, initial: dict, rag_rows: list[dict]) -> dict:
        status = str(initial.get("status") or "not_requested")
        if status != "pending":
            return initial
        candidates = [str(item) for item in initial.get("alternatives", []) if str(item).strip()]
        c1_titles = list(dict.fromkeys(
            str(row.get("title") or "").strip()
            for row in rag_rows
            if str(row.get("layer") or "") == "C1_SYMBOL_CATALOG" and str(row.get("title") or "").strip()
        ))
        if not c1_titles:
            initial["status"] = "not_found"
            return initial
        exact = next((title for title in c1_titles if title in candidates), None)
        if exact:
            return {"status": "resolved", "resolved_symbol": exact, "alternatives": c1_titles[:5], "confidence": 0.99}
        fuzzy_hits: list[str] = []
        for candidate in candidates:
            fuzzy_hits.extend(get_close_matches(candidate, c1_titles, n=3, cutoff=0.7))
        fuzzy_hits = list(dict.fromkeys(fuzzy_hits))
        if len(fuzzy_hits) == 1:
            return {"status": "resolved", "resolved_symbol": fuzzy_hits[0], "alternatives": c1_titles[:5], "confidence": 0.82}
        if len(fuzzy_hits) > 1:
            return {"status": "ambiguous", "resolved_symbol": None, "alternatives": fuzzy_hits[:5], "confidence": 0.55}
        return {"status": "not_found", "resolved_symbol": None, "alternatives": c1_titles[:5], "confidence": 0.0}

    def _retrieve_rows(self, *, rag_session_id: str, rag_query: str, route_result) -> list[dict]:
        path_scope = list(getattr(route_result.retrieval_spec.filters, "path_scope", []) or [])
        expected_repo_id = str(getattr(route_result.retrieval_spec.filters, "repo_id", "") or "").strip() or None
        if route_result.query_plan.sub_intent == "OPEN_FILE" and path_scope:
            return self._rag_adapter.retrieve_exact_files(
                rag_session_id=rag_session_id,
                repo_id=expected_repo_id,
                paths=path_scope,
                layers=["C0_SOURCE_CHUNKS"],
                limit=200,
                query=rag_query,
                ranking_profile=str(getattr(route_result.retrieval_spec, "rerank_profile", "") or ""),
            )
        return self._rag_adapter.retrieve_with_plan(
            rag_session_id,
            rag_query,
            route_result.retrieval_spec,
            route_result.retrieval_constraints,
            query_plan=route_result.query_plan,
        )


def _ms(started: float) -> int:
    return int((perf_counter() - started) * 1000)
