from __future__ import annotations

from fastmcp import FastMCP
from pydantic import BaseModel
from agents import Agent, Runner, WebSearchTool

mcp_web = FastMCP(
    "openai-websearch-server",
    version="0.1.0",
)

websearch_agent = Agent(
    name="Websearch Agent",
    instructions="Search the web and provide a summary (with references) of the results based on the user's query.",
    tools=[WebSearchTool()],
    model="gpt-4.1-mini",  # important: 4.1-mini is 10x cheaper than 4.1
)


class WebSearchArgs(BaseModel):
    query: str


@mcp_web.tool()
async def search_web(args: WebSearchArgs) -> str:
    """Perform a web search using the query, and summarize the results."""
    result = await Runner.run(
        websearch_agent,
        args.query,
    )
    return result.final_output


if __name__ == "__main__":
    import argparse
    import uvicorn

    parser = argparse.ArgumentParser()
    parser.add_argument("--stdio", action="store_true", help="Run on STDIO")
    parser.add_argument("--port", type=int, default=8003, help="Port for SSE")
    opts = parser.parse_args()

    if opts.stdio:
        mcp_web.run(transport="stdio")
    else:
        uvicorn.run(mcp_web.sse_app(), host="127.0.0.1", port=opts.port)