"""A complete durable graph with deterministic or configured-provider planning."""

import json
import os
import sqlite3
import sys
from pathlib import Path
from chio_process import ProcessClient
from chio_langgraph import ChioProcessToolNode, ProcessTool, ChioProcessToolError
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import START, END, MessagesState, StateGraph


def build_graph(client, settings, on_receipts_saved):
    tools = ChioProcessToolNode(
        client,
        [
            ProcessTool(
                "publish_report",
                "reports",
                "publish",
                "Publish one bounded report.",
                {
                    "type": "object",
                    "properties": {"report": {"type": "string"}},
                    "required": ["report"],
                    "additionalProperties": False,
                },
            )
        ],
        namespace="report-graph-v1",
    )
    model = None
    if settings.get("provider"):
        from langchain_openai import ChatOpenAI

        # OPENAI_API_KEY is operator environment, never graph state.
        model = ChatOpenAI(
            model=settings["model"], base_url=settings["endpoint"], max_retries=0
        ).bind_tools(tools.model_schemas())

    def plan(state):
        if model is not None:
            return {"messages": [model.invoke(state["messages"])]}
        if any(isinstance(m, ToolMessage) for m in state["messages"]):
            return {"messages": [AIMessage(content="Report published.", id="finished")]}
        return {
            "messages": [
                AIMessage(
                    content="",
                    id="report-plan",
                    tool_calls=[
                        {
                            "name": "publish_report",
                            "id": "publish-report",
                            "args": {"report": "Three inputs total ten."},
                        }
                    ],
                )
            ]
        }

    def execute(state, config):
        result = tools.invoke(state, config)
        for message in result["messages"]:
            with Path(settings["receipts"]).open("a", encoding="utf-8") as output:
                output.write(message.artifact["chio"]["receipt_json"] + "\n")
                output.flush()
                os.fsync(output.fileno())
        on_receipts_saved()
        return result

    graph = StateGraph(MessagesState)
    graph.add_node("model", plan)
    graph.add_node("tools", execute)
    graph.add_edge(START, "model")
    graph.add_conditional_edges(
        "model", lambda state: "tools" if state["messages"][-1].tool_calls else END
    )
    graph.add_edge("tools", "model")
    return graph


def run_worker(bootstrap, on_receipts_saved=lambda: None):
    if bootstrap["schema"] != "chio.process.worker-bootstrap.v1":
        raise RuntimeError("Unsupported bootstrap")
    connection, settings = bootstrap["connection"], bootstrap["input"]
    client = ProcessClient(connection["socket_path"], connection["credential"])
    graph = build_graph(client, settings, on_receipts_saved)
    config = {"configurable": {"thread_id": settings["threadId"]}}
    with sqlite3.connect(settings["graphDb"], check_same_thread=False) as database:
        app = graph.compile(checkpointer=SqliteSaver(database))
        retained = app.get_state(config)
        initial = (
            None
            if retained.values
            else {"messages": [HumanMessage(content=settings["prompt"], id="request")]}
        )
        try:
            result = app.invoke(initial, config, durability="sync")
        except ChioProcessToolError as error:
            if error.receipt_json:
                with Path(settings["receipts"]).open("a") as output:
                    output.write(error.receipt_json + "\n")
            raise
        print(json.dumps({"complete": True, "text": result["messages"][-1].content}))


if __name__ == "__main__":
    run_worker(json.load(sys.stdin))