from typing import NotRequired, TypedDict
 
class WorkflowState(TypedDict):
    # 입력
    prompt: str
    pod_type: str | None
    retry_count: int
 
    # Orchestrator → Retriever, Generator
    search_query: NotRequired[str]
    style_requirements: NotRequired[list[str]]
    generation_instructions: NotRequired[str]
 
    # Retriever → Analyzer
    retrieved_templates: NotRequired[list[dict]]
    reference_xmls: NotRequired[list[str]]
 
    # Analyzer → Generator
    analysis: NotRequired[dict]
 
    # Generator → Validator
    generated_xml: NotRequired[str]
 
    # Validator → (조건부) Generator 또는 END
    validation_passed: NotRequired[bool]
    validation_issues: NotRequired[list[dict]]
    validation_feedback: NotRequired[str]

async def orchestrate_node(state: WorkflowState) -> dict:
    result = await orchestrator.orchestrate(
        state["prompt"], pod_type=state.get("pod_type")
    )
    return {
        "pod_type": result["pod_type"],
        "search_query": result["search_query"],
        "style_requirements": result.get("style_requirements", []),
        "generation_instructions": result["generation_instructions"],
    }
 
async def retrieve_node(state: WorkflowState) -> dict:
    results = await retriever.retrieve(
        state.get("search_query", state["prompt"])
    )
    return {
        "retrieved_templates": results,
        "reference_xmls": [r["xml_content"] for r in results if r.get("xml_content")],
    }
 
async def generate_node(state: WorkflowState) -> dict:
    feedback = (
        state.get("validation_feedback")
        if state.get("retry_count", 0) > 0
        else None
    )
    xml = await generator.generate(
        state.get("generation_instructions", state["prompt"]),
        state.get("analysis", {}),
        state.get("reference_xmls", []),
        feedback=feedback,
    )
    return {"generated_xml": xml}
 
async def validate_node(state: WorkflowState) -> dict:
    result = validator.validate(state.get("generated_xml", ""))
    if result.valid:
        return {"validation_passed": True, "validation_issues": result.issues}
    feedback = "; ".join(issue["message"] for issue in result.issues)
    return {
        "validation_passed": False,
        "validation_issues": result.issues,
        "validation_feedback": feedback,
        "retry_count": state.get("retry_count", 0) + 1,
    }

from langgraph.graph import END, StateGraph
 
def _should_retry(state: WorkflowState) -> str:
    if state.get("validation_passed"):
        return END
    if state.get("retry_count", 0) >= settings.max_retries:
        return END
    return "generate"
 
def build_graph() -> StateGraph:
    graph = StateGraph(WorkflowState)
 
    graph.add_node("orchestrate", orchestrate_node)
    graph.add_node("retrieve", retrieve_node)
    graph.add_node("analyze", analyze_node)
    graph.add_node("generate", generate_node)
    graph.add_node("validate", validate_node)
 
    graph.set_entry_point("orchestrate")
    graph.add_edge("orchestrate", "retrieve")
    graph.add_edge("retrieve", "analyze")
    graph.add_edge("analyze", "generate")
    graph.add_edge("generate", "validate")
    graph.add_conditional_edges("validate", _should_retry)
 
    return graph

compiled = build_graph().compile()
 
async for chunk in compiled.astream(initial_state, stream_mode="updates"):
    for node_name, state_update in chunk.items():
        stage = NODE_TO_STAGE[node_name]
        yield sse_event({"type": "agent_complete", "stage": stage})