#  Copyright (c) ZenML GmbH 2025. 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.
"""Request management utilities."""

import asyncio
import base64
import json
import random
import time
from contextvars import ContextVar
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
from uuid import UUID, uuid4

from fastapi import Request, Response
from fastapi.responses import JSONResponse

from zenml.constants import MEDIUMTEXT_MAX_LENGTH
from zenml.exceptions import EntityExistsError
from zenml.logger import get_logger
from zenml.models import ApiTransactionRequest, ApiTransactionUpdate
from zenml.utils.json_utils import pydantic_encoder
from zenml.utils.time_utils import utc_now

if TYPE_CHECKING:
    from zenml.zen_server.auth import AuthContext


logger = get_logger(__name__)


class RequestContext:
    """Context for a request."""

    def __init__(self, request: Request) -> None:
        """Initialize the request context.

        Args:
            request: The request.

        Raises:
            ValueError: If the idempotency key is not a valid UUID.
        """
        self.request = request
        self.request_id = request.headers.get("X-Request-ID", str(uuid4())[:8])
        self.transaction_id: Optional[UUID] = None
        transaction_id = request.headers.get("Idempotency-Key")
        if transaction_id:
            try:
                self.transaction_id = UUID(transaction_id)
            except ValueError:
                raise ValueError(
                    f"Invalid UUID idempotency key: {transaction_id}. "
                    "Please use a valid UUID."
                )

        # Use a random trace ID to identify the request internally in the logs.
        self.trace_id = str(uuid4())[:4]

        self.source = request.headers.get("User-Agent") or ""
        self.received_at = utc_now()

        self.auth_context: Optional["AuthContext"] = None

    @property
    def process_time(self) -> float:
        """Get the process time in seconds of the request.

        Returns:
            The request ID.
        """
        return (utc_now() - self.received_at).total_seconds()

    @property
    def log_request_id(self) -> str:
        """Get the full request ID for logging.

        Returns:
            The request ID.
        """
        source_type = self.source.split("/")[0]
        return f"{self.request_id}/{source_type}/{self.trace_id}"

    @property
    def log_request(self) -> str:
        """Get the request details for logging.

        Returns:
            The request details for logging.
        """
        client_ip = (
            self.request.client.host if self.request.client else "unknown"
        )
        url_path = self.request.url.path
        method = self.request.method
        return f"{method} {url_path} from {client_ip}"

    @property
    def log_duration(self) -> str:
        """Get the duration of the request.

        Returns:
            The duration of the request.
        """
        current_time = utc_now()
        duration = (current_time - self.received_at).total_seconds() * 1000
        return f"{duration:.2f}ms"

    @property
    def is_cacheable(self) -> bool:
        """Check if the request is cacheable.

        Returns:
            Whether the request is cacheable.
        """
        # Only cache requests that are authenticated and are part of a
        # transaction.
        return (
            self.auth_context is not None and self.transaction_id is not None
        )


class RequestRecord:
    """A record of an in-flight or cached request."""

    future: asyncio.Future[Any]
    request_context: RequestContext

    def __init__(
        self, future: asyncio.Future[Any], request_context: RequestContext
    ) -> None:
        """Initialize the request record.

        Args:
            future: The future of the request.
            request_context: The request context.
        """
        self.future = future
        self.request_context = request_context
        self.completed = False

    def set_result(self, result: Any) -> None:
        """Set the result of the request.

        Args:
            result: The result of the request.
        """
        self.future.set_result(result)
        self.completed = True

    def set_exception(self, exception: Exception) -> None:
        """Set the exception of the request.

        Args:
            exception: The exception of the request.
        """
        self.future.set_exception(exception)
        self.completed = True


