from __future__ import annotations

import logging
from time import perf_counter

from app.modules.agent.code_qa_runtime.answer_policy import CodeQaAnswerPolicy
from app.modules.agent.code_qa_runtime.models import CodeQaDraftAnswer, CodeQaExecutionState, CodeQaFinalResult
from app.modules.agent.code_qa_runtime.post_gate import CodeQaPostEvidenceGate
from app.modules.agent.code_qa_runtime.prompt_payload_builder import CodeQaPromptPayloadBuilder
from app.modules.agent.code_qa_runtime.prompt_selector import CodeQaPromptSelector
from app.modules.agent.code_qa_runtime.repair import CodeQaAnswerRepairService
from app.modules.agent.code_qa_runtime.repo_context import CodeQaRepoContextFactory
from app.modules.agent.code_qa_runtime.retrieval_adapter import CodeQaRetrievalAdapter
from app.modules.agent.llm import AgentLlmService
from app.modules.rag.code_qa_pipeline.answer_synthesis import build_answer_synthesis_input
from app.modules.rag.code_qa_pipeline.diagnostics import build_diagnostics_report
from app.modules.rag.code_qa_pipeline.evidence_bundle_builder import build_evidence_bundle
from app.modules.rag.code_qa_pipeline.evidence_gate import evaluate_evidence
from app.modules.rag.code_qa_pipeline.retrieval_request_builder import build_retrieval_request
from app.modules.rag.code_qa_pipeline.retrieval_result_builder import build_retrieval_result
from app.modules.rag.intent_router_v2 import ConversationState, IntentRouterV2
from app.modules.rag.intent_router_v2.models import SymbolResolution

LOGGER = logging.getLogger(__name__)


