# Copyright (c) ModelScope Contributors. All rights reserved.
import datetime as dt
import os
import torch.distributed as dist
from dataclasses import dataclass
from typing import Literal, Optional

from swift.rlhf_trainers import VllmArguments
from swift.utils import get_logger, init_process_group, is_dist, to_abspath
from .base_args import BaseArguments
from .merge_args import MergeArguments

logger = get_logger()


@dataclass
class LmdeployArguments:
    """Holds the configuration arguments for lmdeploy.

    Args:
        lmdeploy_tp (int): The tensor parallelism size. Defaults to 1.
        lmdeploy_session_len (Optional[int]): The maximum session length. Defaults to None.
        lmdeploy_cache_max_entry_count (float): The percentage of GPU memory to be used by the K/V cache. Defaults
            to 0.8.
        lmdeploy_quant_policy (int): The quantization policy for the K/V cache. Set to 4 or 8 for 4-bit or 8-bit
            quantization respectively. Defaults to 0, which means no quantization.
        lmdeploy_vision_batch_size (int): The `max_batch_size` parameter to be passed to `VisionConfig`. Defaults to 1.
    """

    # lmdeploy
    lmdeploy_tp: int = 1
    lmdeploy_session_len: Optional[int] = None
    lmdeploy_cache_max_entry_count: float = 0.8
    lmdeploy_quant_policy: int = 0  # e.g. 4, 8
    lmdeploy_vision_batch_size: int = 1  # max_batch_size in VisionConfig

    def get_lmdeploy_engine_kwargs(self):
        kwargs = {
            'tp': self.lmdeploy_tp,
            'session_len': self.lmdeploy_session_len,
            'cache_max_entry_count': self.lmdeploy_cache_max_entry_count,
            'quant_policy': self.lmdeploy_quant_policy,
            'vision_batch_size': self.lmdeploy_vision_batch_size
        }
        if dist.is_initialized():
            kwargs.update({'devices': [dist.get_rank()]})
        return kwargs


