"""Factory for creating FastAPI load-balanced handlers.

This module provides the factory function for generating FastAPI applications
that handle load-balanced serverless endpoints. It supports:
- User-defined HTTP routes
- /execute endpoint for @remote function execution (LiveLoadBalancer only)

Security Model:
    The /execute endpoint accepts and executes serialized function code. This is
    secure because:
    1. The function code originates from the client's @remote decorator
    2. The client (user) controls what function gets sent
    3. This mirrors the trusted client model of LiveServerlessStub
    4. In production, API authentication should protect the /execute endpoint

    Users should NOT expose the /execute endpoint to untrusted clients.
"""

import asyncio
import inspect
import logging
import re
from typing import Any, Callable, Dict, get_type_hints

from fastapi import FastAPI, File, Form, Request
from pydantic import BaseModel, create_model

logger = logging.getLogger(__name__)

_BODY_METHODS = frozenset({"POST", "PUT", "PATCH", "DELETE"})
_PATH_PARAM_RE = re.compile(r"\{(\w+)\}")
_SKIP_KINDS = (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)


def _make_input_model(
    name: str, func: Callable, exclude: set[str] | None = None
) -> type | None:
    """Create a Pydantic model from a function's signature for FastAPI body typing.

    Returns None for zero-param functions or on introspection failure.
    """
    exclude = exclude or set()
    try:
        sig = inspect.signature(func)
        hints = get_type_hints(func)
    except (ValueError, TypeError) as e:
        logger.warning(
            "Failed to introspect signature for %s: %s. "
            "Skipping input model generation.",
            name,
            e,
        )
        return None

    fields: dict[str, Any] = {}
    for param_name, param in sig.parameters.items():
        if param_name == "self" or param_name in exclude or param.kind in _SKIP_KINDS:
            continue
        annotation = hints.get(param_name, Any)
        if param.default is not inspect.Parameter.empty:
            fields[param_name] = (annotation, param.default)
        else:
            fields[param_name] = (annotation, ...)

    if not fields:
        return None

    return create_model(name, **fields)


def _build_file_upload_wrapper(
    handler: Callable,
    file_params: list[tuple[str, type, Any]],
    form_params: list[tuple[str, type, Any]],
    path_params: set[str],
    hints: dict[str, Any],
    sig: inspect.Signature,
) -> Callable:
    """Build a FastAPI-compatible wrapper with File() and Form() annotations.

    Constructs a wrapper function whose signature uses File(...) for bytes
    params and Form(...) for non-file params, enabling multipart upload
    via FastAPI's native support.
    """
    is_async = asyncio.iscoroutinefunction(handler)

    params: list[inspect.Parameter] = []
    annotations: dict[str, Any] = {}

    # Path params first (no default annotation needed)
    for pname in path_params:
        if pname in sig.parameters:
            ann = hints.get(pname, Any)
            params.append(
                inspect.Parameter(
                    pname, inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=ann
                )
            )
            annotations[pname] = ann

    # File params with File(...)
    for pname, ann, default in file_params:
        file_default = (
            File(...) if default is inspect.Parameter.empty else File(default)
        )
        params.append(
            inspect.Parameter(
                pname,
                inspect.Parameter.POSITIONAL_OR_KEYWORD,
                default=file_default,
                annotation=bytes,
            )
        )
        annotations[pname] = bytes

    # Form params with Form(...)
    for pname, ann, default in form_params:
        form_default = (
            Form(...) if default is inspect.Parameter.empty else Form(default)
        )
        params.append(
            inspect.Parameter(
                pname,
                inspect.Parameter.POSITIONAL_OR_KEYWORD,
                default=form_default,
                annotation=ann,
            )
        )
        annotations[pname] = ann

    if is_async:

        async def wrapper(**kwargs):
            return await handler(**kwargs)
    else:

        def wrapper(**kwargs):
            return handler(**kwargs)

    wrapper.__signature__ = inspect.Signature(parameters=params)
    wrapper.__annotations__ = annotations
    wrapper.__name__ = getattr(handler, "__name__", "handler")
    wrapper.__doc__ = handler.__doc__
    return wrapper


