import asyncio
import argparse
from typing import List
from mcp.server.fastmcp import FastMCP
from openai import AsyncOpenAI
from pydantic import BaseModel
from agents import (
    OpenAIChatCompletionsModel,
    Agent,
    Runner,
    set_default_openai_client,
    set_tracing_disabled,
    WebSearchTool,
)
from agents.model_settings import ModelSettings

# 初始化 MCP 服务器
mcp = FastMCP("DeepResearch")
USER_AGENT = "deepresearch-app/1.0"
API_KEY = None

# 创建planner_agent
PROMPT = (
    "You are a helpful research assistant. Given a query, come up with a set of web searches "
    "to perform to best answer the query. Output between 10 and 20 terms to query "
    "for. Output the terms in a comma-separated list."
)


class WebSearchItem(BaseModel):
    reason: str
    """Your reasoning for why this search is important to the query."""
    query: str
    """The search term to use for the web search."""


class WebSearchPlan(BaseModel):
    searches: list[WebSearchItem]
    """A list of web searches to perform to best answer the query."""


planner_agent = Agent(
    name="PlannerAgent",
    instructions=PROMPT,
    model="o3-mini",
    output_type=WebSearchPlan,
)

# 创建search_agent
INSTRUCTIONS = (
    "You are a research assistant. Given a search term, you search the web for "
    "that term and "
    "produce a concise summary of the results. The summary must 2-3 paragraphs "
    "and less than 300 "
    "words. Capture the main points. Write succinctly, no need to have complete "
    "sentences or good "
    "grammar. This will be consumed by someone synthesizing a report, so its "
    "vital you capture the "
    "essence and ignore any fluff. Do not include any additional commentary other "
    "than the summary "
    "itself."
)

search_agent = Agent(
    name="Search agent",
    instructions=INSTRUCTIONS,
    tools=[WebSearchTool()],
    model_settings=ModelSettings(tool_choice="required"),
)

# 创建writer_agent
PROMPT_TEMP = (
    "You are a senior researcher tasked with writing a cohesive report for a "
    "research query. "
    "You will be provided with the original query, and some initial research done "
    "by a research "
    "assistant.\n"
    "You should first come up with an outline for the report that describes the "
    "structure and "
    "flow of the report. Then, generate the report and return that as your final "
    "output.\n"
    "The final output should be in markdown format, and it should be lengthy and "
    "detailed. Aim "
    "for 10-20 pages of content, at least 1500 words."
    "最终结果请用中文输出。"
)


class ReportData(BaseModel):
    short_summary: str
    """A short 2-3 sentence summary of the findings."""
    markdown_report: str
    """The final report"""
    follow_up_questions: list[str]
    """Suggested topics to research further"""


writer_agent = Agent(
    name="WriterAgent",
    instructions=PROMPT_TEMP,
    model="o3-mini",
    output_type=ReportData,
)


# 辅助函数组
async def _plan_searches(query: str) -> WebSearchPlan:
    """
    用于进行某个搜索主题的搜索规划
    """
    result = await Runner.run(
        planner_agent,
        f"Query: {query}",
    )
    return result.final_output_as(WebSearchPlan)


async def _perform_searches(search_plan: WebSearchPlan) -> List[str]:
    """
    用于实际执行搜索，并组成搜索条目列表
    """
    tasks = [asyncio.create_task(_search(item)) for item in search_plan.searches]
    results = []
    for task in asyncio.as_completed(tasks):
        result = await task
        if result is not None:
            results.append(result)
    return results


async def _search(item: WebSearchItem) -> str | None:
    """
    实际负责进行搜索，并完成每个搜索条目的短文本编写
    """
    try:
        result = await Runner.run(
            search_agent,
            input=f"Search term: {item.query}\nReason for searching: {item.reason}",
        )
        return str(result.final_output)
    except Exception:
        return None


async def _write_report(query: str, search_results: List[str]) -> ReportData:
    """
    根据搜索的段文档，编写长文本
    """
    result = await Runner.run(
        writer_agent,
        input=f"Original query: {query}\nSummarized search results: {search_results}",
    )
    return result.final_output_as(ReportData)


# MCP服务器主函数
@mcp.tool()
async def deepresearch(query: str) -> ReportData:
    """
    当用户明确表示需要围绕某个主题进行深入研究时，请调本函数。
    本函数能够围绕用户输入的问题进行联网搜索和深入研究，并创建一篇内容完整的markdown格式的研究
    报告。
    输入参数query:用户提出的研究主题，以字符串形式表示；
    函数返回结果为一个markdown格式的完整的研究报告文档。
    """
    search_plan = await _plan_searches(query)
    search_results = await _perform_searches(search_plan)
    report = await _write_report(query, search_results)
    return report


def main():
    parser = argparse.ArgumentParser(description="DeepResearch Server")
    parser.add_argument(
        "--openai_api_key", type=str, required=True, help="你的OpenAI API Key"
    )
    args = parser.parse_args()

    # 初始化 external_client 和设置
    external_client = AsyncOpenAI(
        base_url="https://ai.devtool.tech/proxy/v1",
        api_key=args.openai_api_key,
    )
    set_default_openai_client(external_client)
    set_tracing_disabled(True)
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
