from langgraph.graph import START, END, StateGraph, MessagesState
from typing import TypedDict, Annotated
from operator import add
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
import os


load_dotenv(override=True)
deepseek_base_url = os.getenv("DEEPSEEK_BASE_URL")
deepseek_base_key = os.getenv("DEEPSEEK_BASE_KEY")
deepseek_base_name = os.getenv("DEEPSEEK_BASE_NAME")

print("1. 三点===============================================================================================================================================================================================================")
class CounterState1(TypedDict):
    count: int

def incr_node_1(state: CounterState1)-> dict:
    count = state["count"] + 1
    return {"count": count}

def incr_node_2(state: CounterState1)-> dict:
    count = state["count"] * 2
    return {"count": count}

build = StateGraph(CounterState1)
build.add_node("incr_node_1", incr_node_1)
build.add_node("incr_node_2", incr_node_2)

build.add_edge(START, "incr_node_1")
build.add_edge("incr_node_1", "incr_node_2")
build.add_edge("incr_node_2", END)

graph = build.compile()
result = graph.invoke({"count": 1})
print(result)


print("2. 条件边===============================================================================================================================================================================================================")
class CounterState2(TypedDict):
    count: int

def sub_node_1(state: CounterState2) -> dict:
    count = state["count"] - 1
    return {"count": count}

def sub_node_2(state: CounterState2) -> dict:
    count = state["count"] - 2
    return {"count": count}

def condition_edge(state: CounterState2) -> str:
    if state["count"] > 5:
        return "sub_node_1"
    return "sub_node_2"

build = StateGraph(CounterState2)

build.add_node("sub_node_1", sub_node_1).add_node("sub_node_2", sub_node_2)

build.add_edge(START, "sub_node_1")
build.add_conditional_edges(
    "sub_node_1",
    condition_edge,
    {"sub_node_1": "sub_node_1", "sub_node_2": "sub_node_2"}
)
build.add_edge("sub_node_2", END)

graph = build.compile()
result = graph.invoke({"count": 20})
print(result)

print("3. 覆盖===============================================================================================================================================================================================================")
class OverwriteState(TypedDict):
    arr:  list[str]

def node_1(state: OverwriteState)->dict:
    return {"arr":["a", "b"]}

def node_2(state: OverwriteState)->dict:
    return {"arr":["c", "d"]}

build = StateGraph(OverwriteState)
build.add_node("node_1", node_1).add_node("node_2", node_2)
build.add_edge(START, "node_1")
build.add_edge("node_1", "node_2")
build.add_edge("node_2", END)

graph = build.compile()
result = graph.invoke({"arr":["e", "f"]})
print(result)

print("4. 追加===============================================================================================================================================================================================================")
class AppendState(TypedDict):
    arr: Annotated[list[str], add]

def node_1(state: AppendState)->dict:
    return {"arr":["aa", "bb"]}

def node_2(state: AppendState)->dict:
    return {"arr":["cc", "dd"]}

build = StateGraph(AppendState)
build.add_node("node_1", node_1).add_node("node_2", node_2)
build.add_edge(START, "node_1")
build.add_edge("node_1", "node_2")
build.add_edge("node_2", END)

graph = build.compile()
result = graph.invoke({"arr":["ee", "ff"]})
print(result)

print("5. API===============================================================================================================================================================================================================")
# llm = ChatOpenAI(
#     base_url= deepseek_base_url,
#     api_key= deepseek_base_key,
#     model= deepseek_base_name
# )

# def chat_llm(state: MessagesState)->dict:
#     response = llm.invoke("你是谁")
#     return {"messages": [response]}

# build = StateGraph(MessagesState)
# build.add_node("chat", chat_llm)
# build.add_edge(START, "chat")
# build.add_edge("chat", END)

# graph = build.compile()
# result = graph.invoke({"messages": ["多学多问"]})
# print(result)

print("6. 自定义===============================================================================================================================================================================================================")

class ReplaceableList(list):
    def __init__(self, *args, replace: bool = False, **kwargs):
        super().__init__(*args, **kwargs)
        self.replace = replace

def smart_merge(current: list, incom: list)->list:
    if isinstance(incom, ReplaceableList) and getattr(incom, "replace", False):
        return list(incom)
    return current + incom

class CostomState(TypedDict):
    arr: Annotated[list[str], smart_merge]