def _wrap_handler_with_body_model(handler: Callable, path: str) -> Callable:
    """Wrap a handler so FastAPI reads its parameters from the JSON body.

    If the handler already accepts a single Pydantic BaseModel parameter,
    or has no eligible body parameters, returns it unchanged.
    """
    try:
        sig = inspect.signature(handler)
        hints = get_type_hints(handler)
    except (ValueError, TypeError) as e:
        logger.warning(
            "Failed to introspect handler %s for body model wrapping: %s. "
            "Returning handler unwrapped.",
            getattr(handler, "__name__", "unknown"),
            e,
        )
        return handler

    path_params = set(_PATH_PARAM_RE.findall(path))

    # Check if any non-path param is already a Pydantic model
    for pname, param in sig.parameters.items():
        if pname in path_params or pname == "self" or param.kind in _SKIP_KINDS:
            continue
        annotation = hints.get(pname, Any)
        if isinstance(annotation, type) and issubclass(annotation, BaseModel):
            return handler

    # Detect bytes params for file upload support
    file_params: list[tuple[str, type, Any]] = []
    form_params: list[tuple[str, type, Any]] = []
    for pname, param in sig.parameters.items():
        if pname in path_params or pname == "self" or param.kind in _SKIP_KINDS:
            continue
        annotation = hints.get(pname, Any)
        if annotation is bytes:
            file_params.append((pname, annotation, param.default))
        else:
            form_params.append((pname, annotation, param.default))

    if file_params:
        return _build_file_upload_wrapper(
            handler, file_params, form_params, path_params, hints, sig
        )

    model_name = handler.__name__.title().replace("_", "") + "Body"
    model = _make_input_model(model_name, handler, exclude=path_params)
    if model is None:
        return handler

    is_async = asyncio.iscoroutinefunction(handler)

    if path_params:
        if is_async:

            async def wrapped_with_path(body, **kwargs):  # type: ignore[valid-type]
                return await handler(**body.model_dump(), **kwargs)
        else:

            def wrapped_with_path(body, **kwargs):  # type: ignore[valid-type]
                return handler(**body.model_dump(), **kwargs)

        # Build explicit signature so FastAPI maps path params correctly
        params = [
            inspect.Parameter(
                "body", inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=model
            )
        ]
        annotations = {"body": model}
        for pname in path_params:
            if pname in sig.parameters:
                ann = hints.get(pname, Any)
                params.append(
                    inspect.Parameter(
                        pname, inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=ann
                    )
                )
                annotations[pname] = ann
        wrapped_with_path.__signature__ = inspect.Signature(parameters=params)
        wrapped_with_path.__annotations__ = annotations
        return wrapped_with_path
    else:
        if is_async:

            async def wrapped(body: model):  # type: ignore[valid-type]
                return await handler(**body.model_dump())
        else:

            def wrapped(body: model):  # type: ignore[valid-type]
                return handler(**body.model_dump())

        return wrapped


