"""MCP server exposing ihyee web intelligence API as tools."""
from __future__ import annotations

import asyncio
import os
import sys
from typing import Any, Sequence

from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import TextContent, Tool

from ihyee import AsyncIhyee
from ihyee.exceptions import IhyeeError

from ihyee_mcp.formatter import format_fetch_response, format_search_response

server = Server("ihyee-mcp")


# ---------------------------------------------------------------------------
# Shared schema fragments
# ---------------------------------------------------------------------------

_FORMAT_PROP = {
    "type": "string",
    "enum": ["text", "markdown"],
    "description": "Output format: 'text' (plain text) or 'markdown' (rich markdown)",
    "default": "text",
}

_CONTENT_MODE_PROP = {
    "type": "string",
    "enum": ["both", "full_text", "summary", "markdown", "rag"],
    "description": (
        "What content to extract: 'both' (summary + full text), 'full_text', "
        "'summary', 'markdown' (structured markdown), or 'rag' (chunked for "
        "retrieval-augmented generation)"
    ),
    "default": "both",
}

_RAG_PROPERTIES = {
    "max_tokens_per_chunk": {
        "type": "integer",
        "description": "Max tokens per RAG chunk (50-2000). Only used when content_mode='rag'",
        "default": 500,
    },
    "max_tokens_per_result": {
        "type": "integer",
        "description": "Token budget per result (100-10000). Only used when content_mode='rag'",
        "default": 2000,
    },
    "max_total_tokens": {
        "type": "integer",
        "description": "Global token budget across all results (500-50000). Only used when content_mode='rag'",
        "default": 8000,
    },
    "min_relevance_score": {
        "type": "number",
        "description": "Minimum BM25 relevance score to include a chunk (0.0-1.0). Only used when content_mode='rag'",
        "default": 0.0,
    },
    "top_k_chunks": {
        "type": "integer",
        "description": "Keep only top-K most relevant chunks per result (1-200). Only used when content_mode='rag'",
    },
    "tokenizer": {
        "type": "string",
        "enum": ["cl100k", "o200k", "approximate"],
        "description": "Tokenizer for counting tokens: 'cl100k' (GPT-4), 'o200k' (GPT-4o), or 'approximate'. Only used when content_mode='rag'",
        "default": "cl100k",
    },
}


# ---------------------------------------------------------------------------
# Tool definitions
# ---------------------------------------------------------------------------

SEARCH_TOOL = Tool(
    name="ihyee_search",
    description=(
        "Search the web and return extracted, summarized content from top results. "
        "Uses Google search under the hood, then fetches and extracts clean text, "
        "summaries, and links from each result page. Supports RAG mode for "
        "chunked, token-budgeted output optimized for retrieval pipelines."
    ),
    inputSchema={
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": "The search query string",
            },
            "max_results": {
                "type": "integer",
                "description": "Maximum number of results to return (1-10)",
                "default": 5,
            },
            "format": _FORMAT_PROP,
            "content_mode": _CONTENT_MODE_PROP,
            "render": {
                "type": "boolean",
                "description": "Whether to use browser rendering for JavaScript-heavy pages",
                "default": False,
            },
            "before": {
                "type": "string",
                "description": "Only return results published before this date (YYYY-MM-DD)",
            },
            "after": {
                "type": "string",
                "description": "Only return results published after this date (YYYY-MM-DD)",
            },
            "must_have": {
                "type": "string",
                "description": "Require this exact phrase to appear in results",
            },
            "allintext": {
                "type": "string",
                "description": "Words that must ALL appear in the page body text",
            },
            **_RAG_PROPERTIES,
        },
        "required": ["query"],
    },
)

FETCH_TOOL = Tool(
    name="ihyee_fetch",
    description=(
        "Fetch and extract clean content from specific web page URLs. "
        "Returns extracted text, summaries, metadata, and links from each page. "
        "Supports RAG mode for chunked, token-budgeted output."
    ),
    inputSchema={
        "type": "object",
        "properties": {
            "urls": {
                "type": "array",
                "items": {"type": "string"},
                "description": "List of URLs to fetch (maximum 10)",
                "maxItems": 10,
            },
            "format": _FORMAT_PROP,
            "content_mode": _CONTENT_MODE_PROP,
            "render": {
                "type": "boolean",
                "description": "Whether to use browser rendering for JavaScript-heavy pages",
                "default": False,
            },
            **_RAG_PROPERTIES,
        },
        "required": ["urls"],
    },
)

