"""MCP Server for Open WebUI Knowledge Bases.

An MCP (Model Context Protocol) server that exposes Open WebUI Knowledge Bases
as tools and resources. Supports both stdio and HTTP transports.

This implementation wraps the REST API endpoints, making it suitable for running
as a separate service or integrated with MCP clients.

Usage:
    Run this as a standalone server that connects to Open WebUI's REST API.
    Supports stdio (default) and HTTP transport modes.
"""

import asyncio
import hashlib
import json
import logging
import os
import re
import uuid
from contextvars import ContextVar
from typing import Any, Callable, Dict, Optional, Sequence

import httpx
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.server.streamable_http import StreamableHTTPServerTransport
from mcp.types import Resource, TextContent, Tool

log = logging.getLogger(__name__)

# Security constants
MAX_REQUEST_SIZE = 10 * 1024 * 1024  # 10MB
MAX_QUERY_LENGTH = 10000
MAX_KB_ID_LENGTH = 255
ALLOWED_KB_ID_PATTERN = re.compile(r"^[a-zA-Z0-9_-]+$")

# Rate limiting defaults (can be overridden via environment variables)
DEFAULT_RATE_LIMIT_PER_IP = "1000/minute"
DEFAULT_RATE_LIMIT_PER_TOKEN = "1000/minute"
DEFAULT_RATE_LIMIT_HEALTH = "10/minute"

# Connection constants
CONNECTION_TIMEOUT = 5.0
CONNECTION_ID_PREFIX_HTTP = "http_"
CONNECTION_ID_PREFIX_STDIO = "conn_"

# HTTP constants
BEARER_PREFIX = "Bearer "
AUTHORIZATION_HEADER = "Authorization"
SESSION_ID_HEADERS = ("mcp-session-id", "x-session-id")

# Token constants
TOKEN_HASH_LENGTH = 16
TOKEN_MASK_LENGTH = 4

# HTTP status codes
HTTP_STATUS_TOO_LARGE = 413
HTTP_STATUS_SERVICE_UNAVAILABLE = 503
HTTP_STATUS_INTERNAL_ERROR = 500
HTTP_STATUS_UNAUTHORIZED = 401
HTTP_STATUS_NOT_FOUND = 404

# Context variable for storing current HTTP request context
_http_request_context: ContextVar[Optional[Dict[str, Any]]] = ContextVar(
    "http_request_context", default=None
)


