"""
agentgraph.py — a minimal, from-scratch re-implementation of the LangGraph core,
so you can SEE what the framework actually does (Knowledge Module 03 §4-5).

A production agent loop fails without four things; this engine provides all four:
  1. SHARED STATE      — a dict that flows through and is updated by each node.
  2. CONDITIONAL EDGES + CYCLES — control flow that branches and loops.
  3. CHECKPOINTER      — persists state after every node => durable, resumable,
                         auditable (resume after a crash; replay for audit).
  4. INTERRUPTS        — pause BEFORE a node (e.g. human approval), return control
                         to the caller, and RESUME later from the saved state.

This is ~120 lines. The real LangGraph adds async, streaming, many checkpoint
backends, sub-graphs, etc. — but the mental model is exactly this.
"""
from __future__ import annotations

import copy
import json
from dataclasses import dataclass, field
from typing import Any, Callable

# A node is a function: (state) -> partial state update (dict, merged into state).
Node = Callable[[dict], dict]
# A router decides the next node name from the current state.
Router = Callable[[dict], str]

START = "__start__"
END = "__end__"


class Interrupt(Exception):
    """Raised by the engine when it reaches an interrupt-before node. The caller
    inspects the checkpoint, gets a human decision, and calls `resume()`."""
    def __init__(self, node: str, state: dict):
        self.node, self.state = node, state
        super().__init__(f"interrupt before '{node}'")


class Checkpointer:
    """In-memory checkpoint store keyed by thread_id. Each thread keeps a LIST of
    snapshots (every node transition) — that history IS the audit trail and
    enables time-travel/replay. A production checkpointer writes to Cosmos DB /
    Postgres / SQLite instead of a dict (see requirements.txt)."""

    def __init__(self) -> None:
        self._threads: dict[str, list[dict]] = {}

    def save(self, thread_id: str, snapshot: dict) -> None:
        self._threads.setdefault(thread_id, []).append(copy.deepcopy(snapshot))

    def latest(self, thread_id: str) -> dict | None:
        hist = self._threads.get(thread_id)
        return copy.deepcopy(hist[-1]) if hist else None

    def history(self, thread_id: str) -> list[dict]:
        return copy.deepcopy(self._threads.get(thread_id, []))

    def serialize(self, thread_id: str) -> str:
        """Prove durability: state survives process death as JSON on disk."""
        return json.dumps(self._threads.get(thread_id, []), default=str)


@dataclass
class StateGraph:
    """Build a graph of nodes + edges, then `.compile()` to a runnable Graph.
    Mirrors LangGraph's StateGraph API (add_node / add_edge /
    add_conditional_edges / compile(checkpointer=..., interrupt_before=...))."""
    nodes: dict[str, Node] = field(default_factory=dict)
    edges: dict[str, str] = field(default_factory=dict)               # node -> next node
    cond_edges: dict[str, tuple[Router, dict[str, str]]] = field(default_factory=dict)
    entry: str = START

    def add_node(self, name: str, fn: Node) -> None:
        self.nodes[name] = fn

    def add_edge(self, src: str, dst: str) -> None:
        if src == START:
            self.entry = dst
        else:
            self.edges[src] = dst

    def add_conditional_edges(self, src: str, router: Router, mapping: dict[str, str]) -> None:
        self.cond_edges[src] = (router, mapping)

    def compile(self, checkpointer: Checkpointer | None = None,
                interrupt_before: list[str] | None = None) -> "Graph":
        return Graph(self, checkpointer or Checkpointer(), set(interrupt_before or []))


class Graph:
    """The runnable graph. `invoke()` starts a thread; `resume()` continues a
    thread that was paused at an interrupt (after a human decision)."""

    def __init__(self, spec: StateGraph, checkpointer: Checkpointer,
                 interrupt_before: set[str]) -> None:
        self.spec = spec
        self.ckpt = checkpointer
        self.interrupt_before = interrupt_before

    def _next(self, node: str, state: dict) -> str:
        if node in self.spec.cond_edges:
            router, mapping = self.spec.cond_edges[node]
            return mapping[router(state)]
        return self.spec.edges.get(node, END)

    def _run_from(self, node: str, state: dict, thread_id: str) -> dict:
        while node != END:
            # Interrupt BEFORE executing a guarded node (e.g. human_review).
            if node in self.interrupt_before and not state.get("_resumed_into") == node:
                state = {**state, "_paused_at": node}
                self.ckpt.save(thread_id, {"node": node, "state": state, "status": "interrupted"})
                raise Interrupt(node, state)
            state.pop("_resumed_into", None)
            update = self.spec.nodes[node](state) or {}
            state = _merge(state, update)
            self.ckpt.save(thread_id, {"node": node, "state": state, "status": "ran"})
            node = self._next(node, state)
        self.ckpt.save(thread_id, {"node": END, "state": state, "status": "done"})
        return state

    def invoke(self, initial_state: dict, thread_id: str) -> dict:
        state = dict(initial_state)
        self.ckpt.save(thread_id, {"node": START, "state": state, "status": "start"})
        return self._run_from(self.spec.entry, state, thread_id)

    def resume(self, thread_id: str, decision: dict) -> dict:
        """Continue a paused thread. `decision` (e.g. {'approved': True}) is merged
        into the saved state; execution restarts AT the interrupted node, this
        time allowed to run through it."""
        snap = self.ckpt.latest(thread_id)
        if not snap or snap["status"] != "interrupted":
            raise ValueError("thread is not paused at an interrupt")
        node = snap["state"]["_paused_at"]
        state = _merge(snap["state"], {**decision, "_resumed_into": node})
        state.pop("_paused_at", None)
        return self._run_from(node, state, thread_id)


def _merge(state: dict, update: dict) -> dict:
    """Merge a node's update into state. Lists under known append-keys (messages,
    audit) are APPENDED (LangGraph calls these 'reducers'); everything else is
    replaced. This is why `messages` and `audit` accumulate across nodes."""
    out = dict(state)
    for k, v in update.items():
        if k in ("messages", "audit") and isinstance(v, list):
            out[k] = out.get(k, []) + v
        else:
            out[k] = v
    return out