"""HTTP server for the routing inference harness.

This module provides the main aiohttp application server with OpenAI-compatible
endpoints for load testing and SLA validation.

Supports optional conductor integration for model orchestration.
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import signal
import sys
from pathlib import Path
from typing import TYPE_CHECKING, Literal

from aiohttp import web

from partner_inference_harness.inference.router import Router, create_router
from partner_inference_harness.server.routes import create_conductor_routes, create_routes

if TYPE_CHECKING:
    from partner_inference_harness.conductor.interface import ConductorInterface


logger = logging.getLogger(__name__)

# Type alias for conductor backend selection
ConductorType = Literal["rnd1", "claude", "none"]


def _create_conductor(
    conductor_type: ConductorType,
    rnd1_url: str | None = None,
    claude_model: str | None = None,
) -> "ConductorInterface | None":
    """Create conductor instance based on type.

    Args:
        conductor_type: Type of conductor to create.
        rnd1_url: Base URL for RND1 server (if using rnd1).
        claude_model: Claude model identifier (if using claude).

    Returns:
        Configured conductor instance or None if type is 'none'.

    Raises:
        ValueError: If conductor_type is unknown.
    """
    if conductor_type == "none":
        return None

    if conductor_type == "rnd1":
        from partner_inference_harness.conductor import RND1Conductor

        return RND1Conductor(
            base_url=rnd1_url or "http://localhost:8080",
        )

    if conductor_type == "claude":
        from partner_inference_harness.conductor import ClaudeConductor

        return ClaudeConductor(
            model=claude_model or "claude-sonnet-4-20250514",
        )

    raise ValueError(f"Unknown conductor type: {conductor_type}")


def create_app(
    artifacts_dir: Path | str | None = None,
    device: str | None = None,
    warmup: bool = True,
    conductor_type: ConductorType = "none",
    rnd1_url: str | None = None,
    claude_model: str | None = None,
) -> web.Application:
    """Create the aiohttp application.

    Args:
        artifacts_dir: Directory containing model artifacts.
        device: Compute device for inference.
        warmup: Whether to run warmup passes on startup.
        conductor_type: Type of conductor backend ('rnd1', 'claude', or 'none').
        rnd1_url: Base URL for RND1 server (when conductor_type='rnd1').
        claude_model: Claude model identifier (when conductor_type='claude').

    Returns:
        Configured aiohttp Application.
    """
    # Initialize router
    logger.info("Initializing router...")
    router = create_router(
        artifacts_dir=artifacts_dir,
        device=device,
        warmup=warmup,
    )
    logger.info(
        f"Router initialized. Memory usage: {router.get_memory_usage_gb():.2f} GB"
    )

    # Initialize conductor if requested
    conductor = _create_conductor(conductor_type, rnd1_url, claude_model)
    if conductor:
        logger.info(
            f"Conductor initialized: type={conductor_type}, "
            f"backend={conductor.backend_type}"
        )

    # Create application
    app = web.Application()

    # Store router in app state
    app["router"] = router
    app["conductor"] = conductor

    # Add routes
    routes = create_routes(router)
    app.add_routes(routes)

    # Add conductor routes if conductor is configured
    if conductor:
        conductor_routes = create_conductor_routes(conductor)
        app.add_routes(conductor_routes)

    # Add middleware for request logging
    @web.middleware
    async def logging_middleware(
        request: web.Request,
        handler: web.RequestHandler,
    ) -> web.StreamResponse:
        """Log requests and responses."""
        start_time = asyncio.get_event_loop().time()
        try:
            response = await handler(request)
            elapsed = (asyncio.get_event_loop().time() - start_time) * 1000
            logger.info(
                f"{request.method} {request.path} - {response.status} - {elapsed:.1f}ms"
            )
            return response
        except Exception as e:
            elapsed = (asyncio.get_event_loop().time() - start_time) * 1000
            logger.error(
                f"{request.method} {request.path} - ERROR - {elapsed:.1f}ms - {e}"
            )
            raise

    app.middlewares.append(logging_middleware)

    # Cleanup on shutdown
    async def cleanup(app: web.Application) -> None:
        """Cleanup resources on shutdown."""
        logger.info("Shutting down router...")
        router = app["router"]
        router.unload()
        logger.info("Router shutdown complete")

    app.on_cleanup.append(cleanup)

    return app


def run_server(
    host: str = "0.0.0.0",
    port: int = 8082,
    artifacts_dir: Path | str | None = None,
    device: str | None = None,
    warmup: bool = True,
    conductor_type: ConductorType = "none",
    rnd1_url: str | None = None,
    claude_model: str | None = None,
) -> None:
    """Run the HTTP server.

    Args:
        host: Host address to bind to.
        port: Port number to listen on.
        artifacts_dir: Directory containing model artifacts.
        device: Compute device for inference.
        warmup: Whether to run warmup passes on startup.
        conductor_type: Type of conductor backend ('rnd1', 'claude', or 'none').
        rnd1_url: Base URL for RND1 server (when conductor_type='rnd1').
        claude_model: Claude model identifier (when conductor_type='claude').
    """
    # Setup logging
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )

    logger.info(f"Starting server on {host}:{port}")
    if conductor_type != "none":
        logger.info(f"Conductor enabled: {conductor_type}")

    app = create_app(
        artifacts_dir=artifacts_dir,
        device=device,
        warmup=warmup,
        conductor_type=conductor_type,
        rnd1_url=rnd1_url,
        claude_model=claude_model,
    )

    # Handle graceful shutdown
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    def handle_signal(sig: signal.Signals) -> None:
        logger.info(f"Received signal {sig.name}, shutting down...")
        loop.stop()

    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, lambda s=sig: handle_signal(s))

    web.run_app(app, host=host, port=port, loop=loop)


def main() -> int:
    """CLI entry point."""
    parser = argparse.ArgumentParser(
        description="Run the routing inference harness server"
    )
    parser.add_argument(
        "--host",
        default="0.0.0.0",
        help="Host address to bind to (default: 0.0.0.0)",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=8082,
        help="Port number to listen on (default: 8082)",
    )
    parser.add_argument(
        "--artifacts-dir",
        type=Path,
        help="Directory containing model artifacts",
    )
    parser.add_argument(
        "--device",
        choices=["cuda", "cpu"],
        help="Compute device (default: cuda if available)",
    )
    parser.add_argument(
        "--no-warmup",
        action="store_true",
        help="Skip warmup inference passes",
    )

    # Conductor configuration
    parser.add_argument(
        "--conductor",
        choices=["rnd1", "claude", "none"],
        default="none",
        help="Conductor backend type (default: none)",
    )
    parser.add_argument(
        "--rnd1-url",
        default="http://localhost:8080",
        help="RND1 server URL (default: http://localhost:8080)",
    )
    parser.add_argument(
        "--claude-model",
        default="claude-sonnet-4-20250514",
        help="Claude model identifier (default: claude-sonnet-4-20250514)",
    )

    args = parser.parse_args()

    try:
        run_server(
            host=args.host,
            port=args.port,
            artifacts_dir=args.artifacts_dir,
            device=args.device,
            warmup=not args.no_warmup,
            conductor_type=args.conductor,
            rnd1_url=args.rnd1_url,
            claude_model=args.claude_model,
        )
        return 0
    except Exception as e:
        logger.exception(f"Server error: {e}")
        return 1


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