from __future__ import annotations

import asyncio
from collections.abc import Callable
from typing import TYPE_CHECKING

from app.modules.agent.engine.graphs.progress_registry import progress_registry
from app.modules.agent.engine.orchestrator.actions import (
    CodeExplainActions,
    DocsActions,
    EditActions,
    ExplainActions,
    GherkinActions,
    ProjectQaActions,
    ReviewActions,
)
from app.modules.agent.engine.orchestrator.execution_context import ExecutionContext
from app.modules.agent.engine.orchestrator.models import ArtifactType, PlanStep

if TYPE_CHECKING:
    from app.modules.rag.explain.retriever_v2 import CodeExplainRetrieverV2

StepFn = Callable[[ExecutionContext], list[str]]


class StepRegistry:
    def __init__(self, code_explain_retriever: CodeExplainRetrieverV2 | None = None) -> None:
        code_explain = CodeExplainActions(code_explain_retriever)
        explain = ExplainActions()
        review = ReviewActions()
        docs = DocsActions()
        edits = EditActions()
        gherkin = GherkinActions()
        project_qa = ProjectQaActions()

        self._functions: dict[str, StepFn] = {
            "collect_state": self._collect_state,
            "finalize_graph_output": self._finalize_graph_output,
            "execute_project_qa_graph": self._collect_state,
            "build_code_explain_pack": code_explain.build_code_explain_pack,
            "collect_sources": explain.collect_sources,
            "extract_logic": explain.extract_logic,
            "summarize": explain.summarize,
            "classify_project_question": project_qa.classify_project_question,
            "collect_project_sources": project_qa.collect_project_sources,
            "analyze_project_sources": project_qa.analyze_project_sources,
            "build_project_answer_brief": project_qa.build_project_answer_brief,
            "compose_project_answer": project_qa.compose_project_answer,
            "fetch_source_doc": review.fetch_source_doc,
            "normalize_document": review.normalize_document,
            "structural_check": review.structural_check,
            "semantic_consistency_check": review.semantic_consistency_check,
            "architecture_fit_check": review.architecture_fit_check,
            "optimization_check": review.optimization_check,
            "compose_review_report": review.compose_review_report,
            "extract_change_intents": docs.extract_change_intents,
            "map_to_doc_tree": docs.map_to_doc_tree,
            "load_current_docs_context": docs.load_current_docs_context,
            "generate_doc_updates": docs.generate_doc_updates,
            "cross_file_validation": docs.cross_file_validation,
            "build_changeset": docs.build_changeset,
            "compose_summary": docs.compose_summary,
            "resolve_target": edits.resolve_target,
            "load_target_context": edits.load_target_context,
            "plan_minimal_patch": edits.plan_minimal_patch,
            "generate_patch": edits.generate_patch,
            "validate_patch_safety": edits.validate_patch_safety,
            "finalize_changeset": edits.finalize_changeset,
            "compose_edit_summary": edits.compose_edit_summary,
            "extract_increment_scope": gherkin.extract_increment_scope,
            "partition_features": gherkin.partition_features,
            "generate_gherkin_bundle": gherkin.generate_gherkin_bundle,
            "lint_gherkin": gherkin.lint_gherkin,
            "validate_coverage": gherkin.validate_coverage,
            "compose_test_model_summary": gherkin.compose_test_model_summary,
        }

    async def execute(self, step: PlanStep, ctx: ExecutionContext) -> list[str]:
        if step.executor == "graph":
            return await self._execute_graph_step(step, ctx)
        fn = self._functions.get(step.action_id)
        if fn is None:
            raise RuntimeError(f"Unknown function action_id: {step.action_id}")
        return fn(ctx)

    def _collect_state(self, ctx: ExecutionContext) -> list[str]:
        state = {
            "task_id": ctx.task.task_id,
            "project_id": ctx.task.rag_session_id,
            "scenario": ctx.task.scenario.value,
            "message": ctx.task.user_message,
            "progress_key": ctx.task.task_id,
            "rag_context": str(ctx.task.metadata.get("rag_context", "")),
            "confluence_context": str(ctx.task.metadata.get("confluence_context", "")),
            "files_map": dict(ctx.task.metadata.get("files_map", {}) or {}),
        }
        item = ctx.artifacts.put(key="agent_state", artifact_type=ArtifactType.STRUCTURED_JSON, content=state)
        return [item.artifact_id]

    async def _execute_graph_step(self, step: PlanStep, ctx: ExecutionContext) -> list[str]:
        graph_key = step.graph_id or "route"
        if graph_key == "route":
            domain_id = ctx.task.routing.domain_id
            process_id = ctx.task.routing.process_id
        elif "/" in graph_key:
            domain_id, process_id = graph_key.split("/", 1)
        else:
            raise RuntimeError(f"Unsupported graph_id: {graph_key}")

        graph = ctx.graph_resolver(domain_id, process_id)
        state = self._build_graph_state(ctx)

        if ctx.progress_cb is not None:
            progress_registry.register(ctx.task.task_id, ctx.progress_cb)
        try:
            result = await asyncio.to_thread(ctx.graph_invoker, graph, state, ctx.task.dialog_session_id)
        finally:
            if ctx.progress_cb is not None:
                progress_registry.unregister(ctx.task.task_id)

        return self._store_graph_outputs(step, ctx, result)

    def _build_graph_state(self, ctx: ExecutionContext) -> dict:
        state = dict(ctx.artifacts.get_content("agent_state", {}) or {})
        for item in ctx.artifacts.all_items():
            state[item.key] = ctx.artifacts.get_content(item.key)
        return state

    def _store_graph_outputs(self, step: PlanStep, ctx: ExecutionContext, result: dict) -> list[str]:
        if not isinstance(result, dict):
            raise RuntimeError("graph_result must be an object")
        if len(step.outputs) == 1 and step.outputs[0].key == "graph_result":
            item = ctx.artifacts.put(key="graph_result", artifact_type=ArtifactType.STRUCTURED_JSON, content=result)
            return [item.artifact_id]

        artifact_ids: list[str] = []
        for output in step.outputs:
            value = result.get(output.key)
            if value is None and output.required:
                raise RuntimeError(f"graph_output_missing:{step.step_id}:{output.key}")
            item = ctx.artifacts.put(key=output.key, artifact_type=output.type, content=value)
            artifact_ids.append(item.artifact_id)
        return artifact_ids

    def _finalize_graph_output(self, ctx: ExecutionContext) -> list[str]:
        raw = ctx.artifacts.get_content("graph_result", {}) or {}
        if not isinstance(raw, dict):
            raise RuntimeError("graph_result must be an object")

        answer = raw.get("answer")
        changeset = raw.get("changeset") or []
        output = [
            ctx.artifacts.put(
                key="final_answer",
                artifact_type=ArtifactType.TEXT,
                content=(str(answer) if answer is not None else ""),
            ).artifact_id,
            ctx.artifacts.put(
                key="final_changeset",
                artifact_type=ArtifactType.CHANGESET,
                content=changeset,
            ).artifact_id,
        ]
        return output
