# Copyright (c) ModelScope Contributors. All rights reserved.
import json
import os
import peft
import shutil
from dataclasses import dataclass, field, fields
from packaging import version
from typing import Any, Dict, List, Literal, Optional, Union

import swift
from swift.dataset import load_dataset
from swift.hub import get_hub
from swift.model import get_ckpt_dir, get_model_processor, load_by_unsloth
from swift.ray_utils import RayArguments
from swift.template import Template, get_template
from swift.tuner_plugin import tuners_map
from swift.utils import (Processor, check_json_format, get_dist_setting, get_logger, import_external_file, is_dist,
                         is_master, json_parse_to_dict, safe_snapshot_download, set_device, use_hf_hub)
from .data_args import DataArguments
from .generation_args import GenerationArguments
from .model_args import ModelArguments
from .quant_args import QuantizeArguments
from .template_args import TemplateArguments

logger = get_logger()


def get_supported_tuners():
    return {'lora', 'full', 'longlora', 'adalora', 'llamapro', 'adapter', 'vera', 'boft', 'fourierft', 'reft', 'bone'
            } | set(tuners_map.keys())


def _patch_peft():
    """Patch peft functions that are incompatible with SWIFT.

    1. _maybe_shard_state_dict_for_tp: TP sharding is not used by SWIFT, and causes errors
       when torch.distributed is initialized (e.g. MoE training with target_parameters).
    2. _maybe_shard_state_dict_for_tp internal logic accesses base_layer.weight.device which
       fails for expert modules that don't have a `weight` attribute.
    """
    if version.parse(peft.__version__) >= version.parse('0.19.0'):
        from peft.utils import save_and_load
        save_and_load._maybe_shard_state_dict_for_tp = lambda model, state_dict, adapter_name: None