class KnowledgeMCPHTTPWrapper:
    """
    MCP Server that wraps Open WebUI Knowledge Base REST API endpoints.

    This approach uses HTTP requests to call the existing REST API,
    making it suitable for running as a separate service.

    Supports multi-user authentication:
    - Each connection can provide its own API token via OPEN_WEBUI_API_TOKEN env var
    - Falls back to default token if not provided (for backward compatibility)
    - Creates per-connection HTTP clients for isolation
    """

    def __init__(self, api_base_url: str = "", api_token: Optional[str] = None):
        """
        Initialize the HTTP wrapper MCP server.

        Args:
            api_base_url: Base URL of Open WebUI API
            api_token: Optional default Bearer token for backward compatibility.
                      Individual connections should provide their own token via
                      OPEN_WEBUI_API_TOKEN environment variable.
        """
        self.api_base_url = api_base_url.rstrip("/")
        self.default_api_token = api_token or os.getenv("OPEN_WEBUI_API_TOKEN")
        self.server = Server("open-webui-knowledge-http")
        self._http_mode = False

        # Per-connection storage
        self.connection_tokens: Dict[str, str] = {}
        self.connection_clients: Dict[str, httpx.AsyncClient] = {}
        self._connection_counter = 0

        self._register_handlers()

    def _get_connection_id(self) -> str:
        """
        Get a unique connection ID for the current connection.

        For stdio transport: Uses process ID
        For HTTP transport: Uses session ID from request context
        """
        # Check if we're in HTTP mode and have request context
        if self._http_mode:
            ctx = _http_request_context.get()
            if ctx and "session_id" in ctx:
                return f"{CONNECTION_ID_PREFIX_HTTP}{ctx['session_id']}"

        # Fallback to process ID for stdio
        return f"{CONNECTION_ID_PREFIX_STDIO}{os.getpid()}"

    def _get_token_for_connection(self, connection_id: str) -> Optional[str]:
        """
        Get API token for a connection.

        Priority:
        1. Connection-specific token (from environment variable per connection)
        2. Default token (from initialization or environment)

        Returns:
            API token string or None if not found
        """
        # Check if we already have a token for this connection
        if connection_id in self.connection_tokens:
            token = self.connection_tokens[connection_id]
            if token:  # Return cached token if valid
                return token

        # For HTTP mode, try to extract token from request context
        if self._http_mode:
            ctx = _http_request_context.get()
            if ctx and "token" in ctx and ctx["token"]:
                token = ctx["token"]
                if self._validate_token(token):
                    self.connection_tokens[connection_id] = token
                    log.debug(f"Using token from HTTP header for connection {connection_id[:8]}...")
                    return token

        # For stdio, each connection is a separate process, so check env var
        # This allows each connection to provide its own token
        token = os.getenv("OPEN_WEBUI_API_TOKEN")

        if token:
            # Validate token format
            if self._validate_token(token):
                self.connection_tokens[connection_id] = token
                log.debug(f"Using token from env var for connection {connection_id[:8]}...")
                return token
            else:
                log.warning(f"Invalid token format for connection {connection_id[:8]}...")

        # Fall back to default token (for backward compatibility)
        if self.default_api_token:
            self.connection_tokens[connection_id] = self.default_api_token
            log.debug(f"Using default token for connection {connection_id[:8]}...")
            return self.default_api_token

        return None

    def _validate_token(self, token: str) -> bool:
        """
        Validate token format.

        Args:
            token: Token string to validate

        Returns:
            True if token format is valid
        """
        if not token or not isinstance(token, str):
            return False

        token = token.strip()

        # API key format: starts with "sk-" and has reasonable length
        if token.startswith("sk-"):
            # Allow minimum length of 3 ("sk-") but prefer longer tokens
            # Maximum length to prevent abuse
            return len(token) >= 3 and len(token) <= 500

        # JWT format: should have 2 or 3 parts separated by dots
        # (unsigned JWTs have 2 parts, signed have 3)
        if "." in token:
            parts = token.split(".")
            if len(parts) in (2, 3):
                # Basic JWT structure check: header.payload or header.payload.signature
                return all(len(part) > 0 for part in parts) and len(token) > 20

        return False

    def _mask_token(self, token: Optional[str]) -> str:
        """
        Mask token for logging (show only first 4 and last 4 chars).

        Args:
            token: Token to mask

        Returns:
            Masked token string
        """
        if not token:
            return "None"
        if len(token) <= TOKEN_MASK_LENGTH * 2:
            return "****"
        return f"{token[:TOKEN_MASK_LENGTH]}...{token[-TOKEN_MASK_LENGTH:]}"

    async def get_client_for_connection(self, connection_id: str) -> httpx.AsyncClient:
        """
        Get or create HTTP client for a specific connection.

        Args:
            connection_id: Unique connection identifier

        Returns:
            httpx.AsyncClient configured with connection's token

        Raises:
            ValueError: If no token is available for the connection
        """
        # Return existing client if available
        if connection_id in self.connection_clients:
            return self.connection_clients[connection_id]

        # Get token for this connection
        token = self._get_token_for_connection(connection_id)

        if not token:
            raise ValueError(
                f"No API token found for connection {connection_id[:8]}... "
                f"Please set OPEN_WEBUI_API_TOKEN environment variable or provide "
                f"Authorization header with Bearer token."
            )

        # Create new client with connection's token
        headers = {"Authorization": f"Bearer {token}"}
        client = httpx.AsyncClient(
            base_url=self.api_base_url,
            headers=headers,
            timeout=httpx.Timeout(30.0, connect=10.0),
            limits=httpx.Limits(max_keepalive_connections=5, max_connections=10),
        )

        self.connection_clients[connection_id] = client
        masked_token = self._mask_token(token)
        log.debug(
            f"Created HTTP client for connection {connection_id[:8]}... " f"(token: {masked_token})"
        )

        return client

    async def _cleanup_connection(self, connection_id: str):
        """
        Clean up resources for a connection.

        Args:
            connection_id: Connection ID to clean up
        """
        # Close and remove HTTP client
        if connection_id in self.connection_clients:
            try:
                await self.connection_clients[connection_id].aclose()
            except Exception as e:
                log.warning(f"Error closing client for connection {connection_id[:8]}...: {e}")
            del self.connection_clients[connection_id]

        # Remove token (don't log the actual token)
        if connection_id in self.connection_tokens:
            del self.connection_tokens[connection_id]

        log.debug(f"Cleaned up resources for connection {connection_id[:8]}...")

    def _validate_knowledge_base_id(self, kb_id: str) -> None:
        """
        Validate knowledge base ID format.

        Args:
            kb_id: Knowledge base ID to validate

        Raises:
            ValueError: If ID is invalid
        """
        if not kb_id or not isinstance(kb_id, str):
            raise ValueError("knowledge_base_id must be a non-empty string")

        kb_id = kb_id.strip()

        if len(kb_id) > MAX_KB_ID_LENGTH:
            raise ValueError(f"knowledge_base_id exceeds maximum length of {MAX_KB_ID_LENGTH}")

        if not ALLOWED_KB_ID_PATTERN.match(kb_id):
            raise ValueError("knowledge_base_id contains invalid characters")

    def _validate_query(self, query: str) -> None:
        """
        Validate search query.

        Args:
            query: Query string to validate

        Raises:
            ValueError: If query is invalid
        """
        if not query or not isinstance(query, str):
            raise ValueError("query must be a non-empty string")

        query = query.strip()

        if len(query) > MAX_QUERY_LENGTH:
            raise ValueError(f"query exceeds maximum length of {MAX_QUERY_LENGTH}")

    def _sanitize_error_message(self, error_text: str) -> str:
        """
        Sanitize error messages to prevent information disclosure.

        Args:
            error_text: Raw error text from backend

        Returns:
            Sanitized error message safe to return to client
        """
        # Remove potential sensitive information patterns
        # Replace HTTP error details with generic message
        if "HTTP error" in error_text or "status" in error_text.lower():
            return "An error occurred while processing the request. Please try again."

        # Limit error message length
        MAX_ERROR_MESSAGE_LENGTH = 500
        if len(error_text) > MAX_ERROR_MESSAGE_LENGTH:
            return error_text[:MAX_ERROR_MESSAGE_LENGTH] + "..."

        return error_text

    async def _handle_http_error(
        self, e: httpx.HTTPStatusError, connection_id: str, kb_id: Optional[str] = None
    ) -> None:
        """
        Handle HTTP errors consistently across all methods.

        Args:
            e: HTTPStatusError from httpx
            connection_id: Connection ID for logging and cleanup
            kb_id: Optional knowledge base ID for 404 errors

        Raises:
            ValueError: For authentication (401) or not found (404) errors
            RuntimeError: For other HTTP errors
        """
        if e.response.status_code == HTTP_STATUS_UNAUTHORIZED:
            conn_id_short = connection_id[:8]
            log.warning(
                f"Authentication failed for connection {conn_id_short}... " f"(401 Unauthorized)"
            )
            await self._cleanup_connection(connection_id)
            raise ValueError("Authentication failed. Please check your API token.")
        elif e.response.status_code == HTTP_STATUS_NOT_FOUND and kb_id:
            raise ValueError(f"Knowledge base not found: {kb_id}")
        else:
            error_msg = self._sanitize_error_message(
                f"HTTP error {e.response.status_code}: Request failed"
            )
            raise RuntimeError(error_msg)

    def _get_tool_definitions(self) -> list[Tool]:
        """
        Get list of available MCP tools.

        Returns:
            List of Tool objects defining available tools
        """
        return [
            Tool(
                name="search_knowledge_base",
                description="Search a knowledge base using semantic search via Open WebUI API",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "knowledge_base_id": {
                            "type": "string",
                            "description": "The ID of the knowledge base to search",
                        },
                        "query": {"type": "string", "description": "The search query"},
                        "k": {
                            "type": "integer",
                            "description": "Number of results to return (default: 5)",
                            "default": 5,
                        },
                    },
                    "required": ["knowledge_base_id", "query"],
                },
            ),
            Tool(
                name="list_knowledge_bases",
                description="List all knowledge bases accessible via Open WebUI API",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "permission": {
                            "type": "string",
                            "enum": ["read", "write"],
                            "description": "Filter by permission level",
                            "default": "read",
                        }
                    },
                },
            ),
            Tool(
                name="get_knowledge_base_info",
                description=("Get detailed information about a knowledge base via Open WebUI API"),
                inputSchema={
                    "type": "object",
                    "properties": {
                        "knowledge_base_id": {
                            "type": "string",
                            "description": "The ID of the knowledge base",
                        }
                    },
                    "required": ["knowledge_base_id"],
                },
            ),
        ]

    def _register_handlers(self):
        """Register MCP protocol handlers."""

        @self.server.list_tools()
        async def list_tools() -> list[Tool]:
            """List available tools."""
            return self._get_tool_definitions()

        @self.server.call_tool()
        async def call_tool(name: str, arguments: Dict[str, Any]) -> Sequence[TextContent]:
            """Handle tool calls via HTTP."""
            # Get connection ID for this request
            connection_id = self._get_connection_id()

            # Get client for this connection
            try:
                client = await self.get_client_for_connection(connection_id)
            except ValueError as e:
                error_msg = (
                    f"Authentication Error: {str(e)}\n\n"
                    f"To use this MCP server, set the OPEN_WEBUI_API_TOKEN "
                    f"environment variable with your Open WebUI API token "
                    f"(starts with 'sk-')."
                )
                return [TextContent(type="text", text=error_msg)]

            if name == "search_knowledge_base":
                return await self._search_knowledge_base_http(arguments, client)
            elif name == "list_knowledge_bases":
                return await self._list_knowledge_bases_http(arguments, client)
            elif name == "get_knowledge_base_info":
                return await self._get_knowledge_base_info_http(arguments, client)
            else:
                raise ValueError(f"Unknown tool: {name}")

        @self.server.list_resources()
        async def list_resources() -> list[Resource]:
            """List available resources (knowledge bases) via HTTP."""
            connection_id = self._get_connection_id()

            try:
                client = await self.get_client_for_connection(connection_id)
            except ValueError as e:
                log.warning(f"Authentication error listing resources: {e}")
                return []

            try:
                response = await client.get("/knowledge/")
                response.raise_for_status()
                knowledge_bases = response.json()

                resources = []
                for kb in knowledge_bases:
                    resources.append(
                        Resource(
                            uri=f"knowledge://{kb['id']}",
                            name=kb.get("name", "Unknown"),
                            description=kb.get(
                                "description", f"Knowledge base: {kb.get('name', 'Unknown')}"
                            ),
                            mimeType="application/json",
                        )
                    )

                return resources
            except httpx.HTTPStatusError as e:
                if e.response.status_code == HTTP_STATUS_UNAUTHORIZED:
                    await self._handle_http_error(e, connection_id)
                log.exception(f"Error listing resources: {e}")
                return []
            except Exception as e:
                log.exception(f"Error listing resources: {e}")
                return []

        @self.server.read_resource()
        async def read_resource(uri: str) -> str:
            """Read a knowledge base resource via HTTP."""
            if not uri.startswith("knowledge://"):
                raise ValueError(f"Invalid resource URI: {uri}")

            kb_id = uri.replace("knowledge://", "")
            self._validate_knowledge_base_id(kb_id)
            connection_id = self._get_connection_id()

            try:
                client = await self.get_client_for_connection(connection_id)
            except ValueError as e:
                raise ValueError(f"Authentication Error: {str(e)}")

            try:
                response = await client.get(f"/knowledge/{kb_id}")
                response.raise_for_status()
                kb_data = response.json()

                return json.dumps(kb_data, indent=2)
            except httpx.HTTPStatusError as e:
                await self._handle_http_error(e, connection_id, kb_id=kb_id)

    async def _search_knowledge_base_http(
        self, arguments: Dict[str, Any], client: httpx.AsyncClient
    ) -> Sequence[TextContent]:
        """Search a knowledge base via HTTP API."""
        kb_id = arguments.get("knowledge_base_id")
        query = arguments.get("query")
        k = arguments.get("k", 5)

        if not kb_id or not query:
            raise ValueError("knowledge_base_id and query are required")

        # Validate inputs
        self._validate_knowledge_base_id(kb_id)
        self._validate_query(query)

        # Validate k parameter
        if not isinstance(k, int) or k < 1 or k > 100:
            raise ValueError("k must be an integer between 1 and 100")

        try:
            # Use the retrieval API endpoint
            response = await client.post(
                "/retrieval/query/collection",
                json={"collection_names": [kb_id], "query": query, "k": k},
            )
            response.raise_for_status()
            result = response.json()

            if not result.get("documents") or not result["documents"][0]:
                return [TextContent(type="text", text=f"No results found for query: {query}")]

            # Format results
            documents = result["documents"][0]
            metadatas = result.get("metadatas", [[]])[0] if result.get("metadatas") else []
            distances = result.get("distances", [[]])[0] if result.get("distances") else []

            results_text = f"Found {len(documents)} results for query: '{query}'\n\n"

            for idx, (doc, metadata) in enumerate(
                zip(documents, metadatas or [None] * len(documents))
            ):
                results_text += f"--- Result {idx + 1} ---\n"
                if distances and idx < len(distances):
                    results_text += f"Relevance Score: {distances[idx]:.4f}\n"
                if metadata:
                    if metadata.get("name"):
                        results_text += f"Source: {metadata.get('name')}\n"
                    if metadata.get("file_id"):
                        results_text += f"File ID: {metadata.get('file_id')}\n"
                results_text += f"\n{doc}\n\n"

            return [TextContent(type="text", text=results_text)]

        except httpx.HTTPStatusError as e:
            if e.response.status_code == 401:
                connection_id = self._get_connection_id()
                conn_id_short = connection_id[:8]
                log.warning(
                    f"Authentication failed for connection {conn_id_short}... "
                    f"(401 Unauthorized)"
                )
                await self._cleanup_connection(connection_id)
                raise ValueError("Authentication failed. Please check your API token.")
            if e.response.status_code == 404:
                raise ValueError(f"Knowledge base not found: {kb_id}")
            error_msg = self._sanitize_error_message(
                f"HTTP error {e.response.status_code}: Request failed"
            )
            raise RuntimeError(error_msg)
        except Exception as e:
            log.exception(f"Error searching knowledge base: {e}")
            raise RuntimeError("An error occurred while searching the knowledge base")

    async def _list_knowledge_bases_http(
        self, arguments: Dict[str, Any], client: httpx.AsyncClient
    ) -> Sequence[TextContent]:
        """List knowledge bases via HTTP API."""
        permission = arguments.get("permission", "read")

        # Validate permission parameter
        if permission not in ("read", "write"):
            raise ValueError("permission must be 'read' or 'write'")

        try:
            # Use /list endpoint for write permission, / for read
            endpoint = "/knowledge/list" if permission == "write" else "/knowledge/"
            response = await client.get(endpoint)
            response.raise_for_status()
            kb_list = response.json()

            if not kb_list:
                return [TextContent(type="text", text="No knowledge bases found.")]

            result_text = f"Found {len(kb_list)} knowledge base(s):\n\n"

            for kb in kb_list:
                file_count = len(kb.get("files", []))
                result_text += f"- {kb.get('name', 'Unknown')} (ID: {kb.get('id')})\n"
                if kb.get("description"):
                    result_text += f"  Description: {kb.get('description')}\n"
                result_text += f"  Files: {file_count}\n"
                result_text += f"  Created: {kb.get('created_at')}\n\n"

            return [TextContent(type="text", text=result_text)]

        except httpx.HTTPStatusError as e:
            connection_id = self._get_connection_id()
            await self._handle_http_error(e, connection_id)
        except Exception as e:
            log.exception(f"Error listing knowledge bases: {e}")
            raise RuntimeError("An error occurred while listing knowledge bases")

    async def _get_knowledge_base_info_http(
        self, arguments: Dict[str, Any], client: httpx.AsyncClient
    ) -> Sequence[TextContent]:
        """Get knowledge base information via HTTP API."""
        kb_id = arguments.get("knowledge_base_id")

        if not kb_id:
            raise ValueError("knowledge_base_id is required")

        # Validate input
        self._validate_knowledge_base_id(kb_id)

        try:
            response = await client.get(f"/knowledge/{kb_id}")
            response.raise_for_status()
            kb_data = response.json()

            return [TextContent(type="text", text=json.dumps(kb_data, indent=2))]
        except httpx.HTTPStatusError as e:
            if e.response.status_code == 401:
                connection_id = self._get_connection_id()
                conn_id_short = connection_id[:8]
                log.warning(
                    f"Authentication failed for connection {conn_id_short}... "
                    f"(401 Unauthorized)"
                )
                await self._cleanup_connection(connection_id)
                raise ValueError("Authentication failed. Please check your API token.")
            if e.response.status_code == 404:
                raise ValueError(f"Knowledge base not found: {kb_id}")
            error_msg = self._sanitize_error_message(
                f"HTTP error {e.response.status_code}: Request failed"
            )
            raise RuntimeError(error_msg)

    async def run_stdio(self):
        """Run the server using stdio transport."""
        try:
            async with stdio_server() as (read_stream, write_stream):
                await self.server.run(
                    read_stream, write_stream, self.server.create_initialization_options()
                )
        finally:
            # Clean up all connection clients
            for connection_id in list(self.connection_clients.keys()):
                await self._cleanup_connection(connection_id)

    async def run_http(self, host: str = "0.0.0.0", port: int = 8001):
        """
        Run the server using HTTP transport with FastAPI.

        This allows remote access via HTTP/SSE endpoints.
        Clients connect via HTTP POST requests with SSE streaming.

        Args:
            host: Host to bind to (default: 0.0.0.0 for all interfaces)
            port: Port to listen on (default: 8001)
        """
        try:
            import uvicorn
            from fastapi import FastAPI, Request
            from fastapi.middleware.cors import CORSMiddleware
            from fastapi.responses import StreamingResponse
            from slowapi import Limiter, _rate_limit_exceeded_handler
            from slowapi.errors import RateLimitExceeded
            from slowapi.util import get_remote_address
        except ImportError:
            raise ImportError(
                "FastAPI, uvicorn, and slowapi are required for HTTP transport. "
                "Install with: pip install fastapi uvicorn[standard] slowapi"
            )

        self._http_mode = True
        app = FastAPI(
            title="Open WebUI Knowledge MCP Server",
            description="MCP Server for Open WebUI Knowledge Bases - HTTP Transport",
            version="1.0.0",
        )

        # Initialize rate limiter
        # Get rate limit configuration from environment or use defaults
        rate_limit_per_ip = os.getenv("MCP_RATE_LIMIT_PER_IP", DEFAULT_RATE_LIMIT_PER_IP)
        rate_limit_per_token = os.getenv("MCP_RATE_LIMIT_PER_TOKEN", DEFAULT_RATE_LIMIT_PER_TOKEN)
        rate_limit_health = os.getenv("MCP_RATE_LIMIT_HEALTH", DEFAULT_RATE_LIMIT_HEALTH)

        limiter = Limiter(
            key_func=get_remote_address,
            default_limits=[rate_limit_per_ip],
            storage_uri="memory://",
        )
        app.state.limiter = limiter
        app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

        # Add CORS middleware for web clients
        # Get allowed origins from environment or default to empty (no CORS)
        allowed_origins_env = os.getenv("MCP_CORS_ORIGINS", "")
        if allowed_origins_env:
            allowed_origins = [origin.strip() for origin in allowed_origins_env.split(",")]
        else:
            # Default: no CORS (empty list) - most secure
            # Set MCP_CORS_ORIGINS environment variable to enable CORS
            allowed_origins = []

        app.add_middleware(
            CORSMiddleware,
            allow_origins=allowed_origins,
            allow_credentials=True,
            allow_methods=["GET", "POST"],
            allow_headers=["Authorization", "Content-Type", "mcp-session-id", "x-session-id"],
        )

        # Store transport instances and connection managers per session
        # Each transport manages its own connection lifecycle
        session_transports: Dict[str, StreamableHTTPServerTransport] = {}
        session_connections: Dict[str, Any] = {}  # Store connection context managers
        session_tasks: Dict[str, asyncio.Task] = {}
        connection_ready: Dict[str, asyncio.Event] = {}  # Signal when connection is ready

        async def run_server_for_transport(
            session_id: str, transport: StreamableHTTPServerTransport
        ):
            """Run MCP server connected to transport streams."""
            connection_id = f"{CONNECTION_ID_PREFIX_HTTP}{session_id}"
            try:
                # Establish connection and keep it alive
                connection_manager = transport.connect()
                session_connections[session_id] = connection_manager

                async with connection_manager as (read_stream, write_stream):
                    # Signal that connection is ready
                    if session_id in connection_ready:
                        connection_ready[session_id].set()

                    # Run the server with the transport's streams
                    await self.server.run(
                        read_stream, write_stream, self.server.create_initialization_options()
                    )
            except asyncio.CancelledError:
                log.debug(f"Server task cancelled for session {session_id[:8]}...")
            except Exception as e:
                log.exception(f"Error running server for session {session_id[:8]}...: {e}")
            finally:
                # Cleanup on disconnect
                if connection_id in self.connection_clients:
                    await self._cleanup_connection(connection_id)
                if session_id in session_transports:
                    del session_transports[session_id]
                if session_id in session_connections:
                    del session_connections[session_id]
                if session_id in session_tasks:
                    del session_tasks[session_id]
                if session_id in connection_ready:
                    del connection_ready[session_id]

        def get_token_for_rate_limit(request: Request) -> str:
            """Extract token from request for rate limiting."""
            auth_header = request.headers.get(AUTHORIZATION_HEADER, "")
            if auth_header.startswith(BEARER_PREFIX):
                token = auth_header[len(BEARER_PREFIX) :].strip()
                if self._validate_token(token):
                    # Use token hash for rate limiting (don't store full token)
                    return hashlib.sha256(token.encode()).hexdigest()[:TOKEN_HASH_LENGTH]
            # Fall back to IP if no valid token
            return get_remote_address(request)

        def _extract_session_id(self, request: Request) -> str:
            """Extract or generate session ID from request headers."""
            session_id = request.headers.get(SESSION_ID_HEADERS[0]) or request.headers.get(
                SESSION_ID_HEADERS[1]
            )
            if not session_id:
                session_id = str(uuid.uuid4())
            return session_id

        def _extract_token_from_request(self, request: Request) -> Optional[str]:
            """Extract and validate token from Authorization header."""
            auth_header = request.headers.get(AUTHORIZATION_HEADER, "")
            if auth_header.startswith(BEARER_PREFIX):
                token = auth_header[len(BEARER_PREFIX) :].strip()
                if self._validate_token(token):
                    return token
                log.warning("Invalid token format in request")
            return None

        async def _ensure_transport_ready(
            self,
            session_id: str,
            session_transports: Dict[str, StreamableHTTPServerTransport],
            session_tasks: Dict[str, asyncio.Task],
            session_connections: Dict[str, Any],
            connection_ready: Dict[str, asyncio.Event],
            run_server_for_transport: Callable,
        ) -> StreamableHTTPServerTransport:
            """Ensure transport is ready for the session."""
            # Get or create transport for this session
            if session_id not in session_transports:
                transport = StreamableHTTPServerTransport(
                    mcp_session_id=session_id, is_json_response_enabled=False
                )
                session_transports[session_id] = transport
                connection_ready[session_id] = asyncio.Event()

                # Start server task for this session (runs in background)
                server_task = asyncio.create_task(run_server_for_transport(session_id, transport))
                session_tasks[session_id] = server_task
                log.info(f"Created new transport and server task for session {session_id[:8]}...")

                # Wait for connection to be established (with timeout)
                try:
                    await asyncio.wait_for(
                        connection_ready[session_id].wait(), timeout=CONNECTION_TIMEOUT
                    )
                    log.debug(f"Connection established for session {session_id[:8]}...")
                except asyncio.TimeoutError:
                    log.warning(
                        f"Connection timeout for session {session_id[:8]}..., proceeding anyway"
                    )

            transport = session_transports[session_id]

            # Ensure server task is still running and connection is active
            if session_id in session_tasks:
                task = session_tasks[session_id]
                if task.done():
                    # Task completed, create new one
                    log.warning(
                        f"Server task completed for session {session_id[:8]}..., recreating..."
                    )
                    if session_id not in connection_ready:
                        connection_ready[session_id] = asyncio.Event()
                    new_task = asyncio.create_task(run_server_for_transport(session_id, transport))
                    session_tasks[session_id] = new_task
                    # Wait for new connection
                    try:
                        await asyncio.wait_for(
                            connection_ready[session_id].wait(), timeout=CONNECTION_TIMEOUT
                        )
                    except asyncio.TimeoutError:
                        log.warning(f"Reconnection timeout for session {session_id[:8]}...")

            # Verify transport is ready (has active connection)
            if session_id not in session_connections:
                raise RuntimeError("Transport connection not established")

            return transport

        @app.post("/mcp")
        @app.get("/mcp")  # Also handle GET for SSE connections
        @limiter.limit(rate_limit_per_token, key_func=get_token_for_rate_limit)
        async def handle_mcp_request(request: Request):
            """Handle MCP requests via HTTP POST with SSE streaming."""
            # Enforce request size limit
            content_length = request.headers.get("content-length")
            if content_length and int(content_length) > MAX_REQUEST_SIZE:
                return {"error": "Request body too large"}, HTTP_STATUS_TOO_LARGE

            # Extract session ID and token
            session_id = self._extract_session_id(request)
            connection_id = f"{CONNECTION_ID_PREFIX_HTTP}{session_id}"
            token = self._extract_token_from_request(request)

            # Set request context for handlers to access
            ctx = {"session_id": session_id, "token": token, "connection_id": connection_id}
            _http_request_context.set(ctx)

            # Store token for this connection
            if token:
                self.connection_tokens[connection_id] = token
                log.debug(f"Stored token for HTTP session {session_id[:8]}...")

            # Ensure transport is ready
            try:
                transport = await self._ensure_transport_ready(
                    session_id,
                    session_transports,
                    session_tasks,
                    session_connections,
                    connection_ready,
                    run_server_for_transport,
                )
            except RuntimeError as e:
                return {"error": str(e)}, HTTP_STATUS_SERVICE_UNAVAILABLE

            # Use ASGI interface - FastAPI requests are ASGI-compatible
            # We'll create an ASGI response handler
            response_chunks = []
            response_started = False

            async def asgi_send(message):
                nonlocal response_started
                response_chunks.append(message)
                if message["type"] == "http.response.start":
                    response_started = True

            async def asgi_receive():
                # Return request body
                body = await request.body()
                return {"type": "http.request", "body": body, "more_body": False}

            # Handle request through transport
            try:
                # Set context before handling
                _http_request_context.set(ctx)

                # Get ASGI scope from FastAPI request
                scope = dict(request.scope)

                # Ensure Accept header includes text/event-stream for SSE
                # MCP transport requires this header
                headers_list = list(scope.get("headers", []))
                accept_header_found = False

                # Check if Accept header exists and update it
                for i, (name, value) in enumerate(headers_list):
                    if name.lower() == b"accept":
                        # Add text/event-stream to Accept header if not present
                        current_accept = (
                            value.decode("latin-1") if isinstance(value, bytes) else value
                        )
                        if "text/event-stream" not in current_accept:
                            new_accept = f"{current_accept}, text/event-stream"
                            headers_list[i] = (name, new_accept.encode("latin-1"))
                        accept_header_found = True
                        break

                # Add Accept header if it doesn't exist
                if not accept_header_found:
                    headers_list.append((b"accept", b"text/event-stream"))

                # Update scope with modified headers
                scope["headers"] = headers_list

                # Create a queue for receive calls
                receive_queue = asyncio.Queue()
                body = await request.body()

                # Enforce request size limit
                if len(body) > MAX_REQUEST_SIZE:
                    return {"error": "Request body too large"}, HTTP_STATUS_TOO_LARGE

                receive_queue.put_nowait({"type": "http.request", "body": body, "more_body": False})

                async def receive():
                    return await receive_queue.get()

                # Handle through transport
                await transport.handle_request(scope, receive, asgi_send)

                # Process response chunks
                if response_chunks:
                    first_chunk = response_chunks[0]
                    if first_chunk.get("type") == "http.response.start":
                        status = first_chunk["status"]
                        # Convert headers from bytes to strings
                        raw_headers = first_chunk["headers"]
                        headers = {}
                        for header_name, header_value in raw_headers:
                            # Headers come as tuples of (bytes, bytes) from ASGI
                            name = (
                                header_name.decode("latin-1")
                                if isinstance(header_name, bytes)
                                else header_name
                            )
                            value = (
                                header_value.decode("latin-1")
                                if isinstance(header_value, bytes)
                                else header_value
                            )
                            headers[name] = value

                        async def generate():
                            for chunk in response_chunks[1:]:
                                if chunk.get("type") == "http.response.body":
                                    body_data = chunk.get("body", b"")
                                    if body_data:
                                        yield body_data
                            # Signal end
                            yield b""

                        return StreamingResponse(
                            generate(),
                            status_code=status,
                            headers=headers,
                            media_type="text/event-stream",
                        )

                # Fallback response
                return {"error": "No response from transport"}, HTTP_STATUS_INTERNAL_ERROR

            except Exception as e:
                log.exception(f"Error handling MCP request: {e}")
                # Don't expose internal error details to client
                return {"error": "Internal server error"}, HTTP_STATUS_INTERNAL_ERROR

        @app.get("/health")
        @limiter.limit(rate_limit_health)
        async def health_check():
            """Health check endpoint."""
            return {
                "status": "ok",
                "service": "open-webui-knowledge-mcp",
                "transport": "http",
                "connections": len(self.connection_clients),
            }

        @app.get("/")
        async def root():
            """Root endpoint with API information."""
            return {
                "service": "Open WebUI Knowledge MCP Server",
                "version": "1.0.0",
                "transport": "http",
                "endpoints": {"mcp": "/mcp", "health": "/health"},
                "usage": {
                    "mcp_endpoint": f"POST http://{host}:{port}/mcp",
                    "headers": {
                        "Authorization": "Bearer sk-your-token",
                        "mcp-session-id": "optional-session-id",
                        "Content-Type": "application/json",
                    },
                },
            }

        # Run the server
        config = uvicorn.Config(app, host=host, port=port, log_level="info", access_log=True)
        server_instance = uvicorn.Server(config)

        try:
            log.info(f"Starting Knowledge MCP Server (HTTP) on http://{host}:{port}")
            log.info(f"MCP endpoint: http://{host}:{port}/mcp")
            log.info(f"Health check: http://{host}:{port}/health")
            await server_instance.serve()
        finally:
            # Cleanup all connections
            log.info("Shutting down HTTP server, cleaning up connections...")
            for connection_id in list(self.connection_clients.keys()):
                await self._cleanup_connection(connection_id)
            session_transports.clear()


