"""langgraph wiring — the one place the framework appears.

Graph: START → plan → execute → synthesize → END

`execute` iterates Plan steps and dispatches each to its Analysis. Sequential
for now; swap to a fan-out with parallel nodes once we want to parallelise.
"""

from __future__ import annotations

import dataclasses
import logging
from typing import Any

from langgraph.graph import END, START, StateGraph

from api import langfuse_client as lf
from api.analyses.registry import REGISTRY as ANALYSES
from api.analyses.types import Finding
from api.llm import chat
from api.plan.planner import plan as run_planner
from api.prompts import load, render
from api.runtime import events
from api.runtime.context import current_run_id
from api.runtime.state import RunState

logger = logging.getLogger("nvi.runtime")


def _plan_node(state: RunState) -> dict[str, Any]:
    plan_obj = run_planner(state["question"])
    return {
        "plan_rationale": plan_obj.rationale,
        "plan_steps": [s.to_dict() for s in plan_obj.steps],
    }


async def _execute_node(state: RunState) -> dict[str, Any]:
    findings: list[dict[str, Any]] = []
    for step in state.get("plan_steps", []):
        name = step["analysis"]
        analysis = ANALYSES.get(name)
        await events.publish(
            state["run_id"],
            events.analysis_start(name, step.get("args", {}), step.get("why", "")),
        )
        if analysis is None:
            f = Finding(analysis=name, summary="", error=f"unknown analysis {name!r}")
        else:
            try:
                f = await analysis.run(step.get("args", {}), state["question"])
            except Exception as e:
                logger.exception("analysis %s raised", name)
                f = Finding(analysis=name, summary="", error=str(e))
        findings.append(dataclasses.asdict(f))
        await events.publish(state["run_id"], events.analysis_end(name, f))
    return {"findings": findings}


async def _synthesize_node(state: RunState) -> dict[str, Any]:
    await events.publish(state["run_id"], events.synth_start())

    findings = state.get("findings", [])
    if not findings or all(f.get("error") for f in findings):
        return {
            "answer": "I couldn't answer this question — see the per-analysis errors above.",
            "error": "all_analyses_failed",
        }

    findings_block = "\n\n".join(
        f"[{f['analysis']}]\n"
        f"summary: {f.get('summary') or '(none)'}\n"
        f"error: {f.get('error') or '(none)'}\n"
        f"sql: {f.get('sql')}"
        for f in findings
    )

    with lf.span("synthesize", as_type="generation", input={"findings": len(findings)}) as span:
        answer = chat(
            system=load("synthesize.system"),
            user=render(
                "synthesize.user",
                question=state["question"],
                findings_block=findings_block,
            ),
            max_tokens=512,
        )
        span.update(output={"answer": answer})
    return {"answer": answer}


def build_graph():
    g: StateGraph = StateGraph(RunState)
    g.add_node("plan", _plan_node)
    g.add_node("execute", _execute_node)
    g.add_node("synthesize", _synthesize_node)
    g.add_edge(START, "plan")
    g.add_edge("plan", "execute")
    g.add_edge("execute", "synthesize")
    g.add_edge("synthesize", END)
    return g.compile()


GRAPH = build_graph()


async def run(run_id: str, question: str) -> RunState:
    """Drive the graph end-to-end, publishing SSE events along the way."""
    initial: RunState = {"run_id": run_id, "question": question}
    final: RunState = initial
    token = current_run_id.set(run_id)
    try:
        await events.publish(run_id, events.run_start(question))
        async for event in GRAPH.astream(initial, stream_mode="updates"):
            for node, update in event.items():
                await events.publish(run_id, events.node_update(node, update))
                if isinstance(update, dict):
                    final = {**final, **update}
                if node == "plan" and isinstance(update, dict):
                    await events.publish(
                        run_id,
                        events.plan_ready(
                            update.get("plan_rationale", ""),
                            update.get("plan_steps", []),
                        ),
                    )
        await events.publish(
            run_id,
            events.run_end(answer=final.get("answer", ""), error=final.get("error")),
        )
    except Exception as e:
        logger.exception("run %s failed", run_id)
        await events.publish(run_id, events.run_end(error=str(e)))
        final = {**final, "error": str(e)}
    finally:
        current_run_id.reset(token)
        await events.close(run_id)
        lf.flush()
    return final