import argparse
import functools
import inspect
import json
import logging
import os
import posixpath
from typing import Any, AsyncGenerator, Callable, Literal, ParamSpec, TypeVar

import httpx
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import Response, StreamingResponse

import mlflow
from mlflow.genai.agent_server.utils import get_request_headers, set_request_headers
from mlflow.genai.agent_server.validator import BaseAgentValidator, ResponsesAgentValidator
from mlflow.pyfunc import ResponsesAgent
from mlflow.tracing.constant import SpanAttributeKey

logger = logging.getLogger(__name__)
STREAM_KEY = "stream"
RETURN_TRACE_HEADER = "x-mlflow-return-trace-id"

AgentType = Literal["ResponsesAgent"]

_P = ParamSpec("_P")
_R = TypeVar("_R")

_invoke_function: Callable[..., Any] | None = None
_stream_function: Callable[..., Any] | None = None


def get_invoke_function():
    return _invoke_function


def get_stream_function():
    return _stream_function


def invoke() -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
    """Decorator to register a function as an invoke endpoint. Can only be used once."""

    def decorator(func: Callable[_P, _R]) -> Callable[_P, _R]:
        global _invoke_function
        if _invoke_function is not None:
            raise ValueError("invoke decorator can only be used once")
        _invoke_function = func

        @functools.wraps(func)
        def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
            return func(*args, **kwargs)

        return wrapper

    return decorator


def stream() -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
    """Decorator to register a function as a stream endpoint. Can only be used once."""

    def decorator(func: Callable[_P, _R]) -> Callable[_P, _R]:
        global _stream_function
        if _stream_function is not None:
            raise ValueError("stream decorator can only be used once")
        _stream_function = func

        @functools.wraps(func)
        def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
            return func(*args, **kwargs)

        return wrapper

    return decorator


