# Copyright (c) ModelScope Contributors. All rights reserved.
import os
import shutil
import torch.distributed as dist
from transformers.utils import strtobool
from typing import List, Optional, Union

from swift.megatron.arguments import MegatronExportArguments
from swift.megatron.convert import test_convert_precision
from swift.megatron.model import get_mcore_model
from swift.megatron.utils import load_mcore_checkpoint, prepare_mcore_model, save_mcore_checkpoint
from swift.pipelines import SwiftPipeline, prepare_model_template
from swift.utils import disable_safe_ddp_context_use_barrier, get_logger, is_master

logger = get_logger()


class MegatronExport(SwiftPipeline):
    args_class = MegatronExportArguments
    args: args_class

    def run(self):
        os.environ['DISABLE_MP_DDP'] = 'true'
        args = self.args
        if args.to_hf:
            self.convert_mcore2hf()
        elif args.to_mcore:
            self.convert_hf2mcore()

    def convert_mcore2hf(self) -> None:
        args = self.args
        args.experts_impl = 'eager'
        download_model = args.model is not None
        _, template = prepare_model_template(args, load_model=False, download_model=download_model)
        self.processor = template.processor
        hf_config = self.processor.model_info.config
        mg_model = get_mcore_model(args, hf_config)[0]
        logger.info('Megatron model created successfully.')
        bridge = mg_model.config.bridge
        if args.mcore_model is not None:
            load_mcore_checkpoint(args, [mg_model], load_arg='mcore_model')
        elif args.model is not None:
            bridge.load_weights([mg_model], args.model_info.model_dir)
        else:
            raise ValueError('Please specify `--mcore_model` or `--model`.')
        if args.adapters or args.mcore_adapter is not None:
            peft_model = prepare_mcore_model(args, mg_model)
            if args.mcore_adapter is not None:
                load_mcore_checkpoint(args, [mg_model], load_arg='mcore_adapter')
            elif args.adapters:
                assert len(args.adapters) == 1, 'Currently only support one adapter'
                bridge.load_weights([mg_model], args.adapters[0], peft_format=True)
            if args.merge_lora:
                logger.info('Merge LoRA...')
                mg_model = peft_model.merge_and_unload()
        logger.info('Converting weights and saving the model...')
        save_peft_format = args.tuner_type == 'lora' and not args.merge_lora
        bridge.save_weights(
            [mg_model],
            args.output_dir,
            peft_format=save_peft_format,
            args=args,
            processor=self.processor,
        )
        if is_master():
            if args.ckpt_dir:
                args_path = os.path.join(args.ckpt_dir, 'args.json')
                shutil.copy(args_path, os.path.join(args.output_dir, 'args.json'))
            else:
                args.save_args(args.output_dir)
        if args.test_convert_precision:
            with disable_safe_ddp_context_use_barrier():
                if save_peft_format:
                    kwargs = {'adapters': [args.output_dir]}
                else:
                    kwargs = {'model': args.output_dir, 'torch_dtype': None, 'adapters': []}
                device_map = args.device_map or 'auto'
                hf_model, template = prepare_model_template(
                    args, device_map=device_map, **kwargs) if is_master() else (None, template)
            test_convert_precision(args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype)
            dist.barrier()

    def convert_hf2mcore(self) -> None:
        args = self.args
        args.experts_impl = 'eager'
        download_model = args.model is not None
        _, template = prepare_model_template(args, load_model=False, download_model=download_model)
        self.processor = template.processor
        hf_config = self.processor.model_info.config
        mg_model = get_mcore_model(args, hf_config)[0]
        logger.info('Megatron model created successfully.')
        bridge = mg_model.config.bridge
        if args.model is not None:
            bridge.load_weights([mg_model], args.model_info.model_dir)
        elif args.mcore_model is not None:
            load_mcore_checkpoint(args, [mg_model], load_arg='mcore_model')
        else:
            raise ValueError('Please specify `--mcore_model` or `--model`.')
        dist.barrier()
        if args.adapters or args.mcore_adapter is not None:
            peft_model = prepare_mcore_model(args, mg_model)
            if args.adapters:
                assert len(args.adapters) == 1, 'Currently only support one adapter'
                bridge.load_weights([mg_model], args.adapters[0], peft_format=True)
            elif args.mcore_adapter is not None:
                load_mcore_checkpoint(args, [mg_model], load_arg='mcore_adapter')
            if args.merge_lora:
                logger.info('Merge LoRA...')
                mg_model = peft_model.merge_and_unload()
        logger.info('Successfully transferred HF model weights to MG model.')
        _test_convert_precision = strtobool(os.getenv('SWIFT_TEST_CONVERT_PRECISION', '0'))
        if not _test_convert_precision:
            args.save_args(args.output_dir)
            logger.info('Saving the model...')
            save_peft_format = args.tuner_type == 'lora' and not args.merge_lora
            save_mcore_checkpoint(args, [mg_model], peft_format=save_peft_format)
        # hf_model does not support loading args.mcore_adapter, so test_convert_precision cannot be performed
        support_convert_precision = args.mcore_adapter is None
        if args.test_convert_precision:
            if support_convert_precision:
                with disable_safe_ddp_context_use_barrier():
                    device_map = args.device_map or 'auto'
                    hf_model, template = prepare_model_template(
                        args, device_map=device_map) if is_master() else (None, template)
                test_convert_precision(args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype)
                dist.barrier()
            else:
                logger.warning('Skip test_convert_precision because `--mcore_adapter` is specified.')


def megatron_export_main(args: Optional[Union[List[str], MegatronExportArguments]] = None):
    return MegatronExport(args).main()