@dataclass
class SglangArguments:
    """Arguments for configuring the SGLang backend.

    Args:
        sglang_tp_size (int): The number of tensor parallel workers. Defaults to 1.
        sglang_pp_size (int): The number of pipeline parallel workers. Defaults to 1.
        sglang_dp_size (int): The number of data parallel workers. Defaults to 1.
        sglang_ep_size (int): The number of expert parallel workers. Defaults to 1.
        sglang_enable_ep_moe (bool): Whether to enable expert parallelism for MoE.
            Note: This argument has been removed in recent versions of SGLang. Defaults to False.
        sglang_mem_fraction_static (Optional[float]): The fraction of GPU memory for the static allocation of model
            weights and the KV cache memory pool. Try lowering this value if you encounter GPU out-of-memory errors.
            Defaults to None.
        sglang_context_length (Optional[int]): The maximum context length for the model. If None, the value from the
            model's `config.json` will be used. Defaults to None.
        sglang_disable_cuda_graph (bool): Disable CUDA graph for inference. Defaults to False.
        sglang_quantization (Optional[str]): The quantization method to use. Defaults to None.
        sglang_kv_cache_dtype (str): The data type for K/V cache storage. 'auto' will use the model's data type.
            'fp8_e5m2' and 'fp8_e4m3' are available for CUDA 11.8 and later. Defaults to 'auto'.
        sglang_enable_dp_attention (bool): Enables data parallelism for the attention mechanism and tensor parallelism
            for the feed-forward network (FFN). The data parallel size (dp_size) must equal the tensor parallel size
            (tp_size). Currently supported for DeepSeek-V2/3 and Qwen2/3 MoE models. Defaults to False.
        sglang_disable_custom_all_reduce (bool): Disable the custom all-reduce kernel and fall back to NCCL. Enabled by
            default (True) for stability. Defaults to True.
        sglang_speculative_algorithm (Optional[str]): The speculative decoding algorithm. Options include "EAGLE",
            "EAGLE3", "NEXTN", "STANDALONE", "NGRAM". Defaults to None.
        sglang_speculative_num_steps (Optional[int]): The number of steps to sample from the draft model during
            speculative decoding. Defaults to None.
        sglang_speculative_eagle_topk (Optional[int]): The number of tokens to sample from the draft model at each step
            for the EAGLE2 algorithm. Defaults to None.
        sglang_speculative_num_draft_tokens (Optional[int]): The number of tokens to sample from the draft model during
            speculative decoding. Defaults to None.
    """
    sglang_tp_size: int = 1
    sglang_pp_size: int = 1
    sglang_dp_size: int = 1
    sglang_ep_size: int = 1
    sglang_enable_ep_moe: bool = False
    sglang_mem_fraction_static: Optional[float] = None
    sglang_context_length: Optional[int] = None
    sglang_disable_cuda_graph: bool = False
    sglang_quantization: Optional[str] = None
    sglang_kv_cache_dtype: str = 'auto'
    sglang_enable_dp_attention: bool = False
    sglang_disable_custom_all_reduce: bool = True
    # speculative decoding
    # e.g. EAGLE, EAGLE3, NEXTN
    sglang_speculative_algorithm: Optional[str] = None
    sglang_speculative_num_steps: Optional[int] = None
    sglang_speculative_eagle_topk: Optional[int] = None
    sglang_speculative_num_draft_tokens: Optional[int] = None

    def get_sglang_engine_kwargs(self):
        kwargs = {
            'tp_size': self.sglang_tp_size,
            'pp_size': self.sglang_pp_size,
            'dp_size': self.sglang_dp_size,
            'ep_size': self.sglang_ep_size,
            'enable_ep_moe': self.sglang_enable_ep_moe,
            'mem_fraction_static': self.sglang_mem_fraction_static,
            'context_length': self.sglang_context_length,
            'disable_cuda_graph': self.sglang_disable_cuda_graph,
            'quantization': self.sglang_quantization,
            'kv_cache_dtype': self.sglang_kv_cache_dtype,
            'enable_dp_attention': self.sglang_enable_dp_attention,
            'disable_custom_all_reduce': self.sglang_disable_custom_all_reduce,
            'speculative_algorithm': self.sglang_speculative_algorithm,
            'speculative_num_steps': self.sglang_speculative_num_steps,
            'speculative_eagle_topk': self.sglang_speculative_eagle_topk,
            'speculative_num_draft_tokens': self.sglang_speculative_num_draft_tokens,
        }
        if self.task_type == 'embedding':
            kwargs['task_type'] = 'embedding'
        return kwargs


