"""MCP 服务器入口"""

import json
import os
from typing import Any

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

from trailbase_mcp.client.client import TrailBaseClient
from trailbase_mcp.tools import (
    auth,
    config,
    db,
    files,
    info,
    jobs,
    logs,
    records,
    register_all_tools,
    schema,
    tx,
)
from trailbase_mcp.tools.registry import get_registry


class TrailBaseMCPServer:
    """TrailBase MCP 服务器"""

    def __init__(self, trailbase_url: str | None = None, trailbase_token: str | None = None):
        """
        初始化 MCP 服务器

        Args:
            trailbase_url: TrailBase 服务器地址，默认从 TRAILBASE_URL 环境变量读取
            trailbase_token: TrailBase 认证令牌，默认从 ADMIN_TOKEN 环境变量读取
        """
        self.server = Server("trailbase-mcp")
        # 如果参数为 None，从环境变量读取
        self.trailbase_url = trailbase_url or os.getenv("TRAILBASE_URL")
        self.trailbase_token = trailbase_token or os.getenv("ADMIN_TOKEN")
        self.client: TrailBaseClient | None = None
        # 注册所有工具
        register_all_tools()
        self._setup_handlers()

    def _setup_handlers(self):
        """设置 MCP 处理器"""

        @self.server.list_tools()
        async def list_tools() -> list[Tool]:
            """列出所有工具"""
            tools = []
            tools.extend(db.get_db_tools())
            tools.extend(records.get_records_tools())
            tools.extend(auth.get_auth_tools())
            tools.extend(schema.get_schema_tools())
            tools.extend(jobs.get_jobs_tools())
            tools.extend(logs.get_logs_tools())
            tools.extend(config.get_config_tools())
            tools.extend(files.get_files_tools())
            tools.extend(tx.get_tx_tools())
            tools.extend(info.get_info_tools())
            return tools

        @self.server.call_tool()
        async def call_tool(name: str, arguments: dict[str, Any] | None) -> list[TextContent]:
            """调用工具 - 使用注册表模式"""
            if not self.client:
                self.client = TrailBaseClient(
                    base_url=self.trailbase_url,
                    token=self.trailbase_token,
                )

            args = arguments or {}
            registry = get_registry()

            # 从注册表获取处理器
            handler = registry.get_handler(name)
            if handler:
                return await handler(self.client, args)
            else:
                return [TextContent(type="text", text=f'{{"error": "Unknown tool: {name}"}}')]

        @self.server.list_resources()
        async def list_resources() -> list[Resource]:
            """列出所有资源"""
            # 不提供资源，避免配置 API 返回 protobuf 导致的显示问题
            # 用户可以通过 config.get 工具获取配置，通过工具查询表结构
            return []

        @self.server.read_resource()
        async def read_resource(uri: str) -> str:
            """读取资源"""
            # 将 uri 转换为字符串（可能是 AnyUrl 对象）
            uri_str = str(uri)
            
            if not self.client:
                self.client = TrailBaseClient(
                    base_url=self.trailbase_url,
                    token=self.trailbase_token,
                )

            if uri_str == "trailbase://config":
                try:
                    # 直接使用 SDK 返回的数据，不做任何转换
                    config_data = self.client.http.get("/api/_admin/config")
                    return json.dumps(config_data, indent=2, ensure_ascii=False)
                except Exception as e:
                    return json.dumps({"error": str(e)}, ensure_ascii=False)

            elif uri_str == "trailbase://openapi":
                try:
                    openapi_data = self.client.http.get("/openapi.json")
                    return json.dumps(openapi_data, indent=2, ensure_ascii=False)
                except Exception:
                    # Fallback: 尝试其他端点
                    try:
                        openapi_data = self.client.http.get("/api/openapi.json")
                        return json.dumps(openapi_data, indent=2, ensure_ascii=False)
                    except Exception as e:
                        return json.dumps({"error": str(e)}, ensure_ascii=False)

            elif uri_str.startswith("trailbase://schema/"):
                table = uri_str.replace("trailbase://schema/", "")
                try:
                    result = self.client.http.post(
                        "/api/_admin/query", json={"query": f"PRAGMA table_info({table})"}
                    )
                    return json.dumps(result, indent=2, ensure_ascii=False)
                except Exception as e:
                    return json.dumps({"error": str(e)}, ensure_ascii=False)

            else:
                return json.dumps({"error": f"Unknown resource: {uri_str}"}, ensure_ascii=False)

    def _get_html_dashboard(self) -> str:
        """获取 HTML 展示页面内容"""
        import importlib.resources

        try:
            # 尝试从包中读取 HTML 文件
            package = importlib.resources.files("trailbase_mcp.static")
            html_file = package / "index.html"
            return html_file.read_text(encoding="utf-8")
        except Exception:
            # 如果文件不存在，返回内联 HTML
            return self._get_inline_html()

    def _get_inline_html(self) -> str:
        """返回内联 HTML（备用方案）"""
        return """<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>TrailBase MCP Server</title>
    <script src="https://cdn.tailwindcss.com"></script>
</head>
<body class="bg-gray-50 min-h-screen">
    <div class="container mx-auto px-4 py-8 max-w-7xl">
        <header class="bg-white rounded-lg shadow-md p-6 mb-6">
            <h1 class="text-3xl font-bold text-gray-900">TrailBase MCP Server</h1>
            <p class="text-gray-600 mt-2">Model Context Protocol Server for TrailBase</p>
        </header>
        <div class="bg-white rounded-lg shadow-md p-6">
            <h2 class="text-2xl font-bold text-gray-900 mb-4">Server Status</h2>
            <p class="text-gray-700">Server is running and ready to accept MCP connections.</p>
            <p class="text-sm text-gray-500 mt-4">Use MCP client to connect via Streamable HTTP protocol.</p>
        </div>
    </div>
</body>
</html>"""

    async def run_stdio(self):
        """通过 stdio 运行服务器"""
        import sys

        # 输出到 stderr，避免干扰 stdio 通信
        print("=" * 60, file=sys.stderr)
        print("TrailBase MCP Server", file=sys.stderr)
        print("=" * 60, file=sys.stderr)
        print("Transport: stdio (MCP 2025-03-26 Standard)", file=sys.stderr)
        print("Note: stdio 模式自动启动，无需单独运行服务器", file=sys.stderr)
        print("Waiting for MCP client connection...", file=sys.stderr)
        print("=" * 60, file=sys.stderr)
        sys.stderr.flush()

        async with stdio_server() as (read_stream, write_stream):
            print("✓ MCP client connected", file=sys.stderr)
            sys.stderr.flush()
            await self.server.run(
                read_stream,
                write_stream,
                self.server.create_initialization_options(),
            )

    async def run_streamable_http(self, host: str = "0.0.0.0", port: int = 9118):
        """通过 Streamable HTTP 运行服务器"""
        import sys
        import os

        # 确定显示的主机地址
        display_host = "localhost" if host == "0.0.0.0" else host
        trailbase_url = self.trailbase_url or os.getenv("TRAILBASE_URL", "未设置")

        print("=" * 60, file=sys.stderr)
        print("TrailBase MCP Server", file=sys.stderr)
        print("=" * 60, file=sys.stderr)
        print(f"Transport: Streamable HTTP (MCP 2025-03-26 Standard)", file=sys.stderr)
        print(f"MCP Server URL: http://{display_host}:{port}", file=sys.stderr)
        print(f"Web Dashboard: http://{display_host}:{port}/", file=sys.stderr)
        print(f"MCP Endpoint: http://{display_host}:{port}/message (Streamable HTTP)", file=sys.stderr)
        print(f"TrailBase URL: {trailbase_url}", file=sys.stderr)
        print("=" * 60, file=sys.stderr)
        print(f"✓ Server is running. Open http://{display_host}:{port} in your browser", file=sys.stderr)
        if trailbase_url == "未设置":
            print("⚠ Warning: TRAILBASE_URL not set. Set it to connect to TrailBase.", file=sys.stderr)
        print("=" * 60, file=sys.stderr)
        sys.stderr.flush()

        import contextlib

        manager = StreamableHTTPSessionManager(
            self.server,
        )

        # 创建 ASGI 应用
        async def asgi_app(scope, receive, send):
            """ASGI 应用包装器"""
            # 处理根路径，返回 HTML 展示页面
            if scope["type"] == "http" and scope["path"] == "/":
                await send(
                    {
                        "type": "http.response.start",
                        "status": 200,
                        "headers": [[b"content-type", b"text/html; charset=utf-8"]],
                    }
                )
                html_content = self._get_html_dashboard()
                await send({"type": "http.response.body", "body": html_content.encode("utf-8")})
                return

            # 处理 favicon.ico 等静态资源请求
            if scope["type"] == "http" and scope["path"] in ["/favicon.ico", "/robots.txt"]:
                await send(
                    {
                        "type": "http.response.start",
                        "status": 404,
                        "headers": [[b"content-type", b"text/plain"]],
                    }
                )
                await send({"type": "http.response.body", "body": b"Not Found"})
                return

            # 其他请求交给 MCP manager 处理
            await manager.handle_request(scope, receive, send)

        # 使用 ASGI 框架运行（如 uvicorn）
        try:
            import uvicorn

            # 创建 lifespan 上下文管理器
            @contextlib.asynccontextmanager
            async def lifespan(app):
                async with manager.run():
                    yield

            # 创建支持 lifespan 的 ASGI 应用
            async def app_with_lifespan(scope, receive, send):
                if scope["type"] == "lifespan":
                    async with lifespan(None) as lifespan_state:
                        while True:
                            message = await receive()
                            if message["type"] == "lifespan.startup":
                                await send({"type": "lifespan.startup.complete"})
                            elif message["type"] == "lifespan.shutdown":
                                await send({"type": "lifespan.shutdown.complete"})
                                break
                else:
                    await asgi_app(scope, receive, send)

            config = uvicorn.Config(
                app_with_lifespan, host=host, port=port, log_level="info"
            )
            server = uvicorn.Server(config)
            await server.serve()
        except ImportError:
            raise ImportError(
                "uvicorn is required for Streamable HTTP transport. "
                "Install it with: uv add uvicorn"
            )

    async def run(self, transport: str | None = None):
        """
        运行服务器

        Args:
            transport: 传输方式，'stdio' 或 'http'。如果为 None，则从环境变量 MCP_TRANSPORT 读取
                     默认使用 Streamable HTTP（最新标准）
        """
        transport = transport or os.getenv("MCP_TRANSPORT", "stdio")

        if transport == "http" or transport == "streamable_http":
            host = os.getenv("MCP_HOST", "0.0.0.0")
            port = int(os.getenv("MCP_PORT", "9118"))
            await self.run_streamable_http(host, port)
        else:
            await self.run_stdio()


async def main():
    """主函数"""
    server = TrailBaseMCPServer()
    await server.run()


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())

