from langgraph.graph import END, START, StateGraph
import logging

from app.modules.agent.engine.graphs.file_targeting import FileTargeting
from app.modules.agent.engine.graphs.docs_graph_logic import DocsContentComposer, DocsContextAnalyzer
from app.modules.agent.engine.graphs.progress import emit_progress_sync
from app.modules.agent.engine.graphs.state import AgentGraphState
from app.modules.agent.llm import AgentLlmService

LOGGER = logging.getLogger(__name__)


class DocsGraphFactory:
    _max_validation_attempts = 2

    def __init__(self, llm: AgentLlmService) -> None:
        self._targeting = FileTargeting()
        self._analyzer = DocsContextAnalyzer(llm, self._targeting)
        self._composer = DocsContentComposer(llm, self._targeting)

    def build(self, checkpointer=None):
        graph = StateGraph(AgentGraphState)
        graph.add_node("collect_code_context", self._collect_code_context)
        graph.add_node("detect_existing_docs", self._detect_existing_docs)
        graph.add_node("decide_strategy", self._decide_strategy)
        graph.add_node("load_rules_and_examples", self._load_rules_and_examples)
        graph.add_node("plan_incremental_changes", self._plan_incremental_changes)
        graph.add_node("plan_new_document", self._plan_new_document)
        graph.add_node("generate_doc_content", self._generate_doc_content)
        graph.add_node("self_check", self._self_check)
        graph.add_node("build_changeset", self._build_changeset)
        graph.add_node("summarize_result", self._summarize_result)

        graph.add_edge(START, "collect_code_context")
        graph.add_edge("collect_code_context", "detect_existing_docs")
        graph.add_edge("detect_existing_docs", "decide_strategy")
        graph.add_edge("decide_strategy", "load_rules_and_examples")
        graph.add_conditional_edges(
            "load_rules_and_examples",
            self._route_after_rules_loading,
            {
                "incremental": "plan_incremental_changes",
                "from_scratch": "plan_new_document",
            },
        )
        graph.add_edge("plan_incremental_changes", "generate_doc_content")
        graph.add_edge("plan_new_document", "generate_doc_content")
        graph.add_edge("generate_doc_content", "self_check")
        graph.add_conditional_edges(
            "self_check",
            self._route_after_self_check,
            {"retry": "generate_doc_content", "ready": "build_changeset"},
        )
        graph.add_edge("build_changeset", "summarize_result")
        graph.add_edge("summarize_result", END)
        return graph.compile(checkpointer=checkpointer)

    def _collect_code_context(self, state: AgentGraphState) -> dict:
        return self._run_node(state, "collect_code_context", "Собираю контекст кода и файлов.", self._analyzer.collect_code_context)

    def _detect_existing_docs(self, state: AgentGraphState) -> dict:
        return self._run_node(
            state,
            "detect_existing_docs",
            "Определяю, есть ли существующая документация проекта.",
            self._analyzer.detect_existing_docs,
        )

    def _decide_strategy(self, state: AgentGraphState) -> dict:
        return self._run_node(state, "decide_strategy", "Выбираю стратегию: инкремент или генерация с нуля.", self._analyzer.decide_strategy)

    def _load_rules_and_examples(self, state: AgentGraphState) -> dict:
        return self._run_node(
            state,
            "load_rules_and_examples",
            "Загружаю правила и примеры формата документации.",
            self._composer.load_rules_and_examples,
        )

    def _plan_incremental_changes(self, state: AgentGraphState) -> dict:
        return self._run_node(
            state,
            "plan_incremental_changes",
            "Планирую точечные изменения в существующей документации.",
            lambda st: self._composer.plan_incremental_changes(st, self._analyzer),
        )

    def _plan_new_document(self, state: AgentGraphState) -> dict:
        return self._run_node(state, "plan_new_document", "Проектирую структуру новой документации.", self._composer.plan_new_document)

    def _generate_doc_content(self, state: AgentGraphState) -> dict:
        return self._run_node(state, "generate_doc_content", "Генерирую содержимое документации.", self._composer.generate_doc_content)

    def _self_check(self, state: AgentGraphState) -> dict:
        return self._run_node(state, "self_check", "Проверяю соответствие результата правилам.", self._composer.self_check)

    def _build_changeset(self, state: AgentGraphState) -> dict:
        return self._run_node(state, "build_changeset", "Формирую итоговый набор изменений файлов.", self._composer.build_changeset)

    def _summarize_result(self, state: AgentGraphState) -> dict:
        return self._run_node(
            state,
            "summarize_result",
            "Формирую краткий обзор выполненных действий и измененных файлов.",
            self._composer.build_execution_summary,
        )

    def _route_after_rules_loading(self, state: AgentGraphState) -> str:
        if state.get("docs_strategy") == "incremental_update":
            return "incremental"
        return "from_scratch"

    def _route_after_self_check(self, state: AgentGraphState) -> str:
        if state.get("validation_passed"):
            return "ready"
        attempts = int(state.get("validation_attempts", 0) or 0)
        return "ready" if attempts >= self._max_validation_attempts else "retry"

    def _run_node(self, state: AgentGraphState, node_name: str, message: str, fn):
        emit_progress_sync(state, stage=f"graph.docs.{node_name}", message=message)
        try:
            result = fn(state)
            emit_progress_sync(state, stage=f"graph.docs.{node_name}.done", message=f"Шаг '{node_name}' завершен.")
            LOGGER.warning("docs graph node completed: node=%s keys=%s", node_name, sorted(result.keys()))
            return result
        except Exception:
            LOGGER.exception("docs graph node failed: node=%s", node_name)
            raise