# Example usage
async def main():
    """Start standalone HTTP wrapper server.

    To get your API token:
    1. Via Web UI:
       - Log into Open WebUI
       - Go to Settings → Account
       - Scroll to "API keys" section
       - Click "Show" to reveal your API key (starts with sk-)
       - Or click the refresh icon to create a new one

    2. Via API (using session token):
       - Login: POST /api/v1/auths/signin
       - Get API key: GET /api/v1/auths/api_key (with Bearer token from login)
       - Or create: POST /api/v1/auths/api_key

    3. Using JWT Token (alternative):
       - Login: POST /api/v1/auths/signin
       - Use the "token" from response as Bearer token
    """
    import os

    api_url = os.getenv("OPEN_WEBUI_API_URL", "https://your-open-webui-instance.com/api/v1")
    # Default token is optional - each connection can provide its own via env var
    api_token = os.getenv("OPEN_WEBUI_API_TOKEN")

    # Check transport mode
    transport_mode = os.getenv("MCP_TRANSPORT", "stdio").lower()  # stdio or http
    http_host = os.getenv("MCP_HTTP_HOST", "0.0.0.0")
    http_port = int(os.getenv("MCP_HTTP_PORT", "8001"))

    server = KnowledgeMCPHTTPWrapper(
        api_base_url=api_url, api_token=api_token  # Optional - can be None for multi-user mode
    )

    if transport_mode == "http":
        print("Starting Knowledge MCP Server in HTTP mode")
        print(f"Connecting to Open WebUI API: {api_url}")
        if not api_token:
            print("Note: No default API token provided.")
            print("Each client should provide their token via Authorization header.")
        print("\nServer will be available at:")
        print(f"  MCP Endpoint: http://{http_host}:{http_port}/mcp")
        print(f"  Health Check: http://{http_host}:{http_port}/health")
        print("\nClients should connect with:")
        print(f"  POST http://{http_host}:{http_port}/mcp")
        print("  Header: Authorization: Bearer sk-your-token")
        await server.run_http(host=http_host, port=http_port)
    else:
        print("Starting Knowledge MCP Server in stdio mode")
        print(f"Connecting to Open WebUI API: {api_url}")
        if not api_token:
            print("Note: No default API token provided.")
            print("Each connection should set OPEN_WEBUI_API_TOKEN environment variable.")
            print("\nTo get your API token:")
            print("1. Via Web UI: Settings → Account → API keys → Show")
            print("2. Via API: POST /api/v1/auths/api_key (requires login)")
            print("3. Set environment variable: export OPEN_WEBUI_API_TOKEN='sk-...'")
        await server.run_stdio()


def _cli_entry():
    """CLI entry point for the owui-mcp-server command."""
    asyncio.run(main())


if __name__ == "__main__":
    _cli_entry()
