"""Orchestrated workflow with classification routing."""
from langgraph.graph import StateGraph, END
from langchain_core.language_models.chat_models import BaseChatModel

from graphs.state import AgentState
from graphs.agents.classifier_agent import classifier_node
from graphs.nodes.retrieval import retrieve_catalogue, retrieve_projects
from graphs.agents.chat_agent import chat_node
from graphs.agents.summarizer_agent import summarizer_llm_node, summarizer_export_node
from tools.pdf import markdown_to_pdf
from tools.storage import upload_pdf_to_supabase


def create_orchestrated_graph(llm: BaseChatModel, checkpointer=None):
    workflow = StateGraph(AgentState)

    # Nodes
    workflow.add_node("classify", classifier_node(llm))
    workflow.add_node("retrieve", retrieve_catalogue)
    # Route to classic vs project retrieval
    def _router_passthrough(state: AgentState) -> AgentState:
        # Must write at least one allowed key; pass through the current query
        q = state.get("query") or ""
        return {"query": q}
    workflow.add_node("retrieve_router", _router_passthrough)
    workflow.add_node("retrieve_project", retrieve_projects)
    workflow.add_node("agent", chat_node(llm))
    workflow.add_node("summarizer_llm", summarizer_llm_node(llm))
    workflow.add_node(
        "summarizer_export",
        summarizer_export_node(
            markdown_to_pdf=markdown_to_pdf,
            upload_pdf=upload_pdf_to_supabase,
        ),
    )

    # Entry
    workflow.set_entry_point("classify")

    # Conditional routing
    workflow.add_conditional_edges(
        "classify",
        lambda s: getattr(s.get("classification"), "classification", "CLASSIC"),
        {
            # Route through a retrieval router to optionally branch to project retrieval
            "CLASSIC": "retrieve_router",
            "SUMMARIZE": "summarizer_llm",
            "UNKNOWN": "retrieve_router",
        },
    )

    # Conditional choice between project vs classic retrieval
    workflow.add_conditional_edges(
        "retrieve_router",
        lambda s: "PROJECT" if s.get("project_id") else "CLASSIC",
        {
            "PROJECT": "retrieve_project",
            "CLASSIC": "retrieve",
        },
    )

    # Linear branches
    # If project path is taken, run project retrieval then classic retrieval
    workflow.add_edge("retrieve_project", "retrieve")
    workflow.add_edge("retrieve", "agent")
    workflow.add_edge("agent", END)
    
    workflow.add_edge("summarizer_llm", "summarizer_export")
    workflow.add_edge("summarizer_export", END)

    return workflow.compile(checkpointer=checkpointer)