# Copyright (c) Alibaba, Inc. and its affiliates.
import os
import re
import signal
import sys
import time
from copy import deepcopy
from datetime import datetime
from functools import partial
from typing import List, Type

import gradio as gr
import json
import torch
from json import JSONDecodeError
from transformers.utils import is_torch_cuda_available, is_torch_npu_available

from swift.llm import DeployArguments, InferArguments, InferClient, InferRequest, RequestConfig
from swift.ui.base import BaseUI
from swift.ui.llm_infer.model import Model
from swift.ui.llm_infer.runtime import Runtime
from swift.utils import get_device_count, get_logger

logger = get_logger()


class LLMInfer(BaseUI):

    group = 'llm_infer'

    is_multimodal = True

    sub_ui = [Model, Runtime]

    locale_dict = {
        'generate_alert': {
            'value': {
                'zh': '请先部署模型',
                'en': 'Please deploy model first',
            }
        },
        'port': {
            'label': {
                'zh': '端口',
                'en': 'port'
            },
        },
        'llm_infer': {
            'label': {
                'zh': 'LLM推理',
                'en': 'LLM Inference',
            }
        },
        'load_alert': {
            'value': {
                'zh': '部署中，请点击"展示部署状态"查看',
                'en': 'Start to deploy model, '
                'please Click "Show running '
                'status" to view details',
            }
        },
        'loaded_alert': {
            'value': {
                'zh': '模型加载完成',
                'en': 'Model loaded'
            }
        },
        'port_alert': {
            'value': {
                'zh': '该端口已被占用',
                'en': 'The port has been occupied'
            }
        },
        'chatbot': {
            'value': {
                'zh': '对话框',
                'en': 'Chat bot'
            },
        },
        'infer_model_type': {
            'label': {
                'zh': 'Lora模块',
                'en': 'Lora module'
            },
            'info': {
                'zh': '发送给server端哪个LoRA，默认为`default`',
                'en': 'Which LoRA to use on server, default value is `default`'
            }
        },
        'prompt': {
            'label': {
                'zh': '请输入：',
                'en': 'Input:'
            },
        },
        'clear_history': {
            'value': {
                'zh': '清除对话信息',
                'en': 'Clear history'
            },
        },
        'submit': {
            'value': {
                'zh': '🚀 发送',
                'en': '🚀 Send'
            },
        },
        'gpu_id': {
            'label': {
                'zh': '选择可用GPU',
                'en': 'Choose GPU'
            },
            'info': {
                'zh': '选择训练使用的GPU号，如CUDA不可用只能选择CPU',
                'en': 'Select GPU to train'
            }
        },
    }

    choice_dict = BaseUI.get_choices_from_dataclass(InferArguments)
    default_dict = BaseUI.get_default_value_from_dataclass(InferArguments)
    arguments = BaseUI.get_argument_names(InferArguments)

    @classmethod
    def do_build_ui(cls, base_tab: Type['BaseUI']):
        with gr.TabItem(elem_id='llm_infer', label=''):
            default_device = 'cpu'
            device_count = get_device_count()
            if device_count > 0:
                default_device = '0'
            with gr.Blocks():
                infer_request = gr.State(None)
                Model.build_ui(base_tab)
                Runtime.build_ui(base_tab)
                with gr.Row():
                    gr.Dropdown(
                        elem_id='gpu_id',
                        multiselect=True,
                        choices=[str(i) for i in range(device_count)] + ['cpu'],
                        value=default_device,
                        scale=8)
                    infer_model_type = gr.Textbox(elem_id='infer_model_type', scale=4)
                    gr.Textbox(elem_id='port', lines=1, value='8000', scale=4)
                chatbot = gr.Chatbot(elem_id='chatbot', elem_classes='control-height')
                with gr.Row():
                    prompt = gr.Textbox(elem_id='prompt', lines=1, interactive=True)
                    with gr.Tabs(visible=cls.is_multimodal):
                        with gr.TabItem(label='Image'):
                            image = gr.Image(type='filepath')
                        with gr.TabItem(label='Video'):
                            video = gr.Video()
                        with gr.TabItem(label='Audio'):
                            audio = gr.Audio(type='filepath')

                with gr.Row():
                    clear_history = gr.Button(elem_id='clear_history')
                    submit = gr.Button(elem_id='submit')

                cls.element('load_checkpoint').click(
                    cls.deploy_model, list(base_tab.valid_elements().values()),
                    [cls.element('runtime_tab'), cls.element('running_tasks')])
                submit.click(
                    cls.send_message,
                    inputs=[
                        cls.element('running_tasks'),
                        cls.element('template'), prompt, image, video, audio, infer_request, infer_model_type,
                        cls.element('system'),
                        cls.element('max_new_tokens'),
                        cls.element('temperature'),
                        cls.element('top_k'),
                        cls.element('top_p'),
                        cls.element('repetition_penalty')
                    ],
                    outputs=[prompt, chatbot, image, video, audio, infer_request],
                    queue=True)

                clear_history.click(
                    fn=cls.clear_session, inputs=[], outputs=[prompt, chatbot, image, video, audio, infer_request])

                base_tab.element('running_tasks').change(
                    partial(Runtime.task_changed, base_tab=base_tab), [base_tab.element('running_tasks')],
                    list(cls.valid_elements().values()) + [cls.element('log')])
                Runtime.element('kill_task').click(
                    Runtime.kill_task,
                    [Runtime.element('running_tasks')],
                    [Runtime.element('running_tasks')] + [Runtime.element('log')],
                )

    @classmethod
    def deploy(cls, *args):
        deploy_args = cls.get_default_value_from_dataclass(DeployArguments)
        kwargs = {}
        kwargs_is_list = {}
        other_kwargs = {}
        more_params = {}
        more_params_cmd = ''
        keys = cls.valid_element_keys()
        for key, value in zip(keys, args):
            compare_value = deploy_args.get(key)
            compare_value_arg = str(compare_value) if not isinstance(compare_value, (list, dict)) else compare_value
            compare_value_ui = str(value) if not isinstance(value, (list, dict)) else value
            if key in deploy_args and compare_value_ui != compare_value_arg and value:
                if isinstance(value, str) and re.fullmatch(cls.int_regex, value):
                    value = int(value)
                elif isinstance(value, str) and re.fullmatch(cls.float_regex, value):
                    value = float(value)
                elif isinstance(value, str) and re.fullmatch(cls.bool_regex, value):
                    value = True if value.lower() == 'true' else False
                kwargs[key] = value if not isinstance(value, list) else ' '.join(value)
                kwargs_is_list[key] = isinstance(value, list) or getattr(cls.element(key), 'is_list', False)
            else:
                other_kwargs[key] = value
            if key == 'more_params' and value:
                try:
                    more_params = json.loads(value)
                except (JSONDecodeError or TypeError):
                    more_params_cmd = value

        kwargs.update(more_params)
        model = kwargs.get('model')
        if os.path.exists(model) and os.path.exists(os.path.join(model, 'args.json')):
            kwargs['ckpt_dir'] = kwargs.pop('model')
            with open(os.path.join(kwargs['ckpt_dir'], 'args.json'), 'r', encoding='utf-8') as f:
                _json = json.load(f)
                kwargs['model_type'] = _json['model_type']
                kwargs['train_type'] = _json['train_type']
        deploy_args = DeployArguments(
            **{
                key: value.split(' ') if key in kwargs_is_list and kwargs_is_list[key] else value
                for key, value in kwargs.items()
            })
        if deploy_args.port in Runtime.get_all_ports():
            raise gr.Error(cls.locale('port_alert', cls.lang)['value'])
        params = ''
        sep = f'{cls.quote} {cls.quote}'
        for e in kwargs:
            if isinstance(kwargs[e], list):
                params += f'--{e} {cls.quote}{sep.join(kwargs[e])}{cls.quote} '
            elif e in kwargs_is_list and kwargs_is_list[e]:
                all_args = [arg for arg in kwargs[e].split(' ') if arg.strip()]
                params += f'--{e} {cls.quote}{sep.join(all_args)}{cls.quote} '
            else:
                params += f'--{e} {cls.quote}{kwargs[e]}{cls.quote} '
        if 'port' not in kwargs:
            params += f'--port "{deploy_args.port}" '
        params += more_params_cmd + ' '
        devices = other_kwargs['gpu_id']
        devices = [d for d in devices if d]
        assert (len(devices) == 1 or 'cpu' not in devices)
        gpus = ','.join(devices)
        cuda_param = ''
        if gpus != 'cpu':
            if is_torch_npu_available():
                cuda_param = f'ASCEND_RT_VISIBLE_DEVICES={gpus}'
            elif is_torch_cuda_available():
                cuda_param = f'CUDA_VISIBLE_DEVICES={gpus}'
            else:
                cuda_param = ''
        now = datetime.now()
        time_str = f'{now.year}{now.month}{now.day}{now.hour}{now.minute}{now.second}'
        file_path = f'output/{deploy_args.model_type}-{time_str}'
        if not os.path.exists(file_path):
            os.makedirs(file_path, exist_ok=True)
        log_file = os.path.join(os.getcwd(), f'{file_path}/run_deploy.log')
        deploy_args.log_file = log_file
        params += f'--log_file "{log_file}" '
        params += '--ignore_args_error true '
        if sys.platform == 'win32':
            if cuda_param:
                cuda_param = f'set {cuda_param} && '
            run_command = f'{cuda_param}start /b swift deploy {params} > {log_file} 2>&1'
        else:
            run_command = f'{cuda_param} nohup swift deploy {params} > {log_file} 2>&1 &'
        return run_command, deploy_args, log_file

    @classmethod
    def deploy_model(cls, *args):
        run_command, deploy_args, log_file = cls.deploy(*args)
        logger.info(f'Running deployment command: {run_command}')
        os.system(run_command)
        gr.Info(cls.locale('load_alert', cls.lang)['value'])
        time.sleep(2)
        running_task = Runtime.refresh_tasks(log_file)
        return gr.update(open=True), running_task

    @classmethod
    def register_clean_hook(cls):
        signal.signal(signal.SIGINT, LLMInfer.signal_handler)
        if os.name != 'nt':
            signal.signal(signal.SIGTERM, LLMInfer.signal_handler)

    @staticmethod
    def signal_handler(*args, **kwargs):
        LLMInfer.clean_deployment()
        sys.exit(0)

    @classmethod
    def clear_session(cls):
        return '', [], gr.update(value=None), gr.update(value=None), gr.update(value=None), []

    @classmethod
    def _replace_tag_with_media(cls, infer_request: InferRequest):
        total_history = []
        messages = deepcopy(infer_request.messages)
        if messages[0]['role'] == 'system':
            messages.pop(0)
        for i in range(0, len(messages), 2):
            slices = messages[i:i + 2]
            if len(slices) == 2:
                user, assistant = slices
            else:
                user = slices[0]
                assistant = {'role': 'assistant', 'content': None}
            user['content'] = (user['content'] or '').replace('<image>', '').replace('<video>',
                                                                                     '').replace('<audio>', '').strip()
            for media in user['medias']:
                total_history.append([(media, ), None])
            if user['content'] or assistant['content']:
                total_history.append((user['content'], assistant['content']))
        return total_history

    @classmethod
    def agent_type(cls, response):
        if not response:
            return None
        if response.lower().endswith('observation:'):
            return 'react'
        if 'observation:' not in response.lower() and 'action input:' in response.lower():
            return 'toolbench'
        return None

    @classmethod
    def send_message(cls, running_task, template_type, prompt: str, image, video, audio, infer_request: InferRequest,
                     infer_model_type, system, max_new_tokens, temperature, top_k, top_p, repetition_penalty):

        if not infer_request:
            infer_request = InferRequest(messages=[])
        if system:
            if not infer_request.messages or infer_request.messages[0]['role'] != 'system':
                infer_request.messages.insert(0, {'role': 'system', 'content': system})
            else:
                infer_request.messages[0]['content'] = system
        if not infer_request.messages or infer_request.messages[-1]['role'] != 'user':
            infer_request.messages.append({'role': 'user', 'content': '', 'medias': []})
        media = image or video or audio
        media_type = 'images' if image else 'videos' if video else 'audios'
        if media:
            _saved_medias: List = getattr(infer_request, media_type)
            if not _saved_medias or _saved_medias[-1] != media:
                _saved_medias.append(media)
                infer_request.messages[-1]['content'] = infer_request.messages[-1]['content'] + f'<{media_type[:-1]}>'
                infer_request.messages[-1]['medias'].append(media)

        if not prompt:
            yield '', cls._replace_tag_with_media(infer_request), gr.update(value=None), gr.update(
                value=None), gr.update(value=None), infer_request
            return
        else:
            infer_request.messages[-1]['content'] = infer_request.messages[-1]['content'] + prompt

        _, args = Runtime.parse_info_from_cmdline(running_task)
        request_config = RequestConfig(
            temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty)
        request_config.stream = True
        request_config.stop = ['Observation:']
        request_config.max_tokens = max_new_tokens
        stream_resp_with_history = ''
        response = ''
        i = len(infer_request.messages) - 1
        for i in range(len(infer_request.messages) - 1, -1, -1):
            if infer_request.messages[i]['role'] == 'assistant':
                response = infer_request.messages[i]['content']
        agent_type = cls.agent_type(response)
        if i != len(infer_request.messages) - 1 and agent_type == 'toolbench':
            infer_request.messages[i + 1]['role'] = 'tool'

        chat = not template_type.endswith('generation')
        _infer_request = deepcopy(infer_request)
        for m in _infer_request.messages:
            if 'medias' in m:
                m.pop('medias')
        model_kwargs = {}
        if infer_model_type:
            model_kwargs = {'model': infer_model_type}
        gen_list = InferClient(
            port=args['port'], ).infer(
                infer_requests=[_infer_request], request_config=request_config, **model_kwargs)
        if infer_request.messages[-1]['role'] != 'assistant':
            infer_request.messages.append({'role': 'assistant', 'content': ''})
        for chunk in gen_list[0]:
            if chunk is None:
                continue
            stream_resp_with_history += chunk.choices[0].delta.content if chat else chunk.choices[0].text
            infer_request.messages[-1]['content'] = stream_resp_with_history
            yield '', cls._replace_tag_with_media(infer_request), gr.update(value=None), gr.update(
                value=None), gr.update(value=None), infer_request