import logging

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver

from src.workflow.state import GovBridgeState
from src.agents.input_agent import input_agent
from src.agents.ideas_agent import ideas_agent
from src.agents.complaint_agent import complaint_agent
from src.agents.email_writer_agent import email_writer_agent

_log = logging.getLogger(__name__)


def _route_after_input(state: GovBridgeState) -> str:
    category = state.get("category")
    if category == "idea":
        _log.info("routing → ideas_agent")
        return "ideas_agent"
    if category == "complaint":
        _log.info("routing → complaint_agent")
        return "complaint_agent"
    _log.info("routing → input_agent (no category yet)")
    return "input_agent"


def _route_after_email_writer(state: GovBridgeState) -> str:
    if state.get("draft_approved"):
        _log.info("draft approved — routing to END")
        return END
    _log.info("draft not yet approved — looping back to email_writer")
    return "email_writer"


def _route_after_conversation(node_name: str):
    def _route(state: GovBridgeState) -> str:
        if state.get("conversation_complete"):
            _log.info("%s conversation complete — routing to email_writer", node_name)
            return "email_writer"
        _log.info("%s conversation continuing — looping back", node_name)
        return node_name
    return _route


def build_graph():
    graph = StateGraph(GovBridgeState)

    graph.add_node("input_agent", input_agent)
    graph.add_node("ideas_agent", ideas_agent)
    graph.add_node("complaint_agent", complaint_agent)
    graph.add_node("email_writer", email_writer_agent)

    graph.set_entry_point("input_agent")
    graph.add_conditional_edges("input_agent", _route_after_input, {
        "ideas_agent": "ideas_agent",
        "complaint_agent": "complaint_agent",
        "input_agent": "input_agent",
    })
    graph.add_conditional_edges("ideas_agent", _route_after_conversation("ideas_agent"), {
        "email_writer": "email_writer",
        "ideas_agent": "ideas_agent",
    })
    graph.add_conditional_edges("complaint_agent", _route_after_conversation("complaint_agent"), {
        "email_writer": "email_writer",
        "complaint_agent": "complaint_agent",
    })
    graph.add_conditional_edges("email_writer", _route_after_email_writer, {
        END: END,
        "email_writer": "email_writer",
    })

    return graph.compile(checkpointer=MemorySaver())