class CodeQaRuntimeExecutor:
    def __init__(
        self,
        llm: AgentLlmService | None,
        *,
        router: IntentRouterV2 | None = None,
        retrieval: CodeQaRetrievalAdapter | None = None,
        repo_context_factory: CodeQaRepoContextFactory | None = None,
        prompt_selector: CodeQaPromptSelector | None = None,
        payload_builder: CodeQaPromptPayloadBuilder | None = None,
        answer_policy: CodeQaAnswerPolicy | None = None,
        post_gate: CodeQaPostEvidenceGate | None = None,
    ) -> None:
        self._llm = llm
        self._router = router or IntentRouterV2()
        self._retrieval = retrieval or CodeQaRetrievalAdapter()
        self._repo_context_factory = repo_context_factory or CodeQaRepoContextFactory()
        self._prompt_selector = prompt_selector or CodeQaPromptSelector()
        self._payload_builder = payload_builder or CodeQaPromptPayloadBuilder()
        self._answer_policy = answer_policy or CodeQaAnswerPolicy()
        self._post_gate = post_gate or CodeQaPostEvidenceGate()
        self._repair = CodeQaAnswerRepairService(llm) if llm is not None else None

    def execute(self, *, user_query: str, rag_session_id: str, files_map: dict[str, dict] | None = None) -> CodeQaFinalResult:
        timings_ms: dict[str, int] = {}
        runtime_trace: list[dict] = []
        state = CodeQaExecutionState(
            user_query=user_query,
            rag_session_id=rag_session_id,
            conversation_state=ConversationState(),
            repo_context=self._repo_context_factory.build(files_map),
        )
        started = perf_counter()
        state.router_result = self._router.route(user_query, state.conversation_state, state.repo_context)
        timings_ms["router"] = self._elapsed_ms(started)
        runtime_trace.append(
            {
                "step": "router",
                "status": "completed",
                "timings_ms": {"router": timings_ms["router"]},
                "output": {
                    "intent": state.router_result.intent,
                    "sub_intent": state.router_result.query_plan.sub_intent,
                    "graph_id": state.router_result.graph_id,
                    "conversation_mode": state.router_result.conversation_mode,
                },
            }
        )
        state.retrieval_request = build_retrieval_request(state.router_result, rag_session_id)
        started = perf_counter()
        raw_rows = self._retrieve(state)
        timings_ms["retrieval"] = self._elapsed_ms(started)
        symbol_resolution = self._resolve_symbol(state.router_result.symbol_resolution.model_dump(), raw_rows)
        state.router_result = state.router_result.model_copy(update={"symbol_resolution": SymbolResolution(**symbol_resolution)})
        retrieval_report = self._retrieval.consume_retrieval_report()
        state.retrieval_result = build_retrieval_result(raw_rows, retrieval_report, symbol_resolution)
        if state.retrieval_request.sub_intent.upper() == "EXPLAIN" and symbol_resolution.get("status") in {"not_found", "ambiguous"}:
            state.retrieval_result = build_retrieval_result([], retrieval_report, symbol_resolution)
        runtime_trace.append(
            {
                "step": "retrieval",
                "status": "completed",
                "timings_ms": {"retrieval": timings_ms["retrieval"]},
                "output": {
                    "rag_count": len(raw_rows),
                    "answer_path_rag_count": len(state.retrieval_result.raw_rows),
                    "resolved_symbol_status": symbol_resolution.get("status"),
                    "resolved_symbol": symbol_resolution.get("resolved_symbol"),
                    "requested_layers": list(state.retrieval_request.requested_layers or []),
                },
                "diagnostics": retrieval_report or {},
            }
        )
        state.evidence_pack = build_evidence_bundle(state.retrieval_result, state.router_result)
        if state.retrieval_request.sub_intent.upper() == "EXPLAIN" and symbol_resolution.get("status") in {"not_found", "ambiguous"}:
            state.evidence_pack.sufficient = False
            state.evidence_pack.failure_reasons = ["target_not_resolved"]
        gate_decision = evaluate_evidence(state.evidence_pack)
        timings_ms["pre_evidence_gate"] = 0
        state.answer_mode = "normal" if gate_decision.passed else "degraded"
        state.degraded_message = gate_decision.degraded_message
        runtime_trace.append(
            {
                "step": "pre_evidence_gate",
                "status": "passed" if gate_decision.passed else "blocked",
                "timings_ms": {"pre_evidence_gate": timings_ms["pre_evidence_gate"]},
                "output": {
                    "passed": gate_decision.passed,
                    "failure_reasons": list(gate_decision.failure_reasons),
                    "degraded_message": gate_decision.degraded_message,
                },
            }
        )
        decision = self._answer_policy.decide(router_result=state.router_result, gate_decision=gate_decision)
        if not decision.should_call_llm:
            state.answer_mode = decision.answer_mode
            runtime_trace.append(
                {
                    "step": "llm",
                    "status": "skipped",
                    "output": {
                        "reason": "policy_short_circuit",
                        "answer_mode": decision.answer_mode,
                    },
                }
            )
            runtime_trace.append(
                {
                    "step": "post_evidence_gate",
                    "status": "skipped",
                    "output": {"reason": "no_draft_answer"},
                }
            )
            return self._finalize(
                state,
                draft=None,
                final_answer=decision.answer,
                repair_used=False,
                llm_used=False,
                timings_ms=timings_ms,
                runtime_trace=runtime_trace,
            )
        if self._llm is None:
            runtime_trace.append(
                {
                    "step": "llm",
                    "status": "skipped",
                    "output": {"reason": "llm_unavailable"},
                }
            )
            return self._finalize(
                state,
                draft=None,
                final_answer="",
                repair_used=False,
                llm_used=False,
                timings_ms=timings_ms,
                runtime_trace=runtime_trace,
            )
        state.synthesis_input = build_answer_synthesis_input(user_query, state.evidence_pack)
        prompt_name = self._prompt_selector.select(sub_intent=state.retrieval_request.sub_intent, answer_mode=state.answer_mode)
        prompt_payload = self._payload_builder.build(
            user_query=user_query,
            synthesis_input=state.synthesis_input,
            evidence_pack=state.evidence_pack,
            answer_mode=state.answer_mode,
        )
        started = perf_counter()
        draft = CodeQaDraftAnswer(
            prompt_name=prompt_name,
            prompt_payload=prompt_payload,
            answer=self._llm.generate(prompt_name, prompt_payload, log_context="graph.project_qa.code_qa.answer").strip(),
        )
        timings_ms["llm"] = self._elapsed_ms(started)
        runtime_trace.append(
            {
                "step": "llm",
                "status": "completed",
                "timings_ms": {"llm": timings_ms["llm"]},
                "output": {
                    "prompt_name": prompt_name,
                    "answer_preview": draft.answer[:300],
                },
            }
        )
        validation = self._post_gate.validate(answer=draft.answer, answer_mode=state.answer_mode, degraded_message=state.degraded_message)
        final_answer = draft.answer
        repair_used = False
        if not validation.passed and self._repair is not None:
            started = perf_counter()
            final_answer = self._repair.repair(draft_answer=draft.answer, validation=validation, prompt_payload=prompt_payload)
            repair_used = True
            timings_ms["repair"] = self._elapsed_ms(started)
            validation = self._post_gate.validate(answer=final_answer, answer_mode=state.answer_mode, degraded_message=state.degraded_message)
            if not validation.passed and state.degraded_message:
                final_answer = state.degraded_message
        runtime_trace.append(
            {
                "step": "post_evidence_gate",
                "status": "passed" if validation.passed else "failed",
                "timings_ms": {
                    "post_evidence_gate": 0,
                    "repair": timings_ms.get("repair", 0),
                },
                "output": {
                    "passed": validation.passed,
                    "reasons": list(validation.reasons),
                    "repair_used": repair_used,
                },
            }
        )
        return self._finalize(
            state,
            draft=draft,
            final_answer=final_answer,
            repair_used=repair_used,
            llm_used=True,
            validation=validation,
            timings_ms=timings_ms,
            runtime_trace=runtime_trace,
        )

    def _retrieve(self, state: CodeQaExecutionState) -> list[dict]:
        assert state.retrieval_request is not None
        if state.retrieval_request.sub_intent == "OPEN_FILE" and state.retrieval_request.path_scope:
            return self._retrieval.retrieve_exact_files(
                state.rag_session_id,
                paths=state.retrieval_request.path_scope,
                layers=["C0_SOURCE_CHUNKS"],
                limit=200,
                query=state.retrieval_request.query,
                ranking_profile=str(getattr(state.retrieval_request.retrieval_spec, "rerank_profile", "") or ""),
            )
        return self._retrieval.retrieve_with_plan(
            state.rag_session_id,
            state.retrieval_request.query,
            state.retrieval_request.retrieval_spec,
            state.retrieval_request.retrieval_constraints,
            query_plan=state.retrieval_request.query_plan,
        )

    def _resolve_symbol(self, initial: dict, rag_rows: list[dict]) -> dict:
        if str(initial.get("status") or "") != "pending":
            return initial
        candidates = [str(item).strip() for item in initial.get("alternatives", []) if str(item).strip()]
        found = [
            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()
        ]
        exact = next((item for item in found if item in candidates), None)
        if exact:
            return {"status": "resolved", "resolved_symbol": exact, "alternatives": found[:5], "confidence": 0.99}
        if found:
            return {"status": "ambiguous", "resolved_symbol": None, "alternatives": found[:5], "confidence": 0.55}
        return {"status": "not_found", "resolved_symbol": None, "alternatives": [], "confidence": 0.0}

    def _finalize(
        self,
        state: CodeQaExecutionState,
        *,
        draft: CodeQaDraftAnswer | None,
        final_answer: str,
        repair_used: bool,
        llm_used: bool,
        validation=None,
        timings_ms: dict[str, int] | None = None,
        runtime_trace: list[dict] | None = None,
    ) -> CodeQaFinalResult:
        diagnostics = build_diagnostics_report(
            router_result=state.router_result,
            retrieval_request=state.retrieval_request,
            retrieval_result=state.retrieval_result,
            evidence_bundle=state.evidence_pack,
            answer_mode=state.answer_mode,
            timings_ms=timings_ms or {},
        )
        result = CodeQaFinalResult(
            final_answer=final_answer.strip(),
            answer_mode=state.answer_mode,
            repair_used=repair_used,
            llm_used=llm_used,
            draft_answer=draft,
            validation=validation or self._post_gate.validate(answer=final_answer, answer_mode=state.answer_mode, degraded_message=state.degraded_message),
            router_result=state.router_result,
            retrieval_request=state.retrieval_request,
            retrieval_result=state.retrieval_result,
            evidence_pack=state.evidence_pack,
            diagnostics=diagnostics,
            runtime_trace=list(runtime_trace or []),
        )
        LOGGER.warning(
            "code qa runtime executed: intent=%s sub_intent=%s answer_mode=%s repair_used=%s llm_used=%s",
            state.router_result.intent,
            state.router_result.query_plan.sub_intent,
            result.answer_mode,
            result.repair_used,
            result.llm_used,
        )
        return result

    def _elapsed_ms(self, started: float) -> int:
        return int((perf_counter() - started) * 1000)