def append(state: CostomState)->dict:
    return {"arr": ["b", "c"]}
                 
def overwrite(state: CostomState)->dict:
    new_list = ReplaceableList(["reset_data"])
    new_list.replace = True
    return {"arr": new_list}

build = StateGraph(CostomState)
build.add_node("append", append)
build.add_node("overwrite", overwrite)
build.add_edge(START, "overwrite")
build.add_edge("overwrite", "append")
build.add_edge("append", END)

graph = build.compile()
result = graph.invoke({"arr": ["a"]})
print(result)


print("7. 子图===============================================================================================================================================================================================================")

class DrinkState(TypedDict):
    logs: Annotated[list[str], add]

class MakeState(TypedDict):
    sub_logs: Annotated[list[str], add]

def add_coffee(state: MakeState) -> dict:
    return {"sub_logs": ["添加咖啡"]}

def add_sugar(state: MakeState) -> dict:
    return {"sub_logs": ["添加粮"]}

def do_make_coffee(state: MakeState) -> dict:
    return {"sub_logs": ["咖啡制作中..."]}

makeBuild = StateGraph(MakeState)
makeBuild.add_node("add_coffee", add_coffee)
makeBuild.add_node("add_sugar", add_sugar)
makeBuild.add_node("do_make_coffee", do_make_coffee)
makeBuild.add_edge(START, "add_coffee")
makeBuild.add_edge("add_coffee", "add_sugar")
makeBuild.add_edge("add_sugar", "do_make_coffee")
makeBuild.add_edge("do_make_coffee", END)
makeGraph = makeBuild.compile()

def order_coffee(state: DrinkState) -> dict:
    return {"logs": ["点咖啡"]}

def make_coffee(state: DrinkState) -> dict:
    result = makeGraph.invoke({"sub_logs": ["开始制作咖啡"]})
    return {"logs": result["sub_logs"]}

def drink_coffee(state: DrinkState) -> dict:
    return {"logs": ["喝咖啡"]}

build = StateGraph(DrinkState)
build.add_node("order_coffee", order_coffee)
build.add_node("make_coffee", make_coffee)
build.add_node("drink_coffee", drink_coffee)
build.add_edge(START, "order_coffee")
build.add_edge("order_coffee", "make_coffee")
build.add_edge("make_coffee", "drink_coffee")
build.add_edge("drink_coffee", END)

graph = build.compile()
result = graph.invoke({"logs": ["喝咖啡"]})
print(result)


print("8. Send===============================================================================================================================================================================================================")
"""
图的拓扑结构：

            START
              │
              ▼
        [parse_query]            ← 节点：解析用户输入，提取品牌列表
              │
              ▼
        route_to_research        ← 条件边路由函数：返回 N 个 Send
              │
              ├─ Send("research_brand", {brand: "苹果"})
              ├─ Send("research_brand", {brand: "华为"})    ← 三个并行
              └─ Send("research_brand", {brand: "小米"})
              │
              ▼  (并行执行)
        [research_brand] × 3     ← 每个品牌独立研究
              │
              ▼  (全部完成后汇总)
        [summarize]              ← 节点：汇总所有报告，输出对比
              │
              ▼
             END
"""
from typing import Annotated, TypedDict
from operator import add
from langgraph.graph import StateGraph, START, END
from langgraph.types import Send


# ══════════════════════════════════════════════════════════════
#  State 定义
# ══════════════════════════════════════════════════════════════
class CompareState(TypedDict):
    """手机对比任务的 State"""
    query: str                                # 用户的原始问题
    brands: list[str]                         # 从 query 中解析出的品牌列表
    reports: Annotated[list[str], add]        # 每个品牌的研究报告（add reducer 自动合并）
    comparison: str                           # 最终的对比总结


# ══════════════════════════════════════════════════════════════
#  节点函数
# ══════════════════════════════════════════════════════════════
def parse_query(state: CompareState) -> dict:
    """
    节点 ①：解析用户输入，提取出要对比的品牌列表。

    实际场景中这里会用 LLM 做实体提取，这里直接硬编码模拟。
    """
    query = state["query"]
    print(f"[parse] 收到问题: {query}")

    # 模拟 LLM 解析：从自然语言中提取品牌名称
    brands = ["苹果", "华为", "小米"]
    print(f"[parse] 解析出 {len(brands)} 个品牌: {brands}")

    return {"brands": brands}


