from langgraph.graph import StateGraph, START, END
from langgraph.types import Send
from alayaflow.clients.alayamem.http_client import HttpAlayaMemClient

from .schemas import (
    WorkflowInitArgs, 
    OverallState, TaskState, WorkflowContext, Input, Output
)
from typing import Dict, Any
from .nodes.memory_nodes import create_retrieve_node, create_memory_retrieve_node
from .nodes.update_memory_node import create_memory_update_node
from .nodes.llm_extract_node import create_generate_node
from .nodes.validate_node import create_validate_node
from .nodes.plan_node import create_plan_node
from .nodes.rewrite_node import create_rewrite_query_node
from .nodes.rerank_node import create_rerank_node


def create_task_graph(client: HttpAlayaMemClient):
    task_g = StateGraph(TaskState)
    task_g.add_node("rewrite", create_rewrite_query_node())
    task_g.add_node("retrieve", create_retrieve_node(client))
    task_g.add_node("rerank", create_rerank_node())
    task_g.add_node("generate", create_generate_node())

    task_g.add_edge(START, "rewrite")
    task_g.add_edge("rewrite", "retrieve")
    task_g.add_edge("retrieve", "rerank")
    task_g.add_edge("rerank", "generate")
    task_g.add_edge("generate", END)

    return task_g.compile()


def fan_out(state: OverallState):
    memory_json = state.get("memory", "{}")
    return [Send("task", {"task": t, "memory_json": memory_json}) for t in state["tasks"]]


def create_graph(init_args: WorkflowInitArgs | Dict[str, Any]):
    if isinstance(init_args, dict):
        init_args = WorkflowInitArgs(**init_args)
    client = HttpAlayaMemClient(init_args.alayamem_url)

    task_flow = create_task_graph(client)

    g = StateGraph(OverallState, WorkflowContext, input_type=Input, output_type=Output)
    g.add_node("plan", create_plan_node())
    g.add_node("memory", create_memory_retrieve_node(client))
    g.add_node("task", task_flow)
    g.add_node("validate", create_validate_node())
    g.add_node("memory_update", create_memory_update_node(client))

    g.add_edge(START, "plan")
    g.add_edge("plan", "memory")
    g.add_conditional_edges("memory", fan_out, ["task"])
    g.add_edge("task", "validate")
    g.add_edge("validate", "memory_update")
    g.add_edge("memory_update", END)

    return g.compile()
