"""LangGraph definition for Gumbo (load from MCP FS -> summarize via LiteLLM)."""

from __future__ import annotations

import os
from typing import TypedDict

from langchain_core.messages import HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph
# AsyncPostgresSaver imported lazily in run_gumbo() to avoid startup crash

from gumbo.mcp_fs import fetch_text_via_mcp


class GumboState(TypedDict):
    object_key: str
    source_text: str
    summary: str


async def _load(state: GumboState) -> dict:
    text = await fetch_text_via_mcp(state["object_key"])
    return {"source_text": text}


async def _summarize(state: GumboState) -> dict:
    llm = ChatOpenAI(
        model=os.environ.get("GUMBO_LLM_MODEL", "ollama-qwen"),
        api_key=os.environ["LITELLM_API_KEY"],
        base_url=os.environ["LITELLM_BASE_URL"],
        temperature=0.2,
    )
    system = SystemMessage(
        content="You are Gumbo, a precise documentation summarizer. Return a concise markdown summary."
    )
    user = HumanMessage(
        content=f"Summarize the following document:\n\n{state['source_text']}",
    )
    resp = await llm.ainvoke([system, user])
    summary = resp.content if isinstance(resp.content, str) else str(resp.content)
    return {"summary": summary}


def build_graph_builder() -> StateGraph:
    builder = StateGraph(GumboState)
    builder.add_node("load", _load)
    builder.add_node("summarize", _summarize)
    builder.set_entry_point("load")
    builder.add_edge("load", "summarize")
    builder.add_edge("summarize", END)
    return builder


async def run_gumbo(object_key: str, thread_id: str, conn_string: str) -> GumboState:
    try:
        from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
    except ImportError:
        from langgraph_checkpoint_postgres.aio import AsyncPostgresSaver
    builder = build_graph_builder()
    async with AsyncPostgresSaver.from_conn_string(conn_string) as checkpointer:
        await checkpointer.setup()
        graph = builder.compile(checkpointer=checkpointer)
        return await graph.ainvoke(
            {"object_key": object_key, "source_text": "", "summary": ""},
            config={"configurable": {"thread_id": thread_id}},
        )
