"""
graph.py — LangGraph 流程編排核心
"""
import operator
import uuid
from typing import TypedDict, Annotated, Optional, Literal
from langgraph.graph import StateGraph, END
from nodes import (
    classify_intent,
    call_crew_strategy,
    call_crew_prd,
    call_crew_marketing,
    call_crew_research,
    call_code_generation,   # 替換原本的 openhands import
    wait_for_job,
    format_response,
    handle_error,
)


class CompanyState(TypedDict):
    user_input: str
    session_id: str
    intent: Optional[str]
    department: Optional[str]
    job_id: Optional[str]
    job_status: Optional[str]
    retry_count: int
    crew_result: Optional[str]
    openhands_result: Optional[str]
    final_response: str
    error: Optional[str]
    messages: Annotated[list[dict], operator.add]


def check_job_status(state: CompanyState) -> Literal["wait", "done", "error"]:
    status = state.get("job_status", "queued")
    if status == "done":
        return "done"
    elif status == "error":
        return "error"
    else:
        return "wait"


def build_company_graph() -> StateGraph:
    workflow = StateGraph(CompanyState)

    workflow.add_node("classify", classify_intent)
    workflow.add_node("crew_strategy", call_crew_strategy)
    workflow.add_node("crew_prd", call_crew_prd)
    workflow.add_node("crew_marketing", call_crew_marketing)
    workflow.add_node("crew_research", call_crew_research)
    workflow.add_node("code_generation", call_code_generation)
    workflow.add_node("wait_job", wait_for_job)
    workflow.add_node("format", format_response)
    workflow.add_node("error_handler", handle_error)

    workflow.set_entry_point("classify")

    workflow.add_conditional_edges(
        "classify",
        lambda s: s.get("intent", "chat"),
        {
            "strategy":  "crew_strategy",
            "prd":       "crew_prd",
            "marketing": "crew_marketing",
            "research":  "crew_research",
            "code":      "code_generation",
            "review":    "code_generation",
            "sales":     "crew_strategy",
            "chat":      "format",
        },
    )

    for node in ["crew_strategy", "crew_prd", "crew_marketing", "crew_research"]:
        workflow.add_edge(node, "wait_job")

    workflow.add_edge("code_generation", "format")

    workflow.add_conditional_edges(
        "wait_job",
        check_job_status,
        {
            "wait":  "wait_job",
            "done":  "format",
            "error": "error_handler",
        },
    )

    workflow.add_edge("format", END)
    workflow.add_edge("error_handler", END)

    return workflow.compile()


company_graph = build_company_graph()
company_graph = company_graph.with_config({"recursion_limit": 120})