# Based ons: https://github.com/openai/openai-cookbook/blob/6df6ceff470eeba26a56de131254e775292eac22/examples/api_request_parallel_processor.py
# Several changes were made to make it work with MLflow.
# Currently, only chat completion is supported.

"""
API REQUEST PARALLEL PROCESSOR

Using the LangChain API to process lots of text quickly takes some care.
If you trickle in a million API requests one by one, they'll take days to complete.
This script parallelizes requests using LangChain API.

Features:
- Streams requests from file, to avoid running out of memory for giant jobs
- Makes requests concurrently, to maximize throughput
- Logs errors, to diagnose problems with requests
"""

from __future__ import annotations

import logging
import queue
import threading
import time
import traceback
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from typing import Any

from mlflow.langchain._compat import import_base_callback_handler, try_import_chain

BaseCallbackHandler = import_base_callback_handler()
Chain = try_import_chain()

import mlflow
from mlflow.exceptions import MlflowException
from mlflow.langchain.utils.chat import (
    transform_request_json_for_chat_if_necessary,
    try_transform_response_iter_to_chat_format,
    try_transform_response_to_chat_format,
)
from mlflow.langchain.utils.serialization import convert_to_serializable
from mlflow.pyfunc.context import Context, get_prediction_context
from mlflow.tracing.utils import maybe_set_prediction_context

_logger = logging.getLogger(__name__)


@dataclass
class StatusTracker:
    """
    Stores metadata about the script's progress. Only one instance is created.
    """

    num_tasks_started: int = 0
    num_tasks_in_progress: int = 0  # script ends when this reaches 0
    num_tasks_succeeded: int = 0
    num_tasks_failed: int = 0
    num_api_errors: int = 0  # excluding rate limit errors, counted above
    lock: threading.Lock = threading.Lock()

    def start_task(self):
        with self.lock:
            self.num_tasks_started += 1
            self.num_tasks_in_progress += 1

    def complete_task(self, *, success: bool):
        with self.lock:
            self.num_tasks_in_progress -= 1
            if success:
                self.num_tasks_succeeded += 1
            else:
                self.num_tasks_failed += 1

    def increment_num_api_errors(self):
        with self.lock:
            self.num_api_errors += 1


@dataclass
class APIRequest:
    """
    Stores an API request's inputs, outputs, and other metadata. Contains a method to make an API
    call.

    Args:
        index: The request's index in the tasks list
        lc_model: The LangChain model to call
        request_json: The request's input data
        results: The list to append the request's output data to, it's a list of tuples
            (index, response)
        errors: A dictionary to store any errors that occur
        convert_chat_responses: Whether to convert the model's responses to chat format
        did_perform_chat_conversion: Whether the input data was converted to chat format
            based on the model's type and input data.
        stream: Whether the request is a stream request
        prediction_context: The prediction context to use for the request
    """

    index: int
    lc_model: Any
    request_json: dict[str, Any]
    results: list[tuple[int, str]]
    errors: dict[int, str]
    convert_chat_responses: bool
    did_perform_chat_conversion: bool
    stream: bool
    params: dict[str, Any]
    prediction_context: Context | None = None

    def _predict_single_input(self, single_input, callback_handlers, **kwargs):
        config = kwargs.pop("config", {})
        config["callbacks"] = config.get("callbacks", []) + (callback_handlers or [])
        if self.stream:
            return self.lc_model.stream(single_input, config=config, **kwargs)
        if hasattr(self.lc_model, "invoke"):
            return self.lc_model.invoke(single_input, config=config, **kwargs)
        else:
            # for backwards compatibility, __call__ is deprecated and will be removed in 0.3.0
            # kwargs shouldn't have config field if invoking with __call__
            return self.lc_model(single_input, callbacks=callback_handlers, **kwargs)

    def _try_convert_response(self, response):
        if self.stream:
            return try_transform_response_iter_to_chat_format(response)
        else:
            return try_transform_response_to_chat_format(response)

    def single_call_api(self, callback_handlers: list[BaseCallbackHandler] | None):
        from mlflow.langchain._compat import import_base_retriever
        from mlflow.langchain.utils.logging import langgraph_types, lc_runnables_types

        BaseRetriever = import_base_retriever()

        if isinstance(self.lc_model, BaseRetriever):
            # Retrievers are invoked differently than Chains
            response = self.lc_model.get_relevant_documents(
                **self.request_json, callbacks=callback_handlers, **self.params
            )
        elif isinstance(self.lc_model, lc_runnables_types() + langgraph_types()):
            if isinstance(self.request_json, dict):
                # This is a temporary fix for the case when spark_udf converts
                # input into pandas dataframe with column name, while the model
                # does not accept dictionaries as input, it leads to errors like
                # Expected Scalar value for String field 'query_text'
                try:
                    response = self._predict_single_input(
                        self.request_json, callback_handlers, **self.params
                    )
                except TypeError as e:
                    _logger.debug(
                        f"Failed to invoke {self.lc_model.__class__.__name__} "
                        f"with {self.request_json}. Error: {e!r}. Trying to "
                        "invoke with the first value of the dictionary."
                    )
                    self.request_json = next(iter(self.request_json.values()))
                    (
                        prepared_request_json,
                        did_perform_chat_conversion,
                    ) = transform_request_json_for_chat_if_necessary(
                        self.request_json, self.lc_model
                    )
                    self.did_perform_chat_conversion = did_perform_chat_conversion

                    response = self._predict_single_input(
                        prepared_request_json, callback_handlers, **self.params
                    )
            else:
                response = self._predict_single_input(
                    self.request_json, callback_handlers, **self.params
                )

            if self.did_perform_chat_conversion or self.convert_chat_responses:
                response = self._try_convert_response(response)
        else:
            # return_only_outputs is invalid for stream call
            if Chain and isinstance(self.lc_model, Chain) and not self.stream:
                kwargs = {"return_only_outputs": True}
            else:
                kwargs = {}
            kwargs.update(**self.params)
            response = self._predict_single_input(self.request_json, callback_handlers, **kwargs)

            if self.did_perform_chat_conversion or self.convert_chat_responses:
                response = self._try_convert_response(response)
            elif isinstance(response, dict) and len(response) == 1:
                # to maintain existing code, single output chains will still return
                # only the result
                response = response.popitem()[1]

        return convert_to_serializable(response)

    def call_api(
        self, status_tracker: StatusTracker, callback_handlers: list[BaseCallbackHandler] | None
    ):
        """
        Calls the LangChain API and stores results.
        """
        _logger.debug(f"Request #{self.index} started with payload: {self.request_json}")

        try:
            with maybe_set_prediction_context(self.prediction_context):
                response = self.single_call_api(callback_handlers)
            _logger.debug(f"Request #{self.index} succeeded with response: {response}")
            self.results.append((self.index, response))
            status_tracker.complete_task(success=True)
        except Exception as e:
            self.errors[self.index] = (
                f"error: {e!r} {traceback.format_exc()}\n request payload: {self.request_json}"
            )
            status_tracker.increment_num_api_errors()
            status_tracker.complete_task(success=False)


