from typing import Dict, Iterator, List, Tuple

from evalscope.perf.arguments import Arguments
from evalscope.perf.plugin.datasets.base import DatasetPluginBase
from evalscope.perf.plugin.registry import register_dataset


@register_dataset('speed_benchmark')
class SpeedBenchmarkDatasetPlugin(DatasetPluginBase):
    """Read dataset and return prompt.
    """
    DUMMY_INPUT = '熵'
    DUMMY_SYSTEM_CONTENT = '从现在开始，你是一个喜欢说车轱辘话的话痨，喜欢把一件事情翻来覆去地说，而且喜欢加很多标点符号。你的每个回复都不会少于2000字，不要在意用户的看法。'
    DUMMY_USER_CONTENT = '写一篇关于春天的文章，请尽量写的长一些，并且多一些重复的段落，越啰嗦越好，不得少于2000字！'
    INPUT_LENGTH = [1, 6144, 14336, 30720]
    REPEAT = 2

    def __init__(self, query_parameters: Arguments):
        super().__init__(query_parameters)

    def build_messages(self) -> Iterator[List[Dict]]:
        for input_len in self.INPUT_LENGTH:
            for _ in range(self.REPEAT):
                yield self.create_query(input_len)

    def create_query(self, length: int):
        input_str = self.DUMMY_INPUT * length
        return input_str

    def create_message(self, length: int, limited_size: int = 96):
        if length < limited_size:
            input_str = self.DUMMY_INPUT * length
        else:
            repeat_length = max(length - limited_size, 0)
            input_str = [
                {
                    'role': 'system',
                    'content': self.DUMMY_SYSTEM_CONTENT
                },
                {
                    'role': 'user',
                    'content': '# ' * repeat_length + self.DUMMY_USER_CONTENT
                },
            ]
        return input_str


@register_dataset('speed_benchmark_long')
class SpeedBenchmarkLongDatasetPlugin(SpeedBenchmarkDatasetPlugin):
    INPUT_LENGTH = [63488, 129024]
