#!/usr/bin/env python3
"""
Ollama Context Proxy - Single port with URL-based context routing + auto-sizing
Use URLs like: http://localhost:11434/proxy-context/4096/api/generate
Or auto-sizing: http://localhost:11434/proxy-context/auto/api/generate
"""

import asyncio
import json
import logging
import os
import re
import urllib.parse
from typing import Optional, Union
import aiohttp
from aiohttp import web, ClientSession
from aiohttp.web_response import StreamResponse
import argparse
import sys


class OllamaContextProxy:
    def __init__(
        self,
        ollama_host: Optional[str] = None,
        ollama_port: int = 11434,
        proxy_port: int = 11434,
    ):
        # Use OLLAMA_BASE_URL environment variable or construct from host/port
        base_url = os.getenv("OLLAMA_BASE_URL")
        if base_url:
            self.ollama_base_url = base_url.rstrip("/")
        else:
            # Fall back to host/port construction
            if ollama_host is None:
                ollama_host = "localhost"
            self.ollama_base_url = f"http://{ollama_host}:{ollama_port}"

        self.proxy_port = proxy_port
        self.session: Optional[ClientSession] = None
        self.logger = logging.getLogger(__name__)

        # Available context sizes (must be sorted ascending)
        self.available_contexts = [2048, 4096, 8192, 16384, 32768]

        # URL pattern to extract context size or 'auto'
        self.context_pattern = re.compile(r"^/proxy-context/(auto|\d+)(/.*)?$")

    async def start(self):
        """Initialize the HTTP session"""
        self.session = ClientSession()

    async def stop(self):
        """Cleanup HTTP session"""
        if self.session:
            await self.session.close()

    def create_app(self) -> web.Application:
        """Create the main web application"""
        app = web.Application()
        app["proxy"] = self

        # Add routes - capture everything under /proxy-context/
        app.router.add_route(
            "*",
            r"/proxy-context/{context_spec:(auto|\d+)}{path:.*}",
            self.proxy_handler,
        )

        # Optional: Add a health check endpoint
        app.router.add_get("/", self.health_check)
        app.router.add_get("/health", self.health_check)

        return app

    async def health_check(self, request: web.Request) -> web.Response:
        """Health check endpoint"""
        return web.Response(
            text="Ollama Context Proxy is running\n"
            "Usage: /proxy-context/{context_size}/api/{endpoint}\n"
            "       /proxy-context/auto/api/{endpoint}\n"
            "Examples:\n"
            "  Fixed:  /proxy-context/4096/api/generate\n"
            "  Auto:   /proxy-context/auto/api/generate\n"
            f"Available contexts: {', '.join(map(str, self.available_contexts))}",
            content_type="text/plain",
        )

    async def proxy_handler(self, request: web.Request) -> web.Response:
        """Handle all proxy requests with context size extraction or auto-detection"""

        # Extract context spec and remaining path
        context_spec = request.match_info["context_spec"]
        remaining_path = request.match_info.get("path", "")

        # Remove leading slash if present
        if remaining_path.startswith("/"):
            remaining_path = remaining_path[1:]

        # Get request data first (needed for auto-sizing)
        if request.content_type == "application/json":
            try:
                data = await request.json()
            except json.JSONDecodeError:
                data = await request.text()
        else:
            data = await request.read()

        # Determine context size
        if context_spec == "auto":
            context_size = self._auto_determine_context_size(data, remaining_path)
        else:
            context_size = int(context_spec)

        # Validate context size
        if context_size not in self.available_contexts:
            # Find the next larger available context
            suitable_context = next(
                (ctx for ctx in self.available_contexts if ctx >= context_size),
                self.available_contexts[-1],
            )
            self.logger.warning(
                f"Requested context {context_size} not available, using {suitable_context}"
            )
            context_size = suitable_context

        # Build target URL
        if not remaining_path:
            target_url = self.ollama_base_url
        else:
            target_url = f"{self.ollama_base_url}/{remaining_path}"

        self.logger.info(f"Routing to context {context_size} -> {target_url}")

        # Inject context if needed
        if self._should_inject_context(remaining_path) and isinstance(data, dict):
            if "options" not in data:
                data["options"] = {}
            data["options"]["num_ctx"] = context_size
            self.logger.info(f"Injected num_ctx={context_size} for {remaining_path}")

        # Prepare headers (exclude hop-by-hop headers)
        headers = {
            key: value
            for key, value in request.headers.items()
            if key.lower() not in ["host", "connection", "upgrade"]
        }

        if not self.session:
            raise RuntimeError("HTTP session not initialized")
        try:
            # Make request to Ollama
            async with self.session.request(
                method=request.method,
                url=target_url,
                data=json.dumps(data) if isinstance(data, dict) else data,
                headers=headers,
                params=request.query,
            ) as response:
                # Handle streaming responses (for generate/chat endpoints)
                if response.headers.get("content-type", "").startswith(
                    "application/x-ndjson"
                ):
                    return await self._handle_streaming_response(request, response)
                else:
                    return await self._handle_regular_response(response)

        except aiohttp.ClientError as e:
            self.logger.error(f"Error proxying request to {target_url}: {e}")
            return web.Response(
                text=f"Proxy error: {str(e)}", status=502, content_type="text/plain"
            )

    def _auto_determine_context_size(
        self, data: Union[dict, str, bytes], endpoint: str
    ) -> int:
        """Automatically determine the required context size based on request content"""

        input_tokens = 0
        max_tokens = 0

        if isinstance(data, dict):
            # Extract text content and max_tokens based on endpoint
            if endpoint.startswith("api/generate"):
                # Ollama generate endpoint
                prompt = data.get("prompt", "")
                input_tokens = self._estimate_tokens(prompt)
                max_tokens = data.get("options", {}).get("num_predict", 0)

            elif endpoint.startswith("api/chat"):
                # Ollama chat endpoint
                messages = data.get("messages", [])
                total_text = ""
                for msg in messages:
                    if isinstance(msg, dict) and "content" in msg:
                        total_text += str(msg["content"]) + " "
                input_tokens = self._estimate_tokens(total_text)
                max_tokens = data.get("options", {}).get("num_predict", 0)

            elif endpoint.startswith("v1/chat/completions"):
                # OpenAI-compatible chat endpoint
                messages = data.get("messages", [])
                total_text = ""
                for msg in messages:
                    if isinstance(msg, dict) and "content" in msg:
                        total_text += str(msg["content"]) + " "
                input_tokens = self._estimate_tokens(total_text)
                max_tokens = data.get("max_tokens", 0)

            elif endpoint.startswith("v1/completions"):
                # OpenAI-compatible completions endpoint
                prompt = data.get("prompt", "")
                input_tokens = self._estimate_tokens(prompt)
                max_tokens = data.get("max_tokens", 0)

        elif isinstance(data, (str, bytes)):
            # Fallback for non-JSON data
            text = (
                data if isinstance(data, str) else data.decode("utf-8", errors="ignore")
            )
            input_tokens = self._estimate_tokens(text)

        # Calculate total tokens needed
        system_overhead = 100  # Buffer for system prompts, formatting, etc.
        response_buffer = max(max_tokens, 512)  # Ensure space for response
        safety_margin = 200  # Additional safety buffer

        total_needed = input_tokens + response_buffer + system_overhead + safety_margin

        # Find the smallest context that can accommodate the request
        suitable_context = next(
            (ctx for ctx in self.available_contexts if ctx >= total_needed),
            self.available_contexts[-1],  # Fall back to largest if none are big enough
        )

        self.logger.info(
            f"Auto-sizing analysis: "
            f"input_tokens={input_tokens}, "
            f"max_tokens={max_tokens}, "
            f"total_needed={total_needed}, "
            f"selected_context={suitable_context}"
        )

        # Log warning if we're using the largest context and it might not be enough
        if (
            suitable_context == self.available_contexts[-1]
            and total_needed > suitable_context
        ):
            self.logger.warning(
                f"Request may exceed largest available context! "
                f"Needed: {total_needed}, Available: {suitable_context}"
            )

        return suitable_context

    def _estimate_tokens(self, text: str) -> int:
        """Estimate token count from text (rough approximation)"""
        if not text:
            return 0

        # Rough estimation: ~4 characters per token for English
        # This is a conservative estimate - actual tokenization varies by model
        char_count = len(str(text))
        estimated_tokens = max(1, char_count // 4)

        self.logger.debug(
            f"Token estimation: {char_count} chars -> ~{estimated_tokens} tokens"
        )
        return estimated_tokens

    def _should_inject_context(self, path: str) -> bool:
        """Determine if we should inject context for this endpoint"""
        # Inject context for endpoints that support the num_ctx parameter
        context_endpoints = [
            "api/generate",
            "api/chat",
            "v1/chat/completions",
            "v1/completions",
        ]
        return any(path.startswith(endpoint) for endpoint in context_endpoints)

    async def _handle_streaming_response(
        self, request: web.Request, response: aiohttp.ClientResponse
    ) -> StreamResponse:
        """Handle streaming responses (NDJSON)"""
        stream_response = StreamResponse(
            status=response.status,
            headers={
                key: value
                for key, value in response.headers.items()
                if key.lower() not in ["content-length", "transfer-encoding"]
            },
        )

        await stream_response.prepare(request)

        async for chunk in response.content.iter_any():
            await stream_response.write(chunk)

        await stream_response.write_eof()
        return stream_response

    async def _handle_regular_response(
        self, response: aiohttp.ClientResponse
    ) -> web.Response:
        """Handle regular (non-streaming) responses"""
        content = await response.read()

        return web.Response(
            body=content,
            status=response.status,
            headers={
                key: value
                for key, value in response.headers.items()
                if key.lower() not in ["content-length", "transfer-encoding"]
            },
        )


async def main():
    parser = argparse.ArgumentParser(
        description="Ollama Context Proxy - URL-based routing with auto-sizing"
    )

    # Get default host from OLLAMA_BASE_URL if available
    default_host = "localhost"
    base_url = os.getenv("OLLAMA_BASE_URL")
    if base_url:
        # Extract host from base URL for backward compatibility with CLI args
        parsed = urllib.parse.urlparse(base_url)
        if parsed.hostname:
            default_host = parsed.hostname

    parser.add_argument(
        "--ollama-host",
        default=default_host,
        help=f"Ollama server host (default: {default_host})",
    )
    parser.add_argument(
        "--ollama-port",
        type=int,
        default=11434,
        help="Ollama server port (default: 11434)",
    )
    parser.add_argument(
        "--proxy-port",
        type=int,
        default=11435,
        help="Proxy server port (default: 11435)",
    )
    parser.add_argument(
        "--log-level",
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        help="Log level (default: INFO)",
    )

    args = parser.parse_args()

    # Setup logging
    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )

    # Create proxy instance
    proxy = OllamaContextProxy(args.ollama_host, args.ollama_port, args.proxy_port)
    await proxy.start()

    # Create and start the web application
    app = proxy.create_app()
    runner = web.AppRunner(app)
    await runner.setup()

    site = web.TCPSite(runner, "0.0.0.0", args.proxy_port)
    await site.start()

    logging.info(f"Ollama Context Proxy started on port {args.proxy_port}")
    logging.info(f"Forwarding to Ollama at {proxy.ollama_base_url}")
    logging.info(f"Available context sizes: {proxy.available_contexts}")
    logging.info("Usage examples:")
    logging.info(
        f"  Auto-size:   http://localhost:{args.proxy_port}/proxy-context/auto"
    )
    logging.info(
        f"  2K context:  http://localhost:{args.proxy_port}/proxy-context/2048"
    )
    logging.info(
        f"  4K context:  http://localhost:{args.proxy_port}/proxy-context/4096"
    )
    logging.info(
        f"  8K context:  http://localhost:{args.proxy_port}/proxy-context/8192"
    )
    logging.info(
        f"  16K context: http://localhost:{args.proxy_port}/proxy-context/16384"
    )
    logging.info(
        f"  32K context: http://localhost:{args.proxy_port}/proxy-context/32768"
    )

    try:
        # Keep running
        while True:
            await asyncio.sleep(1)

    except KeyboardInterrupt:
        logging.info("Shutting down...")
    finally:
        # Cleanup
        await runner.cleanup()
        await proxy.stop()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nShutdown complete.")
        sys.exit(0)