#  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at:
#
#       https://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
#  or implied. See the License for the specific language governing
#  permissions and limitations under the License.
"""Server middlewares."""

import logging
import os
from asyncio import Lock
from asyncio.log import logger
from datetime import datetime, timedelta
from typing import Any, Set

from anyio import CapacityLimiter, to_thread
from fastapi import FastAPI, Request
from fastapi.responses import PlainTextResponse
from starlette.middleware.base import (
    BaseHTTPMiddleware,
    RequestResponseEndpoint,
)
from starlette.middleware.cors import CORSMiddleware
from starlette.responses import (
    JSONResponse,
    Response,
)
from starlette.types import ASGIApp

from zenml.analytics import source_context
from zenml.constants import (
    DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS,
    ENV_ZENML_DEFAULT_ANALYTICS_SOURCE,
    HEALTH,
    READY,
)
from zenml.enums import SourceContextTypes
from zenml.utils.time_utils import utc_now
from zenml.zen_server.request_management import RequestContext
from zenml.zen_server.secure_headers import (
    secure_headers,
)
from zenml.zen_server.utils import (
    get_system_metrics_log_str,
    is_user_request,
    request_manager,
    server_config,
    zen_store,
)

# Track active requests with an atomic counter
active_requests_count = 0
active_requests_lock = Lock()


# Initialize last_user_activity
last_user_activity: datetime = utc_now()
last_user_activity_reported: datetime = last_user_activity + timedelta(
    seconds=-DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS
)
last_user_activity_lock = Lock()
# Create a custom thread pool limiter with a limit of 1 thread for all
# user activity updates
last_user_activity_thread_limiter = CapacityLimiter(1)