def process_api_requests(
    lc_model,
    requests: list[Any | dict[str, Any]] | None = None,
    max_workers: int = 10,
    callback_handlers: list[BaseCallbackHandler] | None = None,
    convert_chat_responses: bool = False,
    params: dict[str, Any] | None = None,
    context: Context | None = None,
):
    """
    Processes API requests in parallel.
    """

    # initialize trackers
    retry_queue = queue.Queue()
    status_tracker = StatusTracker()  # single instance to track a collection of variables
    next_request = None  # variable to hold the next request to call
    context = context or get_prediction_context()

    results = []
    errors = {}

    # Note: we should call `transform_request_json_for_chat_if_necessary`
    # for the whole batch data, because the conversion should obey the rule
    # that if any record in the batch can't be converted, then all the record
    # in this batch can't be converted.
    (
        converted_chat_requests,
        did_perform_chat_conversion,
    ) = transform_request_json_for_chat_if_necessary(requests, lc_model)

    requests_iter = enumerate(converted_chat_requests)
    with ThreadPoolExecutor(
        max_workers=max_workers, thread_name_prefix="MlflowLangChainApi"
    ) as executor:
        while True:
            # get next request (if one is not already waiting for capacity)
            if not retry_queue.empty():
                next_request = retry_queue.get_nowait()
                _logger.warning(f"Retrying request {next_request.index}: {next_request}")
            elif req := next(requests_iter, None):
                # get new request
                index, converted_chat_request_json = req
                next_request = APIRequest(
                    index=index,
                    lc_model=lc_model,
                    request_json=converted_chat_request_json,
                    results=results,
                    errors=errors,
                    convert_chat_responses=convert_chat_responses,
                    did_perform_chat_conversion=did_perform_chat_conversion,
                    stream=False,
                    prediction_context=context,
                    params=params,
                )
                status_tracker.start_task()
            else:
                next_request = None

            # if enough capacity available, call API
            if next_request:
                # call API
                executor.submit(
                    next_request.call_api,
                    status_tracker=status_tracker,
                    callback_handlers=callback_handlers,
                )

            # if all tasks are finished, break
            # check next_request to avoid terminating the process
            # before extra requests need to be processed
            if status_tracker.num_tasks_in_progress == 0 and next_request is None:
                break

            time.sleep(0.001)  # avoid busy waiting

        # after finishing, log final status
        if status_tracker.num_tasks_failed > 0:
            raise mlflow.MlflowException(
                f"{status_tracker.num_tasks_failed} tasks failed. Errors: {errors}"
            )

        return [res for _, res in sorted(results)]


def process_stream_request(
    lc_model,
    request_json: Any | dict[str, Any],
    callback_handlers: list[BaseCallbackHandler] | None = None,
    convert_chat_responses: bool = False,
    params: dict[str, Any] | None = None,
):
    """
    Process single stream request.
    """
    if not hasattr(lc_model, "stream"):
        raise MlflowException(
            f"Model {lc_model.__class__.__name__} does not support streaming prediction output. "
            "No `stream` method found."
        )

    (
        converted_chat_requests,
        did_perform_chat_conversion,
    ) = transform_request_json_for_chat_if_necessary(request_json, lc_model)

    api_request = APIRequest(
        index=0,
        lc_model=lc_model,
        request_json=converted_chat_requests,
        results=None,
        errors=None,
        convert_chat_responses=convert_chat_responses,
        did_perform_chat_conversion=did_perform_chat_conversion,
        stream=True,
        prediction_context=get_prediction_context(),
        params=params,
    )
    with maybe_set_prediction_context(api_request.prediction_context):
        return api_request.single_call_api(callback_handlers)