def research_brand(state: CompareState) -> dict:
    """
    节点 ②：研究单个品牌（被 Send 并行调用的目标节点）。

    它不知道总共有几个品牌，也不关心其他品牌在做什么——
    只专注于自己拿到的这一个。

    实际场景中这里会调搜索引擎 API 或 LLM，这里用数据模拟。
    """
    brand = state["brands"][0]  # Send 保证这里只有一个品牌

    # 模拟：查询该品牌的旗舰机信息
    phone_db = {
        "苹果": "iPhone 16 Pro Max | A18 Pro 芯片 | 6.9吋 OLED | 4K 120fps 视频 | iOS 18",
        "华为": "Mate 70 Pro+ | 麒麟 9100 | 6.8吋 OLED | 物理可变光圈 | 鸿蒙 NEXT",
        "小米": "小米 15 Ultra | 骁龙 8 Gen4 | 6.73吋 AMOLED | 徕卡光学 | 澎湃 OS 2.0",
    }

    info = phone_db.get(brand, f"{brand}旗舰机信息暂缺")
    report = f"【{brand}】{info}"

    print(f"  [research] 并行查询 [{brand}] -> {info}")
    return {"reports": [report]}  # add reducer 自动拼到总 reports 里


def summarize(state: CompareState) -> dict:
    """
    节点 ③：所有品牌都研究完了，汇总生成对比结论。

    实际场景中这里会把所有 reports 喂给 LLM，让它生成对比分析。
    """
    print(f"\n[summarize] 开始汇总 {len(state['reports'])} 份报告...")

    # 模拟 LLM 生成的对比总结
    comparison = (
        "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n"
        "         >> 旗舰手机对比总结\n"
        "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n\n"
    )
    for report in state["reports"]:
        comparison += f"  {report}\n"

    comparison += (
        "\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n"
        "[总结] 三款旗舰各有千秋 --\n"
        "  苹果 iPhone 16 Pro Max：视频拍摄王者，生态闭环体验最佳\n"
        "  华为 Mate 70 Pro+：影像系统物理可变光圈独树一帜\n"
        "  小米 15 Ultra：徕卡光学加持，性价比旗舰首选\n"
        "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n"
    )

    return {"comparison": comparison}


# ══════════════════════════════════════════════════════════════
#  路由函数 —— 核心！返回 Send 列表而不是 dict
# ══════════════════════════════════════════════════════════════
def route_to_research(state: CompareState):
    """
    ★ Send API 的核心：条件边路由函数 ★

    关键规则：
    - 这个函数只能放在 add_conditional_edges 里，不能放进 add_node
    - 普通节点返回 dict（更新 state）
    - 路由函数返回 Send 列表（创建并行任务）或字符串（END / 节点名）

    思想：几个品牌 → 发几个 Send → 启动几次 research_brand

            return [
                Send("research_brand", {"brands": ["苹果"]}),   ─┐
                Send("research_brand", {"brands": ["华为"]}),   ─┤ 全部并行执行
                Send("research_brand", {"brands": ["小米"]}),   ─┘
            ]

    每个 Send 的两个参数：
      参数 1: 目标节点名 —— "派给谁做"
      参数 2: 此分支专属的 state —— "这个任务要什么数据"
    """
    if not state["brands"]:
        return END

    return [
        Send("research_brand", {"brands": [brand]})
        for brand in state["brands"]
    ]


# ══════════════════════════════════════════════════════════════
#  构图
# ══════════════════════════════════════════════════════════════
builder = StateGraph(CompareState)

builder.add_node("parse_query", parse_query)
builder.add_node("research_brand", research_brand)
builder.add_node("summarize", summarize)

# 边
builder.add_edge(START, "parse_query")
builder.add_conditional_edges("parse_query", route_to_research)  # 路由 → 动态 fan-out
builder.add_edge("research_brand", "summarize")                  # 所有并行结果汇聚到 summarize
builder.add_edge("summarize", END)

graph = builder.compile()


# ══════════════════════════════════════════════════════════════
#  运行
# ══════════════════════════════════════════════════════════════

result = graph.invoke({
    "query": "帮我对比一下苹果、华为和小米的旗舰机型",
    "brands": [],
    "reports": [],
    "comparison": "",
})
print("======")
print(result["brands"])
print(result["comparison"])