class AgentServer:
    """FastAPI-based server for hosting agents.

    Args:
        agent_type: An optional parameter to specify the type of agent to serve. If provided,
            input/output validation and streaming tracing aggregation will be done automatically.

            Currently only "ResponsesAgent" is supported.

            If ``None``, no input/output validation and streaming tracing aggregation will be done.
            Default to ``None``.

        enable_chat_proxy: If ``True``, enables a proxy middleware that forwards unmatched requests
            to a chat app running on the port specified by the CHAT_APP_PORT environment variable
            (defaults to 3000) with a timeout specified by the
            CHAT_PROXY_TIMEOUT_SECONDS environment variable, (defaults to 300 seconds).
            ``enable_chat_proxy`` defaults to ``False``.

            The proxy allows requests to ``/``, ``/favicon.ico``, ``/assets/*``, and ``/api/*`` by
            default. Additional paths can be configured via environment variables:

            - ``CHAT_PROXY_ALLOWED_EXACT_PATHS``: Comma-separated list of additional exact paths
              to allow (e.g., ``/custom,/another``).
            - ``CHAT_PROXY_ALLOWED_PATH_PREFIXES``: Comma-separated list of additional path prefixes
              to allow (e.g., ``/custom/,/another/``).

    See https://mlflow.org/docs/latest/genai/serving/agent-server for more information.
    """

    def __init__(self, agent_type: AgentType | None = None, enable_chat_proxy: bool = False):
        self.agent_type = agent_type
        if agent_type == "ResponsesAgent":
            self.validator = ResponsesAgentValidator()
        else:
            self.validator = BaseAgentValidator()

        self.app = FastAPI(title="Agent Server")

        if enable_chat_proxy:
            self._setup_chat_proxy_middleware()

        self._setup_routes()

    def _setup_chat_proxy_middleware(self) -> None:
        """Set up middleware to proxy static asset requests to the chat app.

        Only forwards requests to allowed paths (/, /assets/*, /favicon.ico) to prevent
        SSRF vulnerabilities.
        """
        self.chat_app_port = os.getenv("CHAT_APP_PORT", "3000")
        self.chat_proxy_timeout = float(os.getenv("CHAT_PROXY_TIMEOUT_SECONDS", "300.0"))
        self.proxy_client = httpx.AsyncClient(timeout=self.chat_proxy_timeout)

        # Only proxy static assets to prevent SSRF vulnerabilities
        allowed_exact_paths = {"/", "/favicon.ico"}
        allowed_path_prefixes = {"/assets/", "/api/"}

        # Add additional paths from environment variables
        if additional_exact_paths := os.getenv("CHAT_PROXY_ALLOWED_EXACT_PATHS", ""):
            allowed_exact_paths |= {
                p.strip() for p in additional_exact_paths.split(",") if p.strip()
            }

        if additional_path_prefixes := os.getenv("CHAT_PROXY_ALLOWED_PATH_PREFIXES", ""):
            allowed_path_prefixes |= {
                p.strip() for p in additional_path_prefixes.split(",") if p.strip()
            }

        @self.app.middleware("http")
        async def chat_proxy_middleware(request: Request, call_next):
            """
            Forward static asset requests to the chat app on the port specified by the
            CHAT_APP_PORT environment variable (defaults to 3000).

            Only forwards requests to:
            - / (base path for index.html)
            - /assets/* (Vite static assets)
            - /favicon.ico

            The timeout for the proxy request is specified by the CHAT_PROXY_TIMEOUT_SECONDS
            environment variable (defaults to 300.0 seconds).

            For streaming responses (SSE), the proxy streams chunks as they arrive
            rather than buffering the entire response.
            """
            for route in self.app.routes:
                if hasattr(route, "path_regex") and route.path_regex.match(request.url.path):
                    return await call_next(request)

            # Normalize path to prevent traversal attacks (e.g., /assets/../.env)
            path = posixpath.normpath(request.url.path)

            # Only allow proxying static assets
            is_allowed = path in allowed_exact_paths or any(
                path.startswith(p) for p in allowed_path_prefixes
            )
            if not is_allowed:
                return Response("Not found", status_code=404, media_type="text/plain")

            path = path.lstrip("/")
            try:
                body = await request.body() if request.method in ["POST", "PUT", "PATCH"] else None
                target_url = f"http://localhost:{self.chat_app_port}/{path}"

                # Build and send request with streaming enabled
                req = self.proxy_client.build_request(
                    method=request.method,
                    url=target_url,
                    params=dict(request.query_params),
                    headers={k: v for k, v in request.headers.items() if k.lower() != "host"},
                    content=body,
                )
                proxy_response = await self.proxy_client.send(req, stream=True)

                # Check if this is a streaming response (SSE)
                content_type = proxy_response.headers.get("content-type", "")
                if "text/event-stream" in content_type:
                    # Stream SSE responses chunk by chunk
                    async def stream_generator():
                        try:
                            async for chunk in proxy_response.aiter_bytes():
                                yield chunk
                        except Exception:
                            logger.exception("Error in chat proxy streaming")
                            raise
                        finally:
                            await proxy_response.aclose()

                    return StreamingResponse(
                        stream_generator(),
                        status_code=proxy_response.status_code,
                        headers=dict(proxy_response.headers),
                    )
                else:
                    # Non-streaming response - read fully then close
                    content = await proxy_response.aread()
                    await proxy_response.aclose()
                    return Response(
                        content,
                        proxy_response.status_code,
                        headers=dict(proxy_response.headers),
                    )
            except httpx.ConnectError:
                return Response("Service unavailable", status_code=503, media_type="text/plain")
            except Exception as e:
                return Response(f"Proxy error: {e!s}", status_code=502, media_type="text/plain")

    def _setup_routes(self) -> None:
        @self.app.post("/invocations")
        async def invocations_endpoint(request: Request):
            return await self._handle_invocations_request(request)

        # Only expose /responses endpoint for ResponsesAgent
        if self.agent_type == "ResponsesAgent":

            @self.app.post("/responses")
            async def responses_endpoint(request: Request):
                """
                For compatibility with the OpenAI Client `client.responses.create(...)` method.
                https://platform.openai.com/docs/api-reference/responses/create
                """
                return await self._handle_invocations_request(request)

        @self.app.get("/agent/info")
        async def agent_info_endpoint() -> dict[str, Any]:
            # Get app name from environment or use default
            app_name = os.getenv("DATABRICKS_APP_NAME", "mlflow_agent_server")

            # Base info payload
            info = {
                "name": app_name,
                "use_case": "agent",
                "mlflow_version": mlflow.__version__,
            }

            # Conditionally add agent_api field for ResponsesAgent only
            if self.agent_type == "ResponsesAgent":
                info["agent_api"] = "responses"

            return info

        @self.app.get("/health")
        async def health_check() -> dict[str, str]:
            return {"status": "healthy"}

    async def _handle_invocations_request(
        self, request: Request
    ) -> dict[str, Any] | StreamingResponse:
        # Capture headers such as x-forwarded-access-token
        # https://docs.databricks.com/aws/en/dev-tools/databricks-apps/auth?language=Streamlit#retrieve-user-authorization-credentials
        set_request_headers(dict(request.headers))

        try:
            data = await request.json()
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid JSON in request body: {e!s}")

        # Use actual request path for logging differentiation
        endpoint_path = request.url.path
        logger.debug(
            f"Request received at {endpoint_path}",
            extra={
                "agent_type": self.agent_type,
                "request_size": len(json.dumps(data)),
                "stream_requested": data.get(STREAM_KEY, False),
                "endpoint": endpoint_path,
            },
        )

        is_streaming = data.pop(STREAM_KEY, False)
        return_trace_id = (get_request_headers().get(RETURN_TRACE_HEADER) or "").lower() == "true"

        try:
            request_data = self.validator.validate_and_convert_request(data)
        except ValueError as e:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid parameters for {self.agent_type}: {e}",
            )

        if is_streaming:
            return await self._handle_stream_request(request_data, return_trace_id)
        else:
            return await self._handle_invoke_request(request_data, return_trace_id)

    async def _handle_invoke_request(
        self, request: dict[str, Any], return_trace_id: bool
    ) -> dict[str, Any]:
        if _invoke_function is None:
            raise HTTPException(status_code=500, detail="No invoke function registered")

        func = _invoke_function
        func_name = func.__name__

        try:
            with mlflow.start_span(name=f"{func_name}") as span:
                span.set_inputs(request)
                if inspect.iscoroutinefunction(func):
                    result = await func(request)
                else:
                    result = func(request)

                result = self.validator.validate_and_convert_result(result)
                if self.agent_type == "ResponsesAgent":
                    span.set_attribute(SpanAttributeKey.MESSAGE_FORMAT, "openai")
                    if return_trace_id:
                        result["metadata"] = (result.get("metadata") or {}) | {
                            "trace_id": span.trace_id
                        }
                span.set_outputs(result)

            logger.debug(
                "Response sent",
                extra={
                    "endpoint": "invoke",
                    "response_size": len(json.dumps(result)),
                    "function_name": func_name,
                },
            )

            return result

        except Exception as e:
            logger.exception("Error in invoke endpoint")

            raise HTTPException(status_code=500, detail=str(e))

    async def _generate(
        self,
        func: Callable[..., Any],
        request: dict[str, Any],
        return_trace_id: bool,
    ) -> AsyncGenerator[str, None]:
        func_name = func.__name__
        all_chunks: list[dict[str, Any]] = []
        try:
            with mlflow.start_span(name=f"{func_name}") as span:
                span.set_inputs(request)
                if inspect.iscoroutinefunction(func) or inspect.isasyncgenfunction(func):
                    async for chunk in func(request):
                        chunk = self.validator.validate_and_convert_result(chunk, stream=True)
                        all_chunks.append(chunk)
                        yield f"data: {json.dumps(chunk)}\n\n"
                else:
                    for chunk in func(request):
                        chunk = self.validator.validate_and_convert_result(chunk, stream=True)
                        all_chunks.append(chunk)
                        yield f"data: {json.dumps(chunk)}\n\n"

                if self.agent_type == "ResponsesAgent":
                    span.set_attribute(SpanAttributeKey.MESSAGE_FORMAT, "openai")
                    span.set_outputs(ResponsesAgent.responses_agent_output_reducer(all_chunks))
                    if return_trace_id:
                        yield f"data: {json.dumps({'trace_id': span.trace_id})}\n\n"
                else:
                    span.set_outputs(all_chunks)

                yield "data: [DONE]\n\n"

            logger.debug(
                "Streaming response completed",
                extra={
                    "endpoint": "stream",
                    "total_chunks": len(all_chunks),
                    "function_name": func_name,
                },
            )

        except Exception as e:
            logger.exception(
                "Error in stream endpoint",
                extra={"chunks_sent": len(all_chunks)},
            )

            yield f"data: {json.dumps({'error': str(e)})}\n\n"
            yield "data: [DONE]\n\n"

    async def _handle_stream_request(
        self, request: dict[str, Any], return_trace_id: bool
    ) -> StreamingResponse:
        if _stream_function is None:
            raise HTTPException(status_code=500, detail="No stream function registered")
        return StreamingResponse(
            self._generate(_stream_function, request, return_trace_id),
            media_type="text/event-stream",
        )

    @staticmethod
    def _parse_server_args():
        """Parse command line arguments for the agent server"""
        parser = argparse.ArgumentParser(description="Start the agent server")
        parser.add_argument(
            "--port", type=int, default=8000, help="Port to run the server on (default: 8000)"
        )
        parser.add_argument(
            "--workers",
            type=int,
            default=1,
            help="Number of workers to run the server on (default: 1)",
        )
        parser.add_argument(
            "--reload",
            action="store_true",
            help="Reload the server on code changes (default: False)",
        )
        return parser.parse_args()

    def run(
        self,
        app_import_string: str,
        host: str = "0.0.0.0",
    ) -> None:
        """Run the agent server with command line argument parsing."""
        args = self._parse_server_args()
        uvicorn.run(
            app_import_string, host=host, port=args.port, workers=args.workers, reload=args.reload
        )
