from __future__ import annotations
import logging

from langgraph.graph import END, START, StateGraph

from app.modules.agent.engine.graphs.progress import emit_progress_sync
from app.modules.agent.engine.graphs.state import AgentGraphState
from app.modules.agent.engine.orchestrator.actions.project_qa_analyzer import ProjectQaAnalyzer
from app.modules.agent.engine.orchestrator.actions.project_qa_support import ProjectQaSupport
from app.modules.agent.llm import AgentLlmService
from app.modules.contracts import RagRetriever
from app.modules.rag.explain import ExplainPack, PromptBudgeter

LOGGER = logging.getLogger(__name__)


class ProjectQaConversationGraphFactory:
    def __init__(self, llm: AgentLlmService | None = None) -> None:
        self._support = ProjectQaSupport()

    def build(self, checkpointer=None):
        graph = StateGraph(AgentGraphState)
        graph.add_node("resolve_request", self._resolve_request)
        graph.add_edge(START, "resolve_request")
        graph.add_edge("resolve_request", END)
        return graph.compile(checkpointer=checkpointer)

    def _resolve_request(self, state: AgentGraphState) -> dict:
        emit_progress_sync(state, stage="graph.project_qa.conversation_understanding", message="Нормализую пользовательский запрос.")
        resolved = self._support.resolve_request(str(state.get("message", "") or ""))
        LOGGER.warning("graph step result: graph=project_qa/conversation_understanding normalized=%s", resolved.get("normalized_message", ""))
        return {"resolved_request": resolved}


class ProjectQaClassificationGraphFactory:
    def __init__(self, llm: AgentLlmService | None = None) -> None:
        self._support = ProjectQaSupport()

    def build(self, checkpointer=None):
        graph = StateGraph(AgentGraphState)
        graph.add_node("classify_question", self._classify_question)
        graph.add_edge(START, "classify_question")
        graph.add_edge("classify_question", END)
        return graph.compile(checkpointer=checkpointer)

    def _classify_question(self, state: AgentGraphState) -> dict:
        resolved = state.get("resolved_request", {}) or {}
        message = str(resolved.get("normalized_message") or state.get("message", "") or "")
        profile = self._support.build_profile(message)
        LOGGER.warning("graph step result: graph=project_qa/question_classification domain=%s intent=%s", profile.get("domain"), profile.get("intent"))
        return {"question_profile": profile}


class ProjectQaRetrievalGraphFactory:
    def __init__(self, rag: RagRetriever, llm: AgentLlmService | None = None) -> None:
        self._rag = rag
        self._support = ProjectQaSupport()

    def build(self, checkpointer=None):
        graph = StateGraph(AgentGraphState)
        graph.add_node("retrieve_context", self._retrieve_context)
        graph.add_edge(START, "retrieve_context")
        graph.add_edge("retrieve_context", END)
        return graph.compile(checkpointer=checkpointer)

    def _retrieve_context(self, state: AgentGraphState) -> dict:
        emit_progress_sync(state, stage="graph.project_qa.context_retrieval", message="Собираю контекст по проекту.")
        resolved = state.get("resolved_request", {}) or {}
        profile = state.get("question_profile", {}) or {}
        files_map = dict(state.get("files_map", {}) or {})
        rag_items: list[dict] = []
        source_bundle = self._support.build_source_bundle(profile, list(rag_items), files_map)
        LOGGER.warning(
            "graph step result: graph=project_qa/context_retrieval mode=%s rag_items=%s file_candidates=%s legacy_rag=%s",
            profile.get("domain"),
            len(source_bundle.get("rag_items", []) or []),
            len(source_bundle.get("file_candidates", []) or []),
            False,
        )
        return {"source_bundle": source_bundle}