class RequestManager:
    """A manager for requests.

    This class is used to manage requests by caching the results of requests
    that have already been executed.
    """

    def __init__(
        self,
    ) -> None:
        """Initialize the request manager."""
        from zenml.zen_server.utils import server_config

        cfg = server_config()
        self.deduplicate = cfg.request_deduplication
        self.transactions: Dict[UUID, RequestRecord] = dict()
        self.lock = asyncio.Lock()
        self.transaction_ttl = cfg.request_cache_timeout
        self.request_timeout = cfg.request_timeout

        self.request_contexts: ContextVar[Optional[RequestContext]] = (
            ContextVar("request_contexts", default=None)
        )

        self._cleanup_interval = cfg.api_transaction_cleanup_interval

        self._cleanup_task: Optional[asyncio.Task[None]] = None
        self._shutdown_event = asyncio.Event()

    @property
    def current_request(self) -> RequestContext:
        """Get the current request context.

        Returns:
            The current request context.

        Raises:
            RuntimeError: If no request context is set.
        """
        request_context = self.request_contexts.get()
        if request_context is None:
            raise RuntimeError("No request context set")
        return request_context

    @current_request.setter
    def current_request(self, request_context: RequestContext) -> None:
        """Set the current request context.

        Args:
            request_context: The request context.
        """
        self.request_contexts.set(request_context)

    async def startup(self) -> None:
        """Start the request manager."""
        self._shutdown_event.clear()
        self._cleanup_task = asyncio.create_task(self._run_cleanup_loop())

    async def shutdown(self) -> None:
        """Shutdown the request manager."""
        self._shutdown_event.set()
        if self._cleanup_task is not None:
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass
            self._cleanup_task = None

    async def _run_cleanup_loop(self) -> None:
        """Run the cleanup loop for expired transactions.

        This loop runs continuously in the background, cleaning up expired
        API transactions. It includes:
        - Random jitter on startup and between iterations to reduce overlap
        between pods
        - Graceful shutdown via the shutdown event
        """
        # Random initial delay (0-interval seconds) to stagger pod startup
        initial_jitter = random.uniform(0, self._cleanup_interval)
        logger.debug(
            f"Transaction cleanup will start in {initial_jitter:.1f}s"
        )
        try:
            await asyncio.wait_for(
                self._shutdown_event.wait(), timeout=initial_jitter
            )
            return
        except asyncio.TimeoutError:
            pass

        while not self._shutdown_event.is_set():
            try:
                await asyncio.get_event_loop().run_in_executor(
                    None, self._cleanup_expired_transactions
                )
            except Exception:
                logger.exception("Error during transaction cleanup")

            # Add random jitter (±25%) to interval to reduce pod overlap
            jitter = random.uniform(0.75, 1.25)
            sleep_time = self._cleanup_interval * jitter

            try:
                await asyncio.wait_for(
                    self._shutdown_event.wait(), timeout=sleep_time
                )
                break
            except asyncio.TimeoutError:
                pass

    def _cleanup_expired_transactions(self) -> None:
        """Execute the cleanup operation in a thread pool."""
        from zenml.zen_server.utils import zen_store

        start_time = time.perf_counter()
        zen_store().cleanup_expired_api_transactions()
        duration = time.perf_counter() - start_time

        logger.debug(f"Cleaning up expired transactions took {duration:.2f}s")

    async def async_run_and_cache_result(
        self,
        func: Callable[..., Any],
        deduplicate: bool,
        request_record: RequestRecord,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        """Run a request and cache the result.

        This method is called in the background to run a request and cache the
        result.

        Args:
            func: The function to execute.
            deduplicate: Whether to enable or disable request deduplication for
                this request.
            request_record: The request record to cache the result in.
            *args: The arguments to pass to the function.
            **kwargs: The keyword arguments to pass to the function.
        """
        from starlette.concurrency import run_in_threadpool

        from zenml.zen_server.utils import get_system_metrics_log_str

        request_context = request_record.request_context
        transaction_id = request_context.transaction_id

        logger.debug(
            f"[{request_context.log_request_id}] ENDPOINT STATS - "
            f"{request_context.log_request} "
            f"async {func.__name__} STARTED "
            f"{get_system_metrics_log_str(request_context.request)}"
        )

        def sync_run_and_cache_result(*args: Any, **kwargs: Any) -> Any:
            from zenml.zen_server.utils import zen_store

            # Copy the deduplicate flag to a local variable to avoid modifying
            # the argument in the outer scope.
            deduplicate_request = deduplicate

            logger.debug(
                f"[{request_context.log_request_id}] ENDPOINT STATS - "
                f"{request_context.log_request} "
                f"sync {func.__name__} STARTED "
                f"{get_system_metrics_log_str(request_context.request)}"
            )

            try:
                # Create or get the API transaction from the database
                if deduplicate_request:
                    assert transaction_id is not None
                    try:
                        api_transaction, transaction_created = (
                            zen_store().get_or_create_api_transaction(
                                api_transaction=ApiTransactionRequest(
                                    transaction_id=transaction_id,
                                    method=request_context.request.method,
                                    url=str(request_context.request.url),
                                )
                            )
                        )
                    except EntityExistsError:
                        logger.error(
                            f"[{request_context.log_request_id}] "
                            f"Transaction {transaction_id} already exists "
                            f"with method {request_context.request.method} and "
                            f"URL {str(request_context.request.url)}. Skipping "
                            "caching."
                        )
                        deduplicate_request = False
                    except Exception:
                        logger.exception(
                            f"[{request_context.log_request_id}] "
                            f"Unexpected error getting or creating API "
                            "transaction. Skipping deduplication."
                        )
                        deduplicate_request = False

                    if deduplicate_request:
                        if api_transaction.completed:
                            logger.debug(
                                f"[{request_context.log_request_id}] "
                                "ENDPOINT STATS - "
                                f"{request_context.log_request} "
                                f"sync {func.__name__} CACHE HIT "
                                f"{get_system_metrics_log_str(request_context.request)}"
                            )

                            try:
                                # The transaction already completed, we can return the
                                # result right away.
                                result = (
                                    zen_store().get_api_transaction_result(
                                        api_transaction_id=transaction_id,
                                    )
                                )
                            except Exception:
                                logger.exception(
                                    f"[{request_context.log_request_id}] "
                                    f"Unexpected error getting API transaction "
                                    "result. Skipping deduplication."
                                )
                                deduplicate_request = False
                            else:
                                if result is not None:
                                    return Response(
                                        base64.b64decode(result),
                                        media_type="application/json",
                                    )
                                else:
                                    return

                        elif not transaction_created:
                            logger.debug(
                                f"[{request_context.log_request_id}] "
                                "ENDPOINT STATS - "
                                f"{request_context.log_request} "
                                f"sync {func.__name__} DELAYED "
                                f"{get_system_metrics_log_str(request_context.request)}"
                            )

                            # The transaction is being processed by another server
                            # instance. We need to wait for it to complete. Instead
                            # of blocking this worker thread, we return a 429 error
                            # to the client to force it to retry later.
                            return JSONResponse(
                                {
                                    "error": "Server too busy. Please try again later."
                                },
                                status_code=429,
                            )

                try:
                    result = func(*args, **kwargs)
                except Exception:
                    if deduplicate_request:
                        assert transaction_id is not None
                        try:
                            # We don't cache exceptions. If the client retries, the
                            # request will be executed again and the exception, if
                            # persistent, will be raised again.
                            zen_store().delete_api_transaction(
                                api_transaction_id=transaction_id,
                            )
                        except Exception:
                            logger.exception(
                                f"[{request_context.log_request_id}] "
                                f"Unexpected error deleting API transaction."
                            )
                    raise

                if deduplicate_request:
                    assert transaction_id is not None
                    cache_result = True
                    result_to_cache: Optional[bytes] = None
                    if result is not None:
                        try:
                            result_to_cache = base64.b64encode(
                                json.dumps(
                                    result, default=pydantic_encoder
                                ).encode("utf-8")
                            )
                        except Exception:
                            # If the result is not serializable, we don't cache it.
                            cache_result = False
                            logger.exception(
                                f"Failed to serialize result of {func.__name__} "
                                f"for transaction {transaction_id}. Skipping "
                                "caching."
                            )
                        else:
                            if len(result_to_cache) > MEDIUMTEXT_MAX_LENGTH:
                                # If the result is too large, we also don't cache it.
                                cache_result = False
                                result_to_cache = None
                                logger.error(
                                    f"Result of {func.__name__} "
                                    f"for transaction {transaction_id} is too "
                                    "large. Skipping caching."
                                )

                    if cache_result:
                        api_transaction_update = ApiTransactionUpdate(
                            cache_time=self.transaction_ttl,
                        )
                        if result_to_cache is not None:
                            api_transaction_update.set_result(
                                result_to_cache.decode("utf-8")
                            )
                        try:
                            zen_store().finalize_api_transaction(
                                api_transaction_id=transaction_id,
                                api_transaction_update=api_transaction_update,
                            )
                        except Exception:
                            logger.exception(
                                f"[{request_context.log_request_id}] "
                                f"Unexpected error finalizing API transaction. "
                                "Ignoring."
                            )
                    else:
                        try:
                            # If the result is not cacheable, there is no point in
                            # keeping the transaction around.
                            zen_store().delete_api_transaction(
                                api_transaction_id=transaction_id,
                            )
                        except Exception:
                            logger.exception(
                                f"[{request_context.log_request_id}] "
                                f"Unexpected error deleting API transaction. "
                                "Ignoring."
                            )

                return result
            finally:
                logger.debug(
                    f"[{request_context.log_request_id}] ENDPOINT STATS - "
                    f"{request_context.log_request} "
                    f"sync {func.__name__} COMPLETED "
                    f"{get_system_metrics_log_str(request_context.request)}"
                )

        try:
            result = await run_in_threadpool(
                sync_run_and_cache_result, *args, **kwargs
            )
        except Exception as e:
            result = e

        if deduplicate:
            async with self.lock:
                if transaction_id in self.transactions:
                    del self.transactions[transaction_id]

        logger.debug(
            f"[{request_context.log_request_id}] ENDPOINT STATS - "
            f"{request_context.log_request} "
            f"async {func.__name__} COMPLETED "
            f"{get_system_metrics_log_str(request_context.request)}"
        )

        if isinstance(result, Exception):
            request_record.set_exception(result)
        else:
            request_record.set_result(result)

    async def execute(
        self,
        func: Callable[..., Any],
        deduplicate: Optional[bool],
        *args: Any,
        **kwargs: Any,
    ) -> Any:
        """Execute a request with in-memory de-duplication.

        Call this method to execute a request with in-memory de-duplication.
        If the request (identified by the request_id) has previously been
        executed and the result is still cached, the result will be returned
        immediately. If the request is in-flight, the method will wait for it
        to complete. If the request is not in-flight, the method will start
        execution in the background and cache the result.

        Args:
            func: The function to execute.
            deduplicate: Whether to enable or disable request deduplication for
                this request. If not specified, by default, the deduplication
                is enabled for POST requests and disabled for other requests.
            *args: The arguments to pass to the function.
            **kwargs: The keyword arguments to pass to the function.

        Returns:
            The result of the request.
        """
        from zenml.zen_server.utils import get_system_metrics_log_str

        request_context = self.current_request
        assert request_context is not None

        logger.debug(
            f"[{request_context.log_request_id}] ENDPOINT STATS - "
            f"{request_context.log_request} "
            f"{func.__name__} STARTED "
            f"{get_system_metrics_log_str(request_context.request)}"
        )

        transaction_id = request_context.transaction_id

        if deduplicate is None:
            # If not specified, by default, the deduplication is enabled for
            # POST requests and disabled for other requests.
            deduplicate = request_context.request.method == "POST"

        deduplicate = (
            deduplicate and self.deduplicate and request_context.is_cacheable
        )

        async with self.lock:
            if deduplicate and transaction_id in self.transactions:
                # The transaction is still being processed on the same
                # server instance. We just wait for it to complete.
                fut = self.transactions[transaction_id].future
                logger.debug(
                    f"[{request_context.log_request_id}] ENDPOINT STATS - "
                    f"{request_context.log_request} "
                    f"{func.__name__} RESUMED "
                    f"{get_system_metrics_log_str(request_context.request)}"
                )
            else:
                # Start execution in background, use the future to wait for it
                # to complete.
                fut = asyncio.get_event_loop().create_future()
                request_record = RequestRecord(
                    future=fut, request_context=request_context
                )
                if deduplicate:
                    assert transaction_id is not None
                    # Also record the transaction for deduplication.
                    self.transactions[transaction_id] = request_record
                asyncio.create_task(
                    self.async_run_and_cache_result(
                        func,
                        deduplicate,
                        request_record,
                        *args,
                        **kwargs,
                    )
                )

        # Wait for the request to complete; timeout if deduplication is enabled
        try:
            # We take into account the time that has already elapsed since the
            # request was received to avoid keeping the request for too long.
            timeout = max(
                0, self.request_timeout - request_context.process_time
            )
            result = await asyncio.wait_for(
                # We use asyncio.shield to prevent the request from being
                # cancelled when the timeout is reached.
                asyncio.shield(fut),
                timeout=timeout if deduplicate else None,
            )

            logger.debug(
                f"[{request_context.log_request_id}] ENDPOINT STATS - "
                f"{request_context.log_request} "
                f"{func.__name__} COMPLETED "
                f"{get_system_metrics_log_str(request_context.request)}"
            )
            return result
        except asyncio.TimeoutError:
            logger.debug(
                f"[{request_context.log_request_id}] ENDPOINT STATS - "
                f"{request_context.log_request} "
                f"{func.__name__} TIMEOUT "
                f"{get_system_metrics_log_str(request_context.request)}"
            )
            return JSONResponse(
                {"error": "Server too busy. Please try again later."},
                status_code=429,
            )