class RequestBodyLimit(BaseHTTPMiddleware):
    """Limits the size of the request body."""

    def __init__(self, app: ASGIApp, max_bytes: int) -> None:
        """Limits the size of the request body.

        Args:
            app: The FastAPI app.
            max_bytes: The maximum size of the request body.
        """
        super().__init__(app)
        self.max_bytes = max_bytes

    async def dispatch(
        self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        """Limits the size of the request body.

        Args:
            request: The incoming request.
            call_next: The next function to be called.

        Returns:
            The response to the request.
        """
        if content_length := request.headers.get("content-length"):
            if int(content_length) > self.max_bytes:
                return Response(status_code=413)  # Request Entity Too Large

        try:
            return await call_next(request)
        except Exception:
            logger.exception("An error occurred while processing the request")
            return JSONResponse(
                status_code=500,
                content={"detail": "An unexpected error occurred."},
            )


class RestrictFileUploadsMiddleware(BaseHTTPMiddleware):
    """Restrict file uploads to certain paths."""

    def __init__(self, app: ASGIApp, allowed_paths: Set[str]):
        """Restrict file uploads to certain paths.

        Args:
            app: The FastAPI app.
            allowed_paths: The allowed paths.
        """
        super().__init__(app)
        self.allowed_paths = allowed_paths

    async def dispatch(
        self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        """Restrict file uploads to certain paths.

        Args:
            request: The incoming request.
            call_next: The next function to be called.

        Returns:
            The response to the request.
        """
        if request.method == "POST":
            content_type = request.headers.get("content-type", "")
            if (
                "multipart/form-data" in content_type
                and request.url.path not in self.allowed_paths
            ):
                return JSONResponse(
                    status_code=403,
                    content={
                        "detail": "File uploads are not allowed on this endpoint."
                    },
                )

        try:
            return await call_next(request)
        except Exception:
            logger.exception("An error occurred while processing the request")
            return JSONResponse(
                status_code=500,
                content={"detail": "An unexpected error occurred."},
            )


ALLOWED_FOR_FILE_UPLOAD: Set[str] = set()


async def track_last_user_activity(request: Request, call_next: Any) -> Any:
    """A middleware to track last user activity.

    This middleware checks if the incoming request is a user request and
    updates the last activity timestamp if it is.

    Args:
        request: The incoming request object.
        call_next: A function that will receive the request as a parameter and
            pass it to the corresponding path operation.

    Returns:
        The response to the request.
    """
    global last_user_activity
    global last_user_activity_reported
    global last_user_activity_lock

    now = utc_now()

    try:
        if is_user_request(request):
            report_user_activity = False
            async with last_user_activity_lock:
                last_user_activity = now
                if (
                    (now - last_user_activity_reported).total_seconds()
                    > DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS
                ):
                    last_user_activity_reported = now
                    report_user_activity = True

            if report_user_activity:
                # We don't want to make a DB call here because we're in the
                # context of the asyncio event loop and it would block the
                # entire application for who knows how long.
                # We use the threadpool for it.

                request_context = request_manager().current_request

                def update_last_user_activity_timestamp() -> None:
                    logger.debug(
                        f"[{request_context.log_request_id}] API STATS - "
                        f"{request_context.log_request} "
                        f"UPDATING LAST USER ACTIVITY "
                        f"{get_system_metrics_log_str(request_context.request)}"
                    )

                    try:
                        zen_store()._update_last_user_activity_timestamp(
                            last_user_activity=last_user_activity,
                        )
                    finally:
                        logger.debug(
                            f"[{request_context.log_request_id}] API STATS - "
                            f"{request_context.log_request} "
                            f"UPDATED LAST USER ACTIVITY "
                            f"{get_system_metrics_log_str(request_context.request)}"
                        )

                await to_thread.run_sync(
                    update_last_user_activity_timestamp,
                    limiter=last_user_activity_thread_limiter,
                )

    except Exception as e:
        logger.debug(
            f"An unexpected error occurred while checking user activity: {e}"
        )

    try:
        return await call_next(request)
    except Exception:
        logger.exception("An error occurred while processing the request")
        return JSONResponse(
            status_code=500,
            content={"detail": "An unexpected error occurred."},
        )


async def infer_source_context(request: Request, call_next: Any) -> Any:
    """A middleware to track the source of an event.

    It extracts the source context from the header of incoming requests
    and applies it to the ZenML source context on the API side. This way, the
    outgoing analytics request can append it as an additional field.

    Args:
        request: the incoming request object.
        call_next: a function that will receive the request as a parameter and
            pass it to the corresponding path operation.

    Returns:
        the response to the request.
    """
    try:
        default_source = os.environ.get(
            ENV_ZENML_DEFAULT_ANALYTICS_SOURCE,
            SourceContextTypes.API.value,
        )
        s = request.headers.get(
            source_context.name,
            default=default_source,
        )
        source_context.set(SourceContextTypes(s))
    except Exception as e:
        logger.warning(
            f"An unexpected error occurred while getting the source "
            f"context: {e}"
        )
        source_context.set(SourceContextTypes.API)

    try:
        return await call_next(request)
    except Exception:
        logger.exception("An error occurred while processing the request")
        return JSONResponse(
            status_code=500,
            content={"detail": "An unexpected error occurred."},
        )


async def set_secure_headers(request: Request, call_next: Any) -> Any:
    """Middleware to set secure headers.

    Args:
        request: The incoming request.
        call_next: The next function to be called.

    Returns:
        The response with secure headers set.
    """
    try:
        response = await call_next(request)
    except Exception:
        logger.exception("An error occurred while processing the request")
        response = JSONResponse(
            status_code=500,
            content={"detail": "An unexpected error occurred."},
        )

    # If the request is for the openAPI docs, don't set secure headers
    if request.url.path.startswith("/docs") or request.url.path.startswith(
        "/redoc"
    ):
        return response

    await secure_headers().set_headers_async(response)
    return response


async def log_requests(request: Request, call_next: Any) -> Any:
    """Log requests to the ZenML server.

    Args:
        request: The incoming request object.
        call_next: A function that will receive the request as a parameter and
            pass it to the corresponding path operation.

    Returns:
        The response to the request.
    """
    global active_requests_count

    if not logger.isEnabledFor(logging.DEBUG):
        return await call_next(request)

    async with active_requests_lock:
        active_requests_count += 1

    request_context = request_manager().current_request

    logger.debug(
        f"[{request_context.log_request_id}] API STATS - "
        f"{request_context.log_request} "
        f"RECEIVED {get_system_metrics_log_str(request)}"
    )

    try:
        response = await call_next(request)

        logger.debug(
            f"[{request_context.log_request_id}] API STATS - "
            f"{response.status_code} {request_context.log_request} "
            f"took {request_context.log_duration} "
            f"{get_system_metrics_log_str(request)}"
        )

        return response
    finally:
        async with active_requests_lock:
            active_requests_count -= 1


async def record_requests(request: Request, call_next: Any) -> Any:
    """Record requests to the ZenML server.

    Args:
        request: The incoming request object.
        call_next: A function that will receive the request as a parameter and
            pass it to the corresponding path operation.

    Returns:
        The response to the request.
    """
    # Keep track of the request context in a context variable
    request_context = RequestContext(request=request)
    request_manager().current_request = request_context

    try:
        response = await call_next(request)
    except Exception:
        logger.exception("An error occurred while processing the request")
        response = JSONResponse(
            status_code=500,
            content={"detail": "An unexpected error occurred."},
        )

    return response


async def skip_health_middleware(request: Request, call_next: Any) -> Any:
    """Skip health and ready endpoints.

    Args:
        request: The incoming request.
        call_next: The next function to be called.

    Returns:
        The response to the request.
    """
    if request.url.path in [HEALTH, READY]:
        # Skip expensive processing
        return PlainTextResponse("ok")

    return await call_next(request)


def add_middlewares(app: FastAPI) -> None:
    """Add middlewares to the FastAPI app.

    Args:
        app: The FastAPI app.
    """
    app.add_middleware(BaseHTTPMiddleware, dispatch=track_last_user_activity)
    app.add_middleware(BaseHTTPMiddleware, dispatch=infer_source_context)

    app.add_middleware(
        CORSMiddleware,
        allow_origins=server_config().cors_allow_origins or [],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    app.add_middleware(
        RequestBodyLimit,
        max_bytes=server_config().max_request_body_size_in_bytes,
    )
    app.add_middleware(
        RestrictFileUploadsMiddleware, allowed_paths=ALLOWED_FOR_FILE_UPLOAD
    )

    app.add_middleware(BaseHTTPMiddleware, dispatch=set_secure_headers)
    app.add_middleware(BaseHTTPMiddleware, dispatch=log_requests)
    app.add_middleware(BaseHTTPMiddleware, dispatch=record_requests)
    app.add_middleware(BaseHTTPMiddleware, dispatch=skip_health_middleware)
