# Copyright (c) Alibaba, Inc. and its affiliates.
from typing import Type

import gradio as gr

from swift.ui.base import BaseUI
from swift.utils import get_logger

logger = get_logger()


class Eval(BaseUI):

    group = 'llm_eval'

    locale_dict = {
        'eval_backend': {
            'label': {
                'zh': '评测后端',
                'en': 'Eval backend'
            },
            'info': {
                'zh': '选择评测后端',
                'en': 'Select eval backend'
            }
        },
        'eval_dataset': {
            'label': {
                'zh': '评测数据集',
                'en': 'Evaluation dataset'
            },
            'info': {
                'zh': '选择评测数据集，支持多选 (先选择评测后端)',
                'en': 'Select eval dataset, multiple datasets supported (select eval backend first)'
            }
        },
        'eval_limit': {
            'label': {
                'zh': '评测数据个数',
                'en': 'Eval numbers for each dataset'
            },
            'info': {
                'zh': '每个评测集的取样数',
                'en': 'Number of rows sampled from each dataset'
            }
        },
        'eval_output_dir': {
            'label': {
                'zh': '评测输出目录',
                'en': 'Eval output dir'
            },
            'info': {
                'zh': '评测结果的输出目录',
                'en': 'The dir to save the eval results'
            }
        },
        'custom_eval_config': {
            'label': {
                'zh': '自定义数据集评测配置',
                'en': 'Custom eval config'
            },
            'info': {
                'zh': '可以使用该配置评测自己的数据集，详见github文档的评测部分',
                'en': 'Use this config to eval your own datasets, check the docs in github for details'
            }
        },
        'eval_url': {
            'label': {
                'zh': '评测链接',
                'en': 'The eval url'
            },
            'info': {
                'zh':
                'OpenAI样式的评测链接(如：http://localhost:8080/v1/chat/completions)，用于评测接口（模型类型输入为实际模型类型）',
                'en':
                'The OpenAI style link(like: http://localhost:8080/v1/chat/completions) for '
                'evaluation(Input actual model type into model_type)'
            }
        },
        'api_key': {
            'label': {
                'zh': '接口token',
                'en': 'The url token'
            },
            'info': {
                'zh': 'eval_url的token',
                'en': 'The token used with eval_url'
            }
        },
        'infer_backend': {
            'label': {
                'zh': '推理框架',
                'en': 'Infer backend'
            },
        }
    }

    @classmethod
    def do_build_ui(cls, base_tab: Type['BaseUI']):
        try:
            from swift.llm.argument.eval_args import EvalArguments
            eval_dataset_dict = EvalArguments.list_eval_dataset()
            default_backend = EvalArguments.eval_backend
        except Exception as e:
            logger.warn(e)
            eval_dataset_dict = {}
            default_backend = None

        with gr.Row():
            gr.Dropdown(elem_id='eval_backend', choices=list(eval_dataset_dict.keys()), value=default_backend, scale=20)
            gr.Dropdown(
                elem_id='eval_dataset',
                is_list=True,
                choices=eval_dataset_dict.get(default_backend, []),
                multiselect=True,
                allow_custom_value=True,
                scale=20)
            gr.Textbox(elem_id='eval_limit', scale=20)
            gr.Dropdown(elem_id='infer_backend', scale=20)
        with gr.Row():
            gr.Textbox(elem_id='custom_eval_config', scale=20)
            gr.Textbox(elem_id='eval_output_dir', scale=20)
            gr.Textbox(elem_id='eval_url', scale=20)
            gr.Textbox(elem_id='api_key', scale=20)

        def update_eval_dataset(backend):
            return gr.update(choices=eval_dataset_dict[backend])

        cls.element('eval_backend').change(update_eval_dataset, [cls.element('eval_backend')],
                                           [cls.element('eval_dataset')])