def create_lb_handler(
    route_registry: Dict[tuple[str, str], Callable],
    include_execute: bool = False,
    lifespan: Callable = None,
) -> FastAPI:
    """Create FastAPI app with routes from registry.

    Args:
        route_registry: Mapping of (HTTP_METHOD, path) -> handler_function
                       Example: {("GET", "/api/health"): health_check}
        include_execute: Whether to register /execute endpoint for @remote execution.
                        Only used for LiveLoadBalancer (local development).
                        Deployed endpoints should not expose /execute for security.
        lifespan: Optional lifespan context manager for startup/shutdown hooks.

    Returns:
        Configured FastAPI application with routes registered.
    """
    app = FastAPI(title="Flash Load-Balanced Handler", lifespan=lifespan)

    # Register /execute endpoint for @remote stub execution (if enabled)
    if include_execute:
        from .serialization import deserialize_args, deserialize_kwargs, serialize_arg

        @app.post("/execute")
        async def execute_remote_function(request: Request) -> Dict[str, Any]:
            """Framework endpoint for @remote decorator execution.

            WARNING: This endpoint is INTERNAL to the Flash framework. It should only be
            called by the @remote stub from runpod_flash.stubs.load_balancer_sls. Exposing
            this endpoint to untrusted clients could allow arbitrary code execution.

            Accepts serialized function code and arguments, executes them,
            and returns serialized result.

            Request body:
                {
                    "function_name": "process_data",
                    "function_code": "def process_data(x, y): return x + y",
                    "args": [base64_encoded_arg1, base64_encoded_arg2],
                    "kwargs": {"key": base64_encoded_value}
                }

            Returns:
                {
                    "success": true,
                    "result": base64_encoded_result
                }
                or
                {
                    "success": false,
                    "error": "error message"
                }
            """
            try:
                body = await request.json()
            except Exception as e:
                logger.error(f"Failed to parse request body: {e}")
                return {"success": False, "error": f"Invalid request body: {e}"}

            try:
                # Extract function metadata
                function_name = body.get("function_name")
                function_code = body.get("function_code")

                if not function_name or not function_code:
                    return {
                        "success": False,
                        "error": "Missing function_name or function_code in request",
                    }

                # Deserialize arguments
                try:
                    args = deserialize_args(body.get("args", []))
                    kwargs = deserialize_kwargs(body.get("kwargs", {}))
                except Exception as e:
                    logger.error(f"Failed to deserialize arguments: {e}")
                    return {
                        "success": False,
                        "error": f"Failed to deserialize arguments: {e}",
                    }

                # Execute function in isolated namespace
                namespace: Dict[str, Any] = {}
                try:
                    exec(function_code, namespace)
                except SyntaxError as e:
                    logger.error(f"Syntax error in function code: {e}")
                    return {
                        "success": False,
                        "error": f"Syntax error in function code: {e}",
                    }
                except Exception as e:
                    logger.error(f"Error executing function code: {e}")
                    return {
                        "success": False,
                        "error": f"Error executing function code: {e}",
                    }

                # Get function from namespace
                if function_name not in namespace:
                    return {
                        "success": False,
                        "error": f"Function '{function_name}' not found in executed code",
                    }

                func = namespace[function_name]

                # Execute function
                try:
                    result = func(*args, **kwargs)

                    # Handle async functions
                    if inspect.iscoroutine(result):
                        result = await result
                except Exception as e:
                    logger.error(f"Function execution failed: {e}")
                    return {
                        "success": False,
                        "error": f"Function execution failed: {e}",
                    }

                # Serialize result
                try:
                    result_b64 = serialize_arg(result)
                    return {"success": True, "result": result_b64}
                except Exception as e:
                    logger.error(f"Failed to serialize result: {e}")
                    return {
                        "success": False,
                        "error": f"Failed to serialize result: {e}",
                    }

            except Exception as e:
                logger.error(f"Unexpected error in /execute endpoint: {e}")
                return {"success": False, "error": f"Unexpected error: {e}"}

    # Register user-defined routes from registry
    for (method, path), handler in route_registry.items():
        method_upper = method.upper()

        if method_upper in _BODY_METHODS:
            handler = _wrap_handler_with_body_model(handler, path)

        if method_upper == "GET":
            app.get(path)(handler)
        elif method_upper == "POST":
            app.post(path)(handler)
        elif method_upper == "PUT":
            app.put(path)(handler)
        elif method_upper == "DELETE":
            app.delete(path)(handler)
        elif method_upper == "PATCH":
            app.patch(path)(handler)
        else:
            logger.warning(
                f"Unsupported HTTP method '{method}' for path '{path}'. Skipping."
            )

    return app
