#!/usr/bin/env python3
"""
Hybrid FastAPI + MCP server that exposes both REST endpoints and MCP tooling
backed by the same search indices.

Uses a reverse proxy pattern to avoid 307 redirect issues that some MCP clients
(like Claude) don't handle well. The MCP server runs on an internal port, and
FastAPI proxies requests to it.
"""

import argparse
import sys
import threading
import time
from contextlib import asynccontextmanager

import aiohttp
import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, Response

from ..mcp.mcp import FullTextSearchMCP
from ..core import IndexConfig, load_config_from_yaml
from ..server.api import SearchAPI


# Headers that shouldn't be forwarded between proxy hops
HOP_BY_HOP_HEADERS = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
}


def _filter_headers(headers: dict[str, str]) -> dict[str, str]:
    """Strip hop-by-hop headers; keep everything MCP cares about."""
    return {
        k: v
        for k, v in headers.items()
        if k.lower() not in HOP_BY_HOP_HEADERS and k.lower() != "host"
    }


class MCPPlusAPIServer:
    """Combined server that reuses a single set of indices for API + MCP.

    Uses a reverse proxy pattern: the MCP server runs on an internal port,
    and FastAPI forwards requests to it. This avoids 307 redirect issues
    that occur when mounting MCP directly into FastAPI.
    """

    def __init__(
        self,
        server_name: str,
        server_description: str,
        indices: list[IndexConfig],
        max_workers: int = 4,
        mcp_mount_path: str = "/mcp",
        mcp_internal_port: int = 9001,
        copy_middleware: bool = True,
    ):
        self.server_name = server_name
        self.server_description = server_description
        self.mcp_mount_path = self._normalize_mount_path(mcp_mount_path)
        self.mcp_internal_port = mcp_internal_port
        self.mcp_upstream_url = f"http://127.0.0.1:{mcp_internal_port}/mcp"
        self._copy_middleware = copy_middleware

        # Build indices once via the MCP server
        self.mcp_server = FullTextSearchMCP(
            server_name=server_name,
            server_description=server_description,
            indices=indices,
            max_workers=max_workers,
        )

        # Reuse MCP-built indices inside the REST API server
        self.api_server = SearchAPI(
            server_name=server_name,
            server_description=server_description,
            indices=None,
            max_workers=max_workers,
        )
        self.api_server.indices = self.mcp_server.indices
        self.api_server.citations = self.mcp_server.citations

        # Thread for running the internal MCP server
        self._mcp_thread: threading.Thread | None = None
        self._mcp_started = threading.Event()

        # aiohttp session (created in lifespan)
        self._http_session: aiohttp.ClientSession | None = None

        self.app = self._create_app()

    def _normalize_mount_path(self, mount_path: str) -> str:
        if not mount_path:
            return "/mcp"
        if not mount_path.startswith("/"):
            return f"/{mount_path}"
        return mount_path.rstrip("/")

    def _start_mcp_server_background(self):
        """Start the MCP server in a background thread."""
        mcp_instance = self.mcp_server.create_mcp_server()

        def run_mcp():
            print(f"Starting internal MCP server on 127.0.0.1:{self.mcp_internal_port}")
            self._mcp_started.set()
            mcp_instance.run(
                transport="streamable-http",
                host="127.0.0.1",
                port=self.mcp_internal_port,
                path="/mcp",
                stateless_http=True,
            )

        self._mcp_thread = threading.Thread(target=run_mcp, daemon=True)
        self._mcp_thread.start()

        # Wait for the MCP server to start
        self._mcp_started.wait(timeout=10)
        # Give it a moment to fully initialize
        time.sleep(0.5)

    def _create_app(self) -> FastAPI:
        title = f"{self.server_name} (API + MCP)"
        description = (
            f"{self.server_description}\n\n"
            f"REST API is available at the root paths, and MCP is proxied at {self.mcp_mount_path}."
        )

        @asynccontextmanager
        async def combined_lifespan(app: FastAPI):
            # Start MCP server in background thread
            self._start_mcp_server_background()

            # Create aiohttp session for proxying
            timeout = aiohttp.ClientTimeout(total=None)
            self._http_session = aiohttp.ClientSession(timeout=timeout)

            try:
                yield
            finally:
                # Clean up aiohttp session
                if self._http_session:
                    await self._http_session.close()

                # Clean up API resources
                await self.api_server.shutdown()
                if hasattr(self.mcp_server, "executor"):
                    self.mcp_server.executor.shutdown(wait=True)

        app = FastAPI(
            title=title,
            description=description,
            version="1.0.0",
            lifespan=combined_lifespan,
        )

        if self._copy_middleware:
            # Copy middleware (e.g., CORS) from the API app
            for middleware in self.api_server.app.user_middleware:
                app.add_middleware(middleware.cls, *middleware.args, **middleware.kwargs)

        # Include REST routes
        app.include_router(self.api_server.app.router)

        # Create proxy route for MCP
        # Handle both /mcp and /mcp/ and any sub-paths
        @app.api_route(
            self.mcp_mount_path + "{path:path}",
            methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
        )
        @app.api_route(
            self.mcp_mount_path,
            methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
        )
        async def mcp_proxy(request: Request, path: str = ""):
            """Proxy requests to the internal MCP server."""
            # Build upstream URL
            upstream_url = self.mcp_upstream_url
            if path:
                upstream_url = f"{upstream_url}/{path.lstrip('/')}"

            # Copy incoming headers (keep MCP-related ones like Mcp-Session-Id, Origin, etc.)
            incoming_headers = _filter_headers(dict(request.headers))

            # Read body (MCP payloads are typically small)
            body = await request.body()

            if not self._http_session:
                return JSONResponse(
                    status_code=503,
                    content={"error": "HTTP session not initialized"},
                )

            try:
                async with self._http_session.request(
                    method=request.method,
                    url=upstream_url,
                    headers=incoming_headers,
                    data=body if request.method in ("POST", "PUT", "PATCH") else None,
                ) as upstream_resp:
                    # Prepare downstream headers & content-type
                    resp_headers = _filter_headers(dict(upstream_resp.headers))
                    media_type = upstream_resp.headers.get("content-type")

                    # Read response body within context manager
                    # MCP payloads are typically small, so this is fine
                    response_body = await upstream_resp.read()

                    return Response(
                        content=response_body,
                        status_code=upstream_resp.status,
                        headers=resp_headers,
                        media_type=media_type,
                    )

            except aiohttp.ClientConnectorError:
                return JSONResponse(
                    status_code=503,
                    content={"error": "MCP server not available"},
                )

        return app

    def run(self, host: str = "0.0.0.0", port: int = 8000, **kwargs):
        print(f"Starting combined API + MCP server on {host}:{port}")
        print(f"REST endpoints available at http://{host}:{port}")
        print(
            f"MCP endpoint available at http://{host}:{port}" f"{self.mcp_mount_path}"
        )
        uvicorn.run(self.app, host=host, port=port, **kwargs)