@dataclass
class BaseArguments(GenerationArguments, QuantizeArguments, DataArguments, TemplateArguments, ModelArguments,
                    RayArguments):
    """BaseArguments class is a dataclass that inherits from multiple argument classes.

    This class consolidates arguments from GenerationArguments, QuantizeArguments, DataArguments,
    TemplateArguments, ModelArguments, RayArguments.

    Args:
        tuner_backend (str): The tuner backend to use. Choices are 'peft' or 'unsloth'. Default is 'peft'.
        tuner_type (str): The tuner type. Choices include 'lora', 'full', 'longlora', 'adalora', 'llamapro',
            'adapter', 'vera', 'boft', 'fourierft', 'reft'. Default is 'lora'.
        adapters (List[str]): A list of adapter IDs or paths. This is typically used for inference or deployment.
            It can also resume training by only loading adapter weights, differing from `resume_from_checkpoint`
            which also loads optimizer states. Default is [].
        external_plugins (List[str]): A list of external 'plugin.py' files to be registered and imported into
            the plugin module. Default is [].
        seed (int): The global random seed for reproducibility. Note that this does not affect `data_seed`,
            which controls dataset randomization. Default is 42.
        model_kwargs (Optional[str]): Additional keyword arguments for specific models, passed as a JSON string
            (e.g., '{"key": "value"}'). It's recommended to use the same arguments for inference as for training.
            Default is None.
        enable_npu_model_patch (bool): Whether to enable model-related NPU patches. Default is True.
        load_args (bool): Whether to load `args.json` from a checkpoint when using `--resume_from_checkpoint`,
            `--model`, or `--adapters`. Defaults to True for inference/export and False for training. Usually,
            this does not need to be modified. Default is True.
        load_data_args (bool): If True, will also load data-related arguments from `args.json`. This is useful
            for running inference on the same validation split used during training. Default is False.
        packing (bool): Whether to enable packing of datasets. Default is False.
        packing_length (Optional[int]): Length of packing. Default is None.
        packing_num_proc (int): Number of processes used for packing, Default is 1.
        packing_strategy (Literal['binpack', 'sequential']): Packing algorithm. 'binpack' (default) uses
            best-fit-decreasing bin packing (reorders samples); 'sequential' uses order-preserving greedy
            packing (next-fit: a single open pack, flushed when the next sample doesn't fit) so the sample
            order / pack boundaries follow a sequential sampler (use packing_num_proc=1). Default is 'binpack'.
        lazy_tokenize (Optional[bool]): Whether to enable lazy tokenization. Default is None.
        use_hf (bool): Whether to use Hugging Face for downloading/uploading models and datasets. If False,
            ModelScope is used. Default is False.
        hub_token (Optional[str]): The authentication token for ModelScope or Hugging Face Hub. Default is None.
        ddp_timeout (int): Timeout for DDP (Distributed Data Parallel) operations, in seconds. Default is 18000000.
        ddp_backend (Optional[str]): The backend for DDP. Choices include "nccl", "gloo", "mpi", "ccl", "hccl",
            "cncl", "mccl". If None, it will be automatically selected. Default is None.
        ignore_args_error (bool): Whether to ignore argument errors. This is useful for compatibility with Jupyter
            notebooks. Default is False.
        use_swift_lora (bool): Whether to use swift lora. This is a compatible argument. Default is False.
    """
    tuner_backend: Literal['peft', 'unsloth'] = 'peft'
    tuner_type: str = field(default='lora', metadata={'help': f'tuner_type choices: {list(get_supported_tuners())}'})
    adapters: List[str] = field(default_factory=list)
    external_plugins: List[str] = field(default_factory=list)
    # This parameter is kept for swift3.x compatibility. Please use `external_plugins` as a replacement.
    custom_register_path: List[str] = field(default_factory=list)

    seed: int = 42
    model_kwargs: Optional[Union[dict, str]] = None
    enable_npu_model_patch: bool = True
    load_args: bool = True
    load_data_args: bool = False
    # dataset
    packing: bool = False
    packing_length: Optional[int] = None
    packing_num_proc: int = 1
    packing_strategy: Literal['binpack', 'sequential'] = 'binpack'
    lazy_tokenize: Optional[bool] = None
    # hub
    use_hf: bool = False
    # None: use env var `MODELSCOPE_API_TOKEN`
    hub_token: Optional[str] = field(
        default=None, metadata={'help': 'SDK token can be found in https://modelscope.cn/my/myaccesstoken'})
    # dist
    ddp_timeout: int = 18000000
    ddp_backend: Optional[str] = None

    # extra
    ignore_args_error: bool = False  # True: notebook compatibility
    use_swift_lora: bool = False  # True for using tuner_backend == swift, don't specify this unless you know what you are doing # noqa

    def _prepare_training_args(self, training_args: Dict[str, Any]) -> None:
        pass

    def _init_lazy_tokenize(self):
        if self.lazy_tokenize is None:
            if self.cached_dataset or self.cached_val_dataset:
                self.lazy_tokenize = False
            elif (self.model_meta is not None and self.model_meta.is_multimodal and not self.streaming
                  and not self.packing and not getattr(self, 'group_by_length', False)):
                self.lazy_tokenize = True
            else:
                self.lazy_tokenize = False
            logger.info(f'Setting args.lazy_tokenize: {self.lazy_tokenize}')
        if self.lazy_tokenize:
            if self.packing:
                raise ValueError('Packing and lazy_tokenize are incompatible.')
            if self.streaming:
                raise ValueError('Streaming and lazy_tokenize are incompatible.')

    def _import_external_plugins(self):
        if isinstance(self.external_plugins, str):
            self.external_plugins = [self.external_plugins]
        # swift v3.x compatibility
        if isinstance(self.custom_register_path, str):
            self.custom_register_path = [self.custom_register_path]
        if self.custom_register_path:
            self.external_plugins += self.custom_register_path

        if not self.external_plugins:
            return
        for external_plugin in self.external_plugins:
            import_external_file(external_plugin)
        logger.info(f'Successfully imported external_plugins: {self.external_plugins}.')

    @staticmethod
    def _check_is_adapter(adapter_dir: str) -> bool:
        if (os.path.exists(os.path.join(adapter_dir, 'adapter_config.json'))
                or os.path.exists(os.path.join(adapter_dir, 'default', 'adapter_config.json'))
                or os.path.exists(os.path.join(adapter_dir, 'reft'))):
            return True
        return False

    def _init_adapters(self):
        if isinstance(self.adapters, str):
            self.adapters = [self.adapters]
        self.adapters = [
            safe_snapshot_download(adapter, use_hf=self.use_hf, hub_token=self.hub_token) for adapter in self.adapters
        ]

    def __post_init__(self):
        _patch_peft()
        self.swift_version = swift.__version__
        if self.use_hf or use_hf_hub():
            self.use_hf = True
            os.environ['USE_HF'] = '1'
        self._init_adapters()
        self._init_ckpt_dir()
        self._import_external_plugins()
        self._init_model_kwargs()
        # The Seq2SeqTrainingArguments has a property called world_size, which cannot be assigned a value.
        self.rank, self.local_rank, self.global_world_size, self.local_world_size = get_dist_setting()
        logger.info(f'rank: {self.rank}, local_rank: {self.local_rank}, '
                    f'world_size: {self.global_world_size}, local_world_size: {self.local_world_size}')
        if self.tuner_type not in tuners_map:  # build-in tuner
            for adapter in self.adapters:
                assert self._check_is_adapter(adapter), (
                    f'`{adapter}` is not an adapter, please try using `--model` to pass it.')
        ModelArguments.__post_init__(self)
        QuantizeArguments.__post_init__(self)
        TemplateArguments.__post_init__(self)
        DataArguments.__post_init__(self)
        RayArguments.__post_init__(self)
        self._init_stream()
        if self.max_length is None and self.model_info is not None:
            self.max_length = self.model_info.max_model_len
        if self.packing and self.packing_length is None:
            self.packing_length = self.max_length
        self._init_lazy_tokenize()
        self.hub = get_hub(self.use_hf)
        if self.hub.try_login(self.hub_token):
            logger.info('hub login successful!')

    def _init_model_kwargs(self):
        """Prepare model kwargs and set them to the env"""
        self.model_kwargs: Dict[str, Any] = json_parse_to_dict(self.model_kwargs)
        for k, v in self.model_kwargs.items():
            k = k.upper()
            os.environ[k] = str(v)

    @property
    def is_adapter(self) -> bool:
        return self.tuner_type not in {'full'}

    @property
    def supported_tuners(self):
        return get_supported_tuners()

    @property
    def adapters_can_be_merged(self):
        return {'lora', 'longlora', 'llamapro', 'adalora'}

    @classmethod
    def from_pretrained(cls, checkpoint_dir: str):
        self = super().__new__(cls)
        self.load_data_args = True
        self.ckpt_dir = checkpoint_dir
        self.load_args_from_ckpt()
        all_keys = list(f.name for f in fields(BaseArguments))
        for key in all_keys:
            if not hasattr(self, key):
                setattr(self, key, None)
        return self

    def _init_ckpt_dir(self, adapters=None):
        # compat megatron
        model = self.model or getattr(self, 'mcore_model', None)
        adapters = adapters or self.adapters or getattr(self, 'mcore_adapter', None)
        if isinstance(adapters, str):
            adapters = [adapters]
        self.ckpt_dir = get_ckpt_dir(model, adapters)
        if self.ckpt_dir and self.load_args:
            self.load_args_from_ckpt()

    def load_args_from_ckpt(self) -> None:
        args_path = os.path.join(self.ckpt_dir, 'args.json')
        assert os.path.exists(args_path), f'args_path: {args_path}'
        with open(args_path, 'r', encoding='utf-8') as f:
            old_args = json.load(f)
        force_load_keys = [
            # base_args
            'tuner_type',
            # model_args
            'task_type',
            # quant_args
            'bnb_4bit_quant_type',
            'bnb_4bit_use_double_quant',
        ]
        # If the current value is None or an empty list and it is among the following keys
        load_keys = [
            'external_plugins',
            # model_args
            'model',
            'model_type',
            'model_revision',
            'torch_dtype',
            'attn_impl',
            'experts_impl',
            'new_special_tokens',
            'num_labels',
            'problem_type',
            'rope_scaling',
            'max_model_len',
            # quant_args
            'quant_method',
            'quant_bits',
            'hqq_axis',
            'bnb_4bit_compute_dtype',
            # template_args
            'template',
            'system',
            'truncation_strategy',
            'agent_template',
            'norm_bbox',
            'use_chat_template',
            'response_prefix',
        ]
        data_keys = list(f.name for f in fields(DataArguments))
        swift_version = old_args.get('swift_version')
        if swift_version is None or version.parse(swift_version) < version.parse('4.0.0.dev'):
            load_keys.remove('model_type')
        for key, old_value in old_args.items():
            if old_value is None:
                continue
            if key in force_load_keys or self.load_data_args and key in data_keys:
                setattr(self, key, old_value)
            value = getattr(self, key, None)
            if key in load_keys and (value is None or isinstance(value, (list, tuple)) and len(value) == 0):
                setattr(self, key, old_value)
        logger.info(f'Successfully loaded {args_path}.')

    def save_args(self, output_dir=None) -> None:
        if is_master():
            output_dir = output_dir or self.output_dir
            os.makedirs(output_dir, exist_ok=True)
            fpath = os.path.join(output_dir, 'args.json')
            logger.info(f'The {self.__class__.__name__} will be saved in: {fpath}')
            with open(fpath, 'w', encoding='utf-8') as f:
                json.dump(check_json_format(self.__dict__), f, ensure_ascii=False, indent=2)
            config_file = os.getenv('SWIFT_CONFIG_FILE')
            if config_file:
                shutil.copy(config_file, output_dir)

    def _init_device(self):
        if is_dist():
            set_device()

    def get_template(self, processor: Optional[Processor] = None, **kwargs) -> Template:
        if processor is None:
            processor = self.get_model_processor(load_model=False)[1]
        template_kwargs = self.get_template_kwargs()
        if 'template_type' in kwargs:
            template_type = kwargs.get('template_type')
        else:
            template_type = self.template
        template_kwargs['template_type'] = template_type
        template = get_template(processor, **template_kwargs)
        return template

    def get_model_processor(self,
                            *,
                            model=None,
                            model_type=None,
                            revision=None,
                            task_type=None,
                            num_labels=None,
                            **kwargs):
        if self.tuner_backend == 'unsloth':
            return load_by_unsloth(self)
        res = self.get_model_kwargs()
        res.update(kwargs)
        # compat rlhf
        res['model_id_or_path'] = model or self.model
        res['model_type'] = model_type or self.model_type
        res['revision'] = revision or self.model_revision
        res['task_type'] = task_type or self.task_type
        res['num_labels'] = num_labels or self.num_labels

        return get_model_processor(**res)

    def load_dataset(self):
        dataset_kwargs = self.get_dataset_kwargs()
        train_dataset, val_dataset = None, None
        if self.dataset:
            train_dataset, val_dataset = load_dataset(
                self.dataset,
                split_dataset_ratio=self.split_dataset_ratio,
                shuffle=self.dataset_shuffle,
                **dataset_kwargs)
        if len(self.val_dataset) > 0:
            # Loading val dataset
            dataset_kwargs.pop('interleave_prob', None)
            _, val_dataset = load_dataset(
                self.val_dataset, split_dataset_ratio=1.0, shuffle=self.val_dataset_shuffle, **dataset_kwargs)
            assert self.split_dataset_ratio == 0.
        return train_dataset, val_dataset
