import asyncio
import json
import logging
import os
import signal
import sys

import httpx

from mcp.server import Server
from mcp.server.sse import SseServerTransport
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse, Response
from starlette.routing import Mount, Route
from mcp.types import Tool, TextContent

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    stream=sys.stderr,
)
logger = logging.getLogger("duckduckgo-mcp")

app = Server("duckduckgo-search")
PORT = int(os.environ.get("MCP_PORT", "3002"))
SEARXNG_URL = os.environ.get("SEARXNG_URL", "http://searxng:8080")
RATE_LIMIT_SECONDS = float(os.environ.get("RATE_LIMIT_SECONDS", "3"))


class RateLimiter:
    def __init__(self, min_interval: float):
        self.min_interval = min_interval
        self._last_request = None
        self._lock = asyncio.Lock()

    async def acquire(self):
        async with self._lock:
            now = asyncio.get_event_loop().time()
            if self._last_request:
                elapsed = now - self._last_request
                wait = self.min_interval - elapsed
                if wait > 0:
                    logger.info(f"Rate limit: waiting {wait:.1f}s")
                    await asyncio.sleep(wait)
            self._last_request = asyncio.get_event_loop().time()


rate_limiter = RateLimiter(RATE_LIMIT_SECONDS)


async def do_search(query: str, num_results: int, engine: str):
    """Shared search logic."""
    await rate_limiter.acquire()

    async with httpx.AsyncClient(timeout=15) as client:
        resp = await client.get(
            f"{SEARXNG_URL}/search",
            params={
                "q": query,
                "format": "json",
                "engines": engine,
                "categories": "general",
                "language": "en",
            },
        )
        resp.raise_for_status()
        data = resp.json()

    results = data.get("results", [])[:num_results]
    return [
        {
            "title": r.get("title", ""),
            "url": r.get("url", ""),
            "snippet": r.get("content", "")[:200],
        }
        for r in results
    ]


# --- HTTP endpoint ---

async def search_http(request: Request):
    query = request.query_params.get("q", "")
    if not query:
        return JSONResponse({"error": "'q' parameter required"}, status_code=400)

    num = min(int(request.query_params.get("num", "10")), 20)
    try:
        results = await do_search(query, num, "duckduckgo")
        return JSONResponse({"query": query, "results": results})
    except Exception as e:
        logger.exception("Search failed")
        return JSONResponse({"error": str(e)}, status_code=500)


# --- MCP tool handlers ---

@app.list_tools()
async def list_tools() -> list[Tool]:
    return [
        Tool(
            name="duckduckgo_search",
            description="Search DuckDuckGo and return results with titles, URLs, and snippets.",
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "The search query."},
                    "num_results": {
                        "type": "integer",
                        "description": "Maximum number of results (default 10).",
                        "default": 10,
                    },
                },
                "required": ["query"],
            },
        )
    ]


@app.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    if name != "duckduckgo_search":
        raise ValueError(f"Unknown tool: {name}")

    query = arguments.get("query", "")
    if not query:
        return [TextContent(type="text", text="Error: 'query' is required.")]

    num_results = min(int(arguments.get("num_results", 10)), 20)

    try:
        results = await do_search(query, num_results, "duckduckgo")
    except Exception as e:
        return [TextContent(type="text", text=f"Search failed: {e}")]

    if not results:
        return [TextContent(type="text", text=f"No results found for: {query}")]

    lines = [f"Search results for: {query}\n"]
    for i, r in enumerate(results, 1):
        lines.append(f"{i}. {r['title']}")
        lines.append(f"   URL: {r['url']}")
        if r["snippet"]:
            lines.append(f"   {r['snippet']}")
        lines.append("")

    return [TextContent(type="text", text="\n".join(lines))]


# --- SSE MCP transport ---

sse = SseServerTransport("/messages/")


async def handle_sse(request):
    async with sse.connect_sse(
        request.scope, request.receive, request._send
    ) as (read_stream, write_stream):
        await app.run(
            read_stream,
            write_stream,
            app.create_initialization_options(),
        )
    return Response()


starlette_app = Starlette(
    debug=True,
    routes=[
        Route("/search", endpoint=search_http),
        Route("/sse", endpoint=handle_sse),
        Mount("/messages/", app=sse.handle_post_message),
    ],
)


async def main():
    logger.info(f"DuckDuckGo MCP server ready on port {PORT}")

    import uvicorn

    config = uvicorn.Config(starlette_app, host="0.0.0.0", port=PORT, log_level="info")
    server = uvicorn.Server(config)

    loop = asyncio.get_running_loop()

    def handle_signal():
        server.should_exit = True

    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, handle_signal)

    await server.serve()


if __name__ == "__main__":
    asyncio.run(main())