"""
EvalScope integration utilities for ms-swift models.

This module provides a custom ModelAPI implementation that enables batch inference
for evaluation tasks using ms-swift's TransformersEngine. It implements an asynchronous
batch processing system to improve throughput when evaluating models.
"""

from concurrent.futures import Future
from dataclasses import dataclass
from evalscope.api.messages import ChatMessage as EvalChatMessage
from evalscope.api.model import GenerateConfig, ModelAPI, ModelOutput, ModelUsage
from evalscope.api.registry import register_model_api
from evalscope.api.tool import ToolChoice, ToolInfo
from evalscope.models.utils.openai import chat_choices_from_openai
from queue import Empty, Queue
from threading import Thread
from typing import Any, List, Optional, Tuple

from swift.infer_engine import InferRequest, RequestConfig, TransformersEngine


@dataclass
class BatchInferInput:
    """
    Container for batch inference input data.

    Holds all necessary information for a single inference request
    that will be processed as part of a batch.
    """
    ms_input: InferRequest  # ms-swift format request
    ms_config: RequestConfig  # ms-swift format configuration
    batch_size: int  # desired batch size for this request
    engine: TransformersEngine  # inference engine to use


@dataclass
class _QueueItem:
    """
    Internal queue item for batch processing.

    Pairs a batch input with its corresponding future for result delivery.
    """
    input: BatchInferInput
    future: Future[ModelOutput]  # will be resolved with the inference result


# Global variables for batch processing
# These maintain the shared batch processing infrastructure across all model instances
batch_thread: Optional[Thread] = None  # background thread for processing batches
batch_queue: Queue[_QueueItem] = Queue()  # queue of pending inference requests


@register_model_api('swift_custom')
class EvalModel(ModelAPI):
    """
    Custom ModelAPI implementation for ms-swift models with batch inference support.

    This class integrates ms-swift's TransformersEngine with EvalScope's evaluation framework,
    providing efficient batch processing for improved evaluation throughput.
    """

    def __init__(
            self,
            model_name: str,
            base_url: Optional[str] = None,
            api_key: Optional[str] = None,
            config: GenerateConfig = GenerateConfig(),
            **model_args: Any,
    ):
        """
        Initialize the EvalModel with ms-swift backend.

        Args:
            model_name: Name of the model for identification
            base_url: Not used in this implementation (for API compatibility)
            api_key: Not used in this implementation (for API compatibility)
            config: Generation configuration with batch settings
            **model_args: Additional arguments including 'model' and 'template'
        """
        super().__init__(
            model_name=model_name,
            base_url=base_url,
            api_key=api_key,
            config=config,
        )

        # Extract model-specific arguments from kwargs
        # This pattern allows us to collect known arguments while preserving unknown ones
        def collect_model_arg(name: str) -> Optional[Any]:
            value = model_args.get(name, None)
            if value is not None:
                model_args.pop(name)
            return value

        # Extract required model parameters
        self.model = collect_model_arg('model')  # model path or identifier
        self.template = collect_model_arg('template')  # conversation template
        self.max_batch_size = collect_model_arg('max_batch_size')  # maximum batch size

        # Initialize the inference engine with batch support
        self.engine = TransformersEngine(self.model, template=self.template, max_batch_size=self.max_batch_size)

    def generate(
        self,
        input: List[EvalChatMessage],
        tools: List[ToolInfo],
        tool_choice: ToolChoice,
        config: GenerateConfig,
    ) -> ModelOutput:
        """
        Generate model response using batch inference.

        This method queues the request for batch processing and waits for the result.
        The actual inference is performed asynchronously in a background thread.

        Args:
            input: List of chat messages forming the conversation
            tools: Available tools for function calling (if supported)
            tool_choice: Tool selection strategy
            config: Generation configuration

        Returns:
            ModelOutput containing the generated response
        """
        # Ensure the background batch processing thread is running
        global batch_thread
        if batch_thread is None:
            batch_thread = Thread(target=_process_batches, daemon=True)
            batch_thread.start()

        # Convert EvalScope format to ms-swift format
        ms_input = convert_request(input, tools)
        ms_config = convert_config(config)

        # Package the request for batch processing
        batch_input = BatchInferInput(
            ms_input=ms_input, ms_config=ms_config, batch_size=config.batch_size, engine=self.engine)

        # Create a future to receive the result asynchronously
        future = Future[ModelOutput]()

        # Queue the request for batch processing
        batch_queue.put(_QueueItem(input=batch_input, future=future))

        # Block until the result is available
        return future.result()