def main():
    """CLI entry point for running the combined MCP + API server."""
    parser = argparse.ArgumentParser(description="Full-text search API + MCP server")

    parser.add_argument(
        "--config",
        required=True,
        help="Path to YAML configuration file",
    )
    parser.add_argument("--host", default="0.0.0.0", help="Host for HTTP server")
    parser.add_argument("--port", type=int, default=8000, help="Port for HTTP server")
    parser.add_argument(
        "--workers",
        type=int,
        default=4,
        help="Number of worker threads for search operations",
    )
    parser.add_argument(
        "--mcp-internal-port",
        type=int,
        default=9001,
        help="Internal port for MCP server (default: 9001)",
    )
    parser.add_argument(
        "--reload", action="store_true", help="Enable auto-reload for development"
    )

    args = parser.parse_args()

    try:
        server_name, server_description, index_configs, server_config = (
            load_config_from_yaml(args.config)
        )
        host = server_config.get("host", args.host)
        port = server_config.get("port", args.port)
        max_workers = server_config.get("max_workers", args.workers)
        mount_path = server_config.get("mount", "/mcp")
        mcp_internal_port = server_config.get(
            "mcp_internal_port", args.mcp_internal_port
        )

        server = MCPPlusAPIServer(
            server_name=server_name,
            server_description=server_description,
            indices=index_configs,
            max_workers=max_workers,
            mcp_mount_path=mount_path,
            mcp_internal_port=mcp_internal_port,
        )
        server.run(host=host, port=port, reload=args.reload)

    except Exception as e:
        print(f"Error starting combined server: {e}", file=sys.stderr)
        return 1

    return 0


if __name__ == "__main__":
    sys.exit(main())
