import operator
from typing import TypedDict, Annotated, List
from langgraph.graph import StateGraph, END

# 1. 定义全链路共享的全局状态契约
class DeepResearchState(TypedDict):
    user_input: str                                     # 原始课题
    is_deep_needed: bool                                # 是否启动深度研究
    research_plan: List[str]                            # 拆解出的子任务清单
    search_results: Annotated[List[str], operator.add]  # 累加各研究员的事实
    final_report: str                                   # 最终输出的 Markdown

# 2. Coordinator 节点：意图路由
def coordinator_node(state: DeepResearchState) -> dict:
    prompt = f"判断以下问题是否需要启动深度网络研究（涉及多方对比、最新动态或复杂行业）：{state['user_input']}"
    # 假设 LLM 返回布尔判定
    return {"is_deep_needed": True}

# 3. Planner 节点：拆解子任务提纲
def planner_node(state: DeepResearchState) -> dict:
    question = state["user_input"]
    prompt = f"""请将以下课题拆解为 3~4 个相互独立的子研究方向，每行一条：
课题：{question}"""
    # 模拟 LLM 拆解出的研究提纲
    sub_tasks = [
        "1. 核心技术路线与技术成熟度现状",
        "2. 全球主要头部厂商的量产时间表与产能规划",
        "3. 成本瓶颈、供应链挑战与未来商业化预测",
    ]
    return {"research_plan": sub_tasks}

# 4. Researcher 节点：多轮检索与事实蒸馏
def researcher_node(state: DeepResearchState) -> dict:
    distilled_findings = []
    for task in state["research_plan"]:
        # 实际逻辑中调用 Tavily 搜索并用 LLM 摘要
        summary = f"【研究方向：{task}】\n- 事实 1: 产业界在特定材料方案取得突破。\n- 事实 2: 预计 2027 年实现规模化装车。"
        distilled_findings.append(summary)
    return {"search_results": distilled_findings}

# 5. Writer 节点：长文结构化归纳
def writer_node(state: DeepResearchState) -> dict:
    context = "\n\n".join(state["search_results"])
    prompt = f"""基于以下各维度的详细事实输入，为课题 '{state['user_input']}' 撰写一份深度专业报告：
{context}

报告结构要求：
1. 报告概述与执行摘要
2. 核心技术与产业现状
3. 关键竞争格局与时间表
4. 综合结论与未来展望"""
    
    # 实际调用强写作模型合成
    generated_report = "# 深度研究报告\n\n## 1. 报告概述\n..."
    return {"final_report": generated_report}

# 6. 拓扑装配
workflow = StateGraph(DeepResearchState)
workflow.add_node("coordinator", coordinator_node)
workflow.add_node("planner", planner_node)
workflow.add_node("researcher", researcher_node)
workflow.add_node("writer", writer_node)

workflow.set_entry_point("coordinator")
workflow.add_edge("coordinator", "planner")
workflow.add_edge("planner", "researcher")
workflow.add_edge("researcher", "writer")
workflow.add_edge("writer", END)

app = workflow.compile()

from src.tools.base_search import BaseSearchTool

class EnterpriseKBSearchTool(BaseSearchTool):
    """自定义接入企业内网知识库"""
    
    def search(self, query: str, max_results: int = 5) -> list[dict]:
        # 访问内网向量库或内部搜索引擎
        resp = internal_search_client.query(query, limit=max_results)
        return [
            {
                "title": doc.title,
                "content": doc.summary,
                "url": f"https://wiki.corp.internal/{doc.id}",
            }
            for doc in resp.items
        ]