# Copyright (c) OpenMMLab. All rights reserved.
from functools import partial

from torch.utils.data import ConcatDataset

from xtuner.dataset import MOSSSFTDataset
from xtuner.dataset.collate_fns import default_collate_fn


def moss_003_sft_dataset(tokenizer,
                         plugins_data_file=None,
                         no_plugins_data_file=None,
                         bot_name=None,
                         max_length=2048):
    plugins = moss_003_sft_plugins_dataset(
        tokenizer,
        data_file=plugins_data_file,
        bot_name=bot_name,
        max_length=max_length)
    no_plugins = moss_003_sft_no_plugins_dataset(
        tokenizer,
        data_file=no_plugins_data_file,
        bot_name=bot_name,
        max_length=max_length)
    dataset = ConcatDataset([plugins, no_plugins])
    return dataset


def moss_003_sft_data_collator(return_hf_format=False):
    return partial(default_collate_fn, return_hf_format=return_hf_format)


def moss_003_sft_no_plugins_dataset(tokenizer,
                                    data_file=None,
                                    bot_name=None,
                                    max_length=2048):

    # Download data from https://huggingface.co/datasets/fnlp/moss-003-sft-data
    if data_file is None:
        data_file = './data/moss-003-sft-no-tools.jsonl'
    dataset = MOSSSFTDataset(
        data_file=data_file,
        bot_name=bot_name,
        tokenizer=tokenizer,
        max_length=max_length)

    return dataset


def moss_003_sft_no_plugins_data_collator(return_hf_format=False):
    return partial(default_collate_fn, return_hf_format=return_hf_format)


def moss_003_sft_plugins_dataset(tokenizer,
                                 data_file=None,
                                 bot_name=None,
                                 max_length=2048):

    # Download data from https://huggingface.co/datasets/fnlp/moss-003-sft-data
    if data_file is None:
        data_file = './data/conversations_with_tools_with_inner_instruction_no_text2image_train_all_random_meta0.5_0.1_0.01_moss_0709.jsonl'  # noqa: E501
    dataset = MOSSSFTDataset(
        data_file=data_file,
        bot_name=bot_name,
        tokenizer=tokenizer,
        max_length=max_length)

    return dataset


def moss_003_sft_plugins_data_collator(return_hf_format=False):
    return partial(default_collate_fn, return_hf_format=return_hf_format)