import os
from typing import Literal

from google.oauth2 import service_account
from langchain_core.messages import AIMessage
from langchain_google_genai import ChatGoogleGenerativeAI
from langgraph.graph import StateGraph, END, START
from pydantic import BaseModel
from pymongo import MongoClient

from src.core.config import settings, MONGO_URI
from src.server.agent.state import AgentState
from src.server.agent.agents import designer_node, researcher_node, visualizer_node, suggester_node
from langgraph.checkpoint.mongodb import MongoDBSaver


# --- Supervisor (路由逻辑) ---
# 定义路由的输出结构，强制 LLM 选择一个
class RouteResponse(BaseModel):
    # 将 FINISH 替换或增加 Suggester
    next: Literal["Designer", "Researcher", "Visualizer", "Suggester"]


creds = service_account.Credentials.from_service_account_file(
    settings.GOOGLE_GENAI_USE_VERTEXAI,
    scopes=["https://www.googleapis.com/auth/cloud-platform"],
)

llm_supervisor = ChatGoogleGenerativeAI(
    model="gemini-2.0-flash", temperature=0, credentials=creds,
    project="aida-461108", location='us-central1', vertexai=True, api_key=settings.GOOGLE_API_KEY
)


def supervisor_node(state: AgentState):
    messages = state["messages"]
    if not messages:
        return {"next": "Suggester"}

    last_message = messages[-1]

    # --- 拦截逻辑修改 ---
    # 如果专家已经回复完了（AIMessage 且无工具调用），则交给 Suggester 生成按钮
    if isinstance(last_message, AIMessage) and not last_message.tool_calls:
        return {"next": "Suggester"}

    system_prompt = """你是家具设计主管。分配任务给专家：
    - Designer: 设计建议、参数细化。
    - Visualizer: 绘图需求。
    - Researcher: 市场报告。
    """

    chain = llm_supervisor.with_structured_output(RouteResponse)
    decision = chain.invoke([{"role": "system", "content": system_prompt}] + messages)
    return {"next": decision.next}


# --- 构建 Graph ---
workflow = StateGraph(AgentState)

workflow.add_node("Supervisor", supervisor_node)
workflow.add_node("Designer", designer_node)
workflow.add_node("Researcher", researcher_node)
workflow.add_node("Visualizer", visualizer_node)
workflow.add_node("Suggester", suggester_node)  # 新增节点

workflow.add_edge(START, "Supervisor")

# 修改条件边映射
workflow.add_conditional_edges(
    "Supervisor",
    lambda state: state["next"],
    {
        "Designer": "Designer",
        "Researcher": "Researcher",
        "Visualizer": "Visualizer",
        "Suggester": "Suggester"  # 原本的 FINISH 现在指向 Suggester
    }
)

# 专家执行完依然回到 Supervisor
workflow.add_edge("Designer", "Supervisor")
workflow.add_edge("Researcher", "Supervisor")
workflow.add_edge("Visualizer", "Supervisor")

# 重点：Suggester 是整个流程的终点
workflow.add_edge("Suggester", END)

client = MongoClient(MONGO_URI)
checkpointer = MongoDBSaver(
    client=client["furniture_agent_db"],
    db_name="langgraph",
    collection_name="checkpoints"
)
app = workflow.compile(checkpointer=checkpointer)