"""Generator for FastAPI handlers for LoadBalancerSlsResource endpoints."""

import ast
import logging
from pathlib import Path
from typing import Any, Dict, List, Union

from runpod_flash.runtime.models import Manifest

logger = logging.getLogger(__name__)

LB_HANDLER_TEMPLATE = '''"""
Auto-generated FastAPI handler for LoadBalancerSlsResource: {resource_name}
Generated at: {timestamp}

This file is generated by the Flash build process. Do not edit manually.

Load-balanced endpoints expose HTTP servers directly to clients, enabling:
- REST APIs with custom HTTP routing
- WebSocket servers
- Real-time communication patterns
"""

import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI
from runpod_flash.runtime.lb_handler import create_lb_handler

logger = logging.getLogger(__name__)

# Import all functions/classes that belong to this resource
{imports}

# Route registry: (method, path) -> function
ROUTE_REGISTRY = {{
{registry}
}}


# Lifespan context manager for startup/shutdown
@asynccontextmanager
async def lifespan(app: FastAPI):
    """Handle application startup and shutdown."""
    logger.info("Starting {resource_name} endpoint")
    yield
    logger.info("Shutting down {resource_name} endpoint")


# Create FastAPI app with routes and lifespan
# Note: include_execute={include_execute} for this endpoint type
# - LiveLoadBalancer (local): include_execute=True for /execute endpoint
# - LoadBalancerSlsResource (deployed): include_execute=False (security)
app = create_lb_handler(ROUTE_REGISTRY, include_execute={include_execute}, lifespan=lifespan)


# Health check endpoint (required for RunPod load-balancer endpoints)
@app.get("/ping")
def ping():
    """Health check endpoint for RunPod load-balancer.

    Returns:
        dict: Status response
    """
    return {{"status": "healthy"}}


if __name__ == "__main__":
    import uvicorn
    # Local development server for testing
    uvicorn.run(app, host="0.0.0.0", port=8000)
'''


class LBHandlerGenerator:
    """Generates FastAPI handlers for LoadBalancerSlsResource endpoints."""

    def __init__(self, manifest: Union[Dict[str, Any], Manifest], build_dir: Path):
        self.manifest = manifest
        self.build_dir = build_dir

    def generate_handlers(self) -> List[Path]:
        """Generate all LB handler files."""
        handler_paths = []

        # Handle both dict and Manifest types
        resources = (
            self.manifest.resources
            if isinstance(self.manifest, Manifest)
            else self.manifest.get("resources", {})
        )

        for resource_name, resource_data in resources.items():
            # Generate for both LiveLoadBalancer (local dev) and LoadBalancerSlsResource (deployed)
            # Use flag determined by isinstance() at scan time
            is_load_balanced = (
                resource_data.is_load_balanced
                if hasattr(resource_data, "is_load_balanced")
                else resource_data.get("is_load_balanced", False)
            )
            if not is_load_balanced:
                continue

            handler_path = self._generate_handler(resource_name, resource_data)
            handler_paths.append(handler_path)

        return handler_paths

    def _generate_handler(self, resource_name: str, resource_data: Any) -> Path:
        """Generate a single FastAPI handler file."""
        handler_filename = f"handler_{resource_name}.py"
        handler_path = self.build_dir / handler_filename

        # Get timestamp from manifest
        timestamp = (
            self.manifest.generated_at
            if isinstance(self.manifest, Manifest)
            else self.manifest.get("generated_at", "")
        )

        # Determine if /execute endpoint should be included
        # LiveLoadBalancer (local dev) includes /execute, deployed LoadBalancerSlsResource does not
        # Use flag determined by isinstance() at scan time
        include_execute = (
            resource_data.is_live_resource
            if hasattr(resource_data, "is_live_resource")
            else resource_data.get("is_live_resource", False)
        )

        # Get functions from resource (handle both dict and ResourceConfig)
        functions = (
            resource_data.functions
            if hasattr(resource_data, "functions")
            else resource_data.get("functions", [])
        )

        # Generate imports section
        imports = self._generate_imports(functions)

        # Generate route registry
        registry = self._generate_route_registry(functions)

        # Format template
        handler_code = LB_HANDLER_TEMPLATE.format(
            resource_name=resource_name,
            timestamp=timestamp,
            imports=imports,
            registry=registry,
            include_execute=str(include_execute),
        )

        handler_path.write_text(handler_code)

        # Validate that generated handler can be imported
        self._validate_handler_imports(handler_path)

        return handler_path

    def _generate_imports(self, functions: List[Any]) -> str:
        """Generate import statements for functions.

        Uses importlib to handle module paths with any characters,
        including numeric prefixes that aren't valid Python identifiers.

        Args:
            functions: List of function metadata (dicts or FunctionMetadata objects)

        Returns:
            Import statements as string
        """
        if not functions:
            return "# No functions to import"

        imports = ["import importlib"]

        for func in functions:
            # Handle both dict and FunctionMetadata
            module = func.module if hasattr(func, "module") else func.get("module")
            name = func.name if hasattr(func, "name") else func.get("name")

            if module and name:
                # Use importlib to handle module names with invalid identifiers
                imports.append(f"{name} = importlib.import_module('{module}').{name}")

        return "\n".join(imports)

    def _generate_route_registry(self, functions: List[Any]) -> str:
        """Generate route registry for FastAPI app.

        Creates mapping of (method, path) tuples to function names.

        Args:
            functions: List of function metadata dicts with http_method and http_path

        Returns:
            Registry dictionary as string
        """
        if not functions:
            return "    # No functions registered"

        registry_lines = []

        for func in functions:
            # Handle both dict and FunctionMetadata
            name = func.name if hasattr(func, "name") else func.get("name")
            method = (
                func.http_method
                if hasattr(func, "http_method")
                else func.get("http_method")
            )
            path = (
                func.http_path if hasattr(func, "http_path") else func.get("http_path")
            )

            if name and method and path:
                # Create tuple key: ("GET", "/api/process")
                registry_lines.append(f'    ("{method}", "{path}"): {name},')
            elif name:
                # Skip if method or path missing (shouldn't happen with validation)
                logger.warning(
                    f"Function '{name}' missing http_method or http_path. Skipping."
                )

        return (
            "\n".join(registry_lines)
            if registry_lines
            else "    # No routes registered"
        )

    def _validate_handler_imports(self, handler_path: Path) -> None:
        """Validate that generated handler has valid Python syntax.

        Uses ast.parse() to check syntax without executing the module.
        This avoids ImportErrors from modules that only resolve at runtime
        inside Docker (e.g., numeric-prefixed module paths).

        Args:
            handler_path: Path to generated handler file

        Raises:
            ValueError: If handler has syntax errors
        """
        try:
            source = handler_path.read_text(encoding="utf-8")
            ast.parse(source, filename=str(handler_path))
        except SyntaxError as e:
            raise ValueError(f"Handler has syntax errors: {e}") from e