@dataclass
class InferArguments(MergeArguments, LmdeployArguments, SglangArguments, VllmArguments, BaseArguments):
    """Arguments for model inference.

    A dataclass that extends BaseArguments, MergeArguments, VllmArguments, and LmdeployArguments to define all
    arguments required for model inference.

    Args:
        infer_backend (Literal['transformers', 'vllm', 'sglang', 'lmdeploy']): The inference acceleration
            backend to use. Defaults to 'transformers'.
        result_path (Optional[str]): The path to store inference results in JSONL format. If the file already exists,
            new results will be appended. If None, results are saved in the checkpoint directory (if available) or
            './result'. The final path will be printed to the console. Defaults to None.
        write_batch_size (int): The batch size for writing results to `result_path`. A value of -1 means no limit.
            Defaults to 1000.
        metric (Optional[str]): The metric to use for evaluating inference results. Supported values are 'acc' and
            'rouge'. If None, no evaluation is performed. Defaults to None.
        max_batch_size (int): The maximum batch size for inference, effective only when `infer_backend` is
            'transformers'. A value of -1 means no limit. Defaults to 1.
        val_dataset_sample (Optional[int]): The number of samples to use from the inference dataset. If None, the
            entire dataset is used. Defaults to None.
        reranker_use_activation (bool): Whether to apply a sigmoid activation to the scores during reranker inference.
            Defaults to True.
    """
    # `pt` is used for swift3.x shell script compatibility.
    infer_backend: Literal['vllm', 'transformers', 'sglang', 'lmdeploy', 'pt'] = 'transformers'

    result_path: Optional[str] = None
    write_batch_size: int = 1000
    metric: Literal['acc', 'rouge'] = None
    # for transformers engine
    max_batch_size: int = 1

    # only for inference
    val_dataset_sample: Optional[int] = None

    # for reranker
    reranker_use_activation: bool = True

    def _get_result_path(self, folder_name: str) -> str:
        result_dir = self.ckpt_dir or f'result/{self.model_suffix}'
        os.makedirs(result_dir, exist_ok=True)
        result_dir = to_abspath(os.path.join(result_dir, folder_name))
        os.makedirs(result_dir, exist_ok=True)
        time = dt.datetime.now().strftime('%Y%m%d-%H%M%S')
        return os.path.join(result_dir, f'{time}.jsonl')

    def _init_result_path(self, folder_name: str) -> None:
        if self.result_path is not None:
            self.result_path = to_abspath(self.result_path)
            return
        # By default, a result_path file is automatically created
        # when a validation or evaluation dataset is present.
        if self._val_dataset_exists or getattr(self, 'eval_dataset', None):
            self.result_path = self._get_result_path(folder_name)
            logger.info(f'args.result_path: {self.result_path}')

    def _init_stream(self):
        self.eval_human = not self._val_dataset_exists
        logger.info(f'Setting args.eval_human: {self.eval_human}')
        if self.stream is None:
            self.stream = self.eval_human
        if self.stream and self.num_beams != 1:
            self.stream = False
            logger.info('Setting args.stream: False')

    def _init_ddp(self):
        if not is_dist():
            return
        eval_human = getattr(self, 'eval_human', False)
        assert not eval_human and not self.stream, (
            'In DDP scenarios, interactive interfaces and streaming output are not supported.'
            f'args.eval_human: {eval_human}, args.stream: {self.stream}')
        self._init_device()
        init_process_group(backend=self.ddp_backend, timeout=self.ddp_timeout)

    def __post_init__(self) -> None:
        if self.infer_backend == 'pt':
            self.infer_backend = 'transformers'  # compat swift3.x
            logger.warning('args.infer_backend: `pt` is deprecated, please use args.infer_backend: `transformers`.')
        BaseArguments.__post_init__(self)
        VllmArguments.__post_init__(self)
        self._init_vllm_async_engine()
        # Default to False for swift infer (non-encode tasks)
        if self.vllm_use_async_engine is None:
            self.vllm_use_async_engine = False
        self._init_result_path('infer_result')
        self._init_ddp()

    def _init_vllm_async_engine(self):
        """Initialize vllm_use_async_engine based on task_type.

        Encode tasks (embedding, seq_cls, reranker, generative_reranker) require
        async engine because vLLM's synchronous LLMEngine does not have the `encode` method.

        Note: This method only handles encode tasks. For non-encode tasks, the default value
        should be set by subclasses (DeployArguments sets True, RolloutArguments uses
        _set_default_engine_type, InferArguments defaults to False).
        """
        # Task types that require vLLM's encode() method, which is only available in AsyncLLMEngine
        encode_task_types = ('embedding', 'seq_cls', 'reranker', 'generative_reranker')
        is_vllm_encode_task = self.infer_backend == 'vllm' and self.task_type in encode_task_types

        if is_vllm_encode_task:
            if self.vllm_use_async_engine is None:
                self.vllm_use_async_engine = True
            elif not self.vllm_use_async_engine:
                raise ValueError(
                    f'task_type={self.task_type} requires vllm_use_async_engine=True. '
                    f'The synchronous vLLM LLMEngine does not support the `encode` method for encode tasks. '
                    f'Please set --vllm_use_async_engine true or remove the explicit false setting.')