def _process_batches() -> None:
    """
    Background thread function that processes batched inference requests.

    This function runs continuously, collecting requests from the queue and
    processing them in batches for improved efficiency. It uses a timeout-based
    approach to balance between batch size and latency.
    """
    while True:
        # Collect requests from the queue until timeout or batch size limit
        inputs: List[Tuple[BatchInferInput, Future[ModelOutput]]] = []

        while True:
            try:
                # Wait for new requests with a 2-second timeout
                item = batch_queue.get(timeout=2)
                inputs.append((item.input, item.future))

                # Check if we've reached the desired batch size
                if len(inputs) == item.input.batch_size:
                    break  # Process this batch now

            except Empty:
                # No more requests in queue, process what we have
                break

        # Skip processing if no requests were collected
        if len(inputs) == 0:
            continue

        try:
            # Prepare batch inputs for ms-swift inference
            ms_inputs = [item[0].ms_input for item in inputs]
            ms_config = inputs[0][0].ms_config  # use first config for the batch
            engine = inputs[0][0].engine  # use first engine for the batch

            # Perform batch inference using ms-swift engine
            completions = engine.infer(ms_inputs, ms_config, use_tqdm=False)

            # Process results and deliver them to waiting futures
            for i, (batch_input, future) in enumerate(inputs):
                completion = completions[i]

                # Convert ms-swift response to EvalScope format
                choices = chat_choices_from_openai(completion, tools=[])
                result = ModelOutput(
                    model=completion.model,
                    choices=choices,
                    usage=(ModelUsage(
                        input_tokens=completion.usage.prompt_tokens,
                        output_tokens=completion.usage.completion_tokens,
                        total_tokens=completion.usage.total_tokens,
                    ) if completion.usage else None),
                )

                # Deliver the result to the waiting caller
                future.set_result(result)

        except Exception as ex:
            # If batch processing fails, propagate the error to all waiting futures
            for _, future in inputs:
                future.set_exception(ex)


def convert_config(config: GenerateConfig) -> RequestConfig:
    """
    Convert EvalScope GenerateConfig to ms-swift RequestConfig.

    Maps configuration parameters between the two frameworks, ensuring
    compatibility while maintaining the same generation behavior.

    Args:
        config: EvalScope generation configuration

    Returns:
        RequestConfig: ms-swift compatible configuration
    """
    return RequestConfig(
        max_tokens=config.max_tokens,
        temperature=config.temperature,
        top_k=config.top_k,
        top_p=config.top_p,
        presence_penalty=config.presence_penalty,
        frequency_penalty=config.frequency_penalty,
        seed=config.seed,
        stream=False,  # batch processing doesn't support streaming
        logprobs=config.logprobs,
        top_logprobs=config.top_logprobs)


def convert_request(messages: List[EvalChatMessage], tools: List[ToolInfo]) -> InferRequest:
    """
    Convert EvalScope request format to ms-swift InferRequest format.

    Transforms the message and tool format from EvalScope's representation
    to the format expected by ms-swift's inference engine.

    Args:
        messages: List of chat messages in EvalScope format
        tools: List of available tools in EvalScope format

    Returns:
        InferRequest: ms-swift compatible request object
    """
    # Convert tools to ms-swift format
    tools_list = []
    if len(tools) > 0:
        tools_list = [tool.model_dump(exclude_none=True) for tool in tools]

    # Convert messages to ms-swift format
    ms_messages = []
    for message in messages:
        ms_messages.append(message.model_dump(exclude_none=True))

    return InferRequest(
        messages=ms_messages,
        tools=tools_list,
    )