RENDER_TOOL = Tool(
    name="ihyee_render",
    description=(
        "Force full browser rendering of a JavaScript-heavy web page using Playwright, "
        "then extract its content. Use this for single-page apps, pages behind JS frameworks, "
        "or when a regular fetch returns incomplete content. Supports RAG mode for "
        "chunked, token-budgeted output."
    ),
    inputSchema={
        "type": "object",
        "properties": {
            "url": {
                "type": "string",
                "description": "The URL to render",
            },
            "format": _FORMAT_PROP,
            "content_mode": _CONTENT_MODE_PROP,
            "wait_for": {
                "type": "string",
                "enum": ["networkidle", "domcontentloaded", "load"],
                "description": "When to consider the page loaded",
                "default": "networkidle",
            },
            "wait_selector": {
                "type": "string",
                "description": "CSS selector to wait for before extracting content",
            },
            "timeout_ms": {
                "type": "integer",
                "description": "Maximum time to wait for page load in milliseconds",
                "default": 30000,
            },
            **_RAG_PROPERTIES,
        },
        "required": ["url"],
    },
)

ALL_TOOLS = [SEARCH_TOOL, FETCH_TOOL, RENDER_TOOL]


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _extract_rag_kwargs(args: dict[str, Any]) -> dict[str, Any]:
    """Extract RAG-related parameters from tool arguments.

    Returns a dict suitable for passing as **kwargs to the SDK client methods.
    Only includes keys that were explicitly provided (non-None).
    """
    rag_keys = (
        "max_tokens_per_chunk",
        "max_tokens_per_result",
        "max_total_tokens",
        "min_relevance_score",
        "top_k_chunks",
        "tokenizer",
    )
    return {k: args[k] for k in rag_keys if args.get(k) is not None}


# ---------------------------------------------------------------------------
# Handlers
# ---------------------------------------------------------------------------

@server.list_tools()
async def list_tools() -> list[Tool]:
    """Return the list of available tools."""
    return ALL_TOOLS


@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]) -> Sequence[TextContent]:
    """Execute a tool call and return formatted results."""
    api_key = os.environ.get("IHYEE_API_KEY")
    if not api_key:
        return [TextContent(
            type="text",
            text="Error: IHYEE_API_KEY environment variable is not set. "
                 "Please configure it in your MCP server settings.",
        )]

    try:
        async with AsyncIhyee(api_key=api_key) as client:
            if name == "ihyee_search":
                return await _handle_search(client, arguments)
            elif name == "ihyee_fetch":
                return await _handle_fetch(client, arguments)
            elif name == "ihyee_render":
                return await _handle_render(client, arguments)
            else:
                return [TextContent(
                    type="text",
                    text=f"Error: Unknown tool '{name}'",
                )]
    except IhyeeError as exc:
        return [TextContent(
            type="text",
            text=f"ihyee API error ({exc.code}): {exc.message}",
        )]
    except Exception as exc:
        return [TextContent(
            type="text",
            text=f"Unexpected error: {type(exc).__name__}: {exc}",
        )]


async def _handle_search(
    client: AsyncIhyee, args: dict[str, Any]
) -> Sequence[TextContent]:
    """Handle the ihyee_search tool call."""
    response = await client.search(
        query=args["query"],
        max_results=args.get("max_results", 5),
        content_mode=args.get("content_mode", "both"),
        format=args.get("format", "text"),
        render=args.get("render", False),
        before=args.get("before"),
        after=args.get("after"),
        must_have=args.get("must_have"),
        allintext=args.get("allintext"),
        **_extract_rag_kwargs(args),
    )
    text = format_search_response(response)
    return [TextContent(type="text", text=text)]


async def _handle_fetch(
    client: AsyncIhyee, args: dict[str, Any]
) -> Sequence[TextContent]:
    """Handle the ihyee_fetch tool call."""
    urls = args["urls"]
    if len(urls) > 10:
        return [TextContent(
            type="text",
            text="Error: Maximum 10 URLs allowed per request.",
        )]

    response = await client.fetch(
        urls=urls,
        content_mode=args.get("content_mode", "both"),
        format=args.get("format", "text"),
        render=args.get("render", False),
        **_extract_rag_kwargs(args),
    )
    text = format_fetch_response(response, label="Fetched Content")
    return [TextContent(type="text", text=text)]


async def _handle_render(
    client: AsyncIhyee, args: dict[str, Any]
) -> Sequence[TextContent]:
    """Handle the ihyee_render tool call."""
    response = await client.render(
        url=args["url"],
        wait_for=args.get("wait_for", "networkidle"),
        wait_selector=args.get("wait_selector"),
        timeout_ms=args.get("timeout_ms", 30000),
        content_mode=args.get("content_mode", "both"),
        format=args.get("format", "text"),
        **_extract_rag_kwargs(args),
    )
    text = format_fetch_response(response, label="Rendered Page Content")
    return [TextContent(type="text", text=text)]


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main() -> None:
    """Run the MCP server on stdio transport."""
    async def _run() -> None:
        async with stdio_server() as (read_stream, write_stream):
            await server.run(
                read_stream,
                write_stream,
                server.create_initialization_options(),
            )

    asyncio.run(_run())


if __name__ == "__main__":
    main()