class ProjectQaAnalysisGraphFactory:
    def __init__(self, llm: AgentLlmService | None = None) -> None:
        self._support = ProjectQaSupport()
        self._analyzer = ProjectQaAnalyzer()

    def build(self, checkpointer=None):
        graph = StateGraph(AgentGraphState)
        graph.add_node("analyze_context", self._analyze_context)
        graph.add_edge(START, "analyze_context")
        graph.add_edge("analyze_context", END)
        return graph.compile(checkpointer=checkpointer)

    def _analyze_context(self, state: AgentGraphState) -> dict:
        explain_pack = state.get("explain_pack")
        if explain_pack:
            analysis = self._analysis_from_pack(explain_pack)
            LOGGER.warning(
                "graph step result: graph=project_qa/context_analysis findings=%s evidence=%s",
                len(analysis.get("findings", []) or []),
                len(analysis.get("evidence", []) or []),
            )
            return {"analysis_brief": analysis}
        bundle = state.get("source_bundle", {}) or {}
        profile = bundle.get("profile", {}) or state.get("question_profile", {}) or {}
        rag_items = list(bundle.get("rag_items", []) or [])
        file_candidates = list(bundle.get("file_candidates", []) or [])
        analysis = self._analyzer.analyze_code(profile, rag_items, file_candidates) if str(profile.get("domain")) == "code" else self._analyzer.analyze_docs(profile, rag_items)
        LOGGER.warning(
            "graph step result: graph=project_qa/context_analysis findings=%s evidence=%s",
            len(analysis.get("findings", []) or []),
            len(analysis.get("evidence", []) or []),
        )
        return {"analysis_brief": analysis}

    def _analysis_from_pack(self, raw_pack) -> dict:
        pack = ExplainPack.model_validate(raw_pack)
        findings: list[str] = []
        evidence: list[str] = []
        for entrypoint in pack.selected_entrypoints[:3]:
            findings.append(f"Entrypoint `{entrypoint.title}` maps to handler `{entrypoint.metadata.get('handler_symbol_id', '')}`.")
            if entrypoint.source:
                evidence.append(entrypoint.source)
        for path in pack.trace_paths[:3]:
            if path.symbol_ids:
                findings.append(f"Trace path: {' -> '.join(path.symbol_ids)}")
        for excerpt in pack.code_excerpts[:4]:
            evidence.append(f"{excerpt.path}:{excerpt.start_line}-{excerpt.end_line} [{excerpt.evidence_id}]")
        return {
            "subject": pack.intent.normalized_query,
            "findings": findings or ["No explain trace was built from the available code evidence."],
            "evidence": evidence,
            "gaps": list(pack.missing),
            "answer_mode": "summary",
        }


class ProjectQaAnswerGraphFactory:
    def __init__(self, llm: AgentLlmService | None = None) -> None:
        self._support = ProjectQaSupport()
        self._llm = llm
        self._budgeter = PromptBudgeter()

    def build(self, checkpointer=None):
        graph = StateGraph(AgentGraphState)
        graph.add_node("compose_answer", self._compose_answer)
        graph.add_edge(START, "compose_answer")
        graph.add_edge("compose_answer", END)
        return graph.compile(checkpointer=checkpointer)

    def _compose_answer(self, state: AgentGraphState) -> dict:
        profile = state.get("question_profile", {}) or {}
        analysis = state.get("analysis_brief", {}) or {}
        brief = self._support.build_answer_brief(profile, analysis)
        explain_pack = state.get("explain_pack")
        answer = self._compose_explain_answer(state, explain_pack)
        if not answer:
            answer = self._support.compose_answer(brief)
        LOGGER.warning("graph step result: graph=project_qa/answer_composition answer_len=%s", len(answer or ""))
        return {"answer_brief": brief, "final_answer": answer}

    def _compose_explain_answer(self, state: AgentGraphState, raw_pack) -> str:
        if raw_pack is None or self._llm is None:
            return ""
        pack = ExplainPack.model_validate(raw_pack)
        prompt_input = self._budgeter.build_prompt_input(str(state.get("message", "") or ""), pack)
        return self._llm.generate(
            "code_explain_answer_v2",
            prompt_input,
            log_context="graph.project_qa.answer_v2",
        ).strip()
