from typing import TypedDict, Optional
from langgraph.errors import GraphInterrupt
from langgraph.graph import StateGraph, START, END
from langgraph.types import interrupt
from quillwright.models import Capture, Observation, LineItem, Estimate, TraceStep
from quillwright.catalog import Catalog
from quillwright.resolver import Model
from quillwright.tools import perceive, lookup_price, compute, draft_line_item
from quillwright.brain_loop import run_brain


class AgentState(TypedDict):
    capture: Capture
    observations: list[Observation]
    line_items: list[LineItem]
    trace: list[TraceStep]
    estimate: Optional[Estimate]


def build_agent(perception_model: Model, catalog: Catalog, checkpointer, brain_model=None):
    def perceive_node(state: AgentState) -> dict:
        obs: list[Observation] = []
        for path in state["capture"].image_paths:
            obs.extend(perceive(path, perception_model))
        trace = state["trace"] + [
            TraceStep(
                action="perceive",
                model=getattr(perception_model, "name", "model"),
                detail=f"found {len(obs)} observation(s)",
            )
        ]
        return {"observations": obs, "trace": trace}

    def deterministic_price(state: AgentState) -> dict:
        items = list(state["line_items"])
        trace = list(state["trace"])
        for ob in state["observations"]:
            res = lookup_price(ob.text, catalog)
            if not res["found"]:
                # Agent Pause: ask the human for a price; resumed value comes back here.
                human_rate = interrupt({"reason": f"No price for '{ob.text}'", "item": ob.text})
                items.append(
                    draft_line_item(
                        ob.text, qty=1, unit="ea", rate=float(human_rate), source="user"
                    )
                )
                trace.append(
                    TraceStep(action="price", detail=f"user priced {ob.text}", status="ok")
                )
                continue
            subtotal = compute(f"1 * {res['rate']}")  # qty defaults to 1 in the core
            items.append(
                draft_line_item(
                    res["description"], qty=1, unit=res["unit"], rate=res["rate"], source="catalog"
                )
            )
            trace.append(
                TraceStep(
                    action="price",
                    model="lookup_price",
                    detail=f"{res['description']} -> {subtotal}",
                )
            )
        return {"line_items": items, "trace": trace}

    def brain_price(state: AgentState) -> dict:
        # The LLM brain decides which items to add; deterministic tools own the numbers.
        obs_text = ", ".join(ob.text for ob in state["observations"])
        priced_extra: list[LineItem] = []
        while True:
            try:
                items, brain_trace, pause = run_brain(
                    brain_model,
                    catalog,
                    observations_text=obs_text,
                    transcript=state["capture"].transcript,
                )
            except GraphInterrupt:
                raise  # an Agent Pause is control flow, not a failure — let it propagate
            except Exception as exc:  # noqa: BLE001 — brain/model failure: degrade, don't crash
                # The LLM brain is unavailable (e.g. Ollama 500). Fall back to the
                # deterministic catalog pricer so the forge still produces an estimate
                # instead of crashing the stream. Facts-from-Tools still holds.
                print(f"[quillwright] brain failed ({exc}); falling back to deterministic pricing.")
                out = deterministic_price(state)
                out["trace"] = out["trace"] + [
                    TraceStep(
                        action="price",
                        model="fallback",
                        detail="Brain unavailable — priced from the catalog directly.",
                        status="ok",
                    )
                ]
                return out
            if pause is None:
                trace = state["trace"] + brain_trace
                return {
                    "line_items": list(state["line_items"]) + priced_extra + items,
                    "trace": trace,
                }
            # Missing price -> ask the human, record it on the catalog, then re-run the brain.
            human_rate = interrupt(
                {"reason": f"No price for '{pause['item']}'", "item": pause["item"]}
            )
            priced_extra.append(
                draft_line_item(
                    pause["item"], qty=1, unit="ea", rate=float(human_rate), source="user"
                )
            )
            catalog.add(pause["item"], pause["item"], "ea", float(human_rate))

    price_node = brain_price if brain_model is not None else deterministic_price

    def assemble_node(state: AgentState) -> dict:
        est = Estimate(
            job_title=state["capture"].trade_hint or "Job",
            line_items=state["line_items"],
            tax_rate=0.13,
        )
        trace = state["trace"] + [TraceStep(action="assemble", detail=f"total {est.total}")]
        return {"estimate": est, "trace": trace}

    g = StateGraph(AgentState)
    g.add_node("perceive", perceive_node)
    g.add_node("price", price_node)
    g.add_node("assemble", assemble_node)
    g.add_edge(START, "perceive")
    g.add_edge("perceive", "price")
    g.add_edge("price", "assemble")
    g.add_edge("assemble", END)
    return g.compile(checkpointer=checkpointer)