# coding=utf-8

"""
@fileName       :   runner.py
@data           :   2024/8/27
@author         :   jiangmenggui@hosonsoft.com
"""
import datetime
import json
import os.path
import queue
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from jinja2 import Environment, FileSystemLoader, select_autoescape

from lljz_tools.color import Color
from lljz_tools.console_table import ConsoleTable
from lljz_tools.simple_pref_test._task import find_tasks, TaskGroup
from lljz_tools.simple_pref_test.core import TaskResultSet, init_result, strf_time


class PrefRunner:

    def __init__(
            self,
            __name: str = '性能测试', /, *,
            virtual_users: int | None = 1, user_add_interval: float | None = 0.1,
            pre_task: float | None = None,
            run_seconds: int | float | None = 10,
            save_directory: str | None = None,
    ):
        assert pre_task is not None or (virtual_users is not None and user_add_interval is not None)
        assert pre_task is None or (isinstance(pre_task, float | int) and pre_task > 0)
        assert virtual_users is None or (isinstance(virtual_users, int) and virtual_users > 0)
        assert user_add_interval is None or (isinstance(user_add_interval, float | int) and user_add_interval >= 0)
        self.name = __name
        self.save_directory = save_directory or os.path.abspath(f'./pref_test_result/{self.name}')
        self.run_seconds = run_seconds
        self.virtual_users = virtual_users
        self.user_add_interval = user_add_interval
        self.pre_task = pre_task
        self.pool = ThreadPoolExecutor(max_workers=max(10, virtual_users, pre_task or 0))
        self.tasks = find_tasks()
        self.running = True

    def _run_with_thread_num(self):
        """按照固定线程数启动"""

        def run():
            group = TaskGroup(self.tasks)
            while True:
                task_set = group.get()
                for task in task_set.tasks:
                    if time.time() - TaskResultSet.start_time > self.run_seconds:
                        group.teardown_all()
                        return

                    task_set.setup(task.name)
                    task()
                    task_set.teardown(task.name)

        for _ in range(self.virtual_users):
            self.pool.submit(run)
            time.sleep(self.user_add_interval)

    def _run_with_pre_task(self):
        """按照每秒启动的线程数启动"""
        group_queue = queue.Queue()
        k = 1
        pre_task = self.pre_task
        while pre_task > 10:
            pre_task = pre_task / 10
            k = k * 10
        interval = 1 / pre_task

        def run():
            nonlocal group_queue
            if group_queue.empty():
                group = TaskGroup(self.tasks)
            else:
                group = group_queue.get()
            task_set = group.get()
            for task in task_set.tasks:
                if time.time() - TaskResultSet.start_time > self.run_seconds:
                    group.teardown_all()
                    return
                task_set.setup(task.name)
                task()
                task_set.teardown(task.name)
            group_queue.put(group)

        while time.time() - TaskResultSet.start_time < self.run_seconds and self.running:
            for _ in range(k):
                self.pool.submit(run)
            time.sleep(interval)

    def run_task(self):
        if self.pre_task:
            self._run_with_pre_task()
        else:
            self._run_with_thread_num()

    def print_start_info(self):
        print(f"开始测试：{Color.green(self.name)}")
        if self.pre_task:
            print(
                f"{Color.yellow('==========测试任务启动参数==========')}\n"
                f"   {Color.thin_magenta('每秒任务数')} : {self.pre_task}\n"
                f"     {Color.thin_magenta('运行时间')} : {self.run_seconds}s\n"
                f"     {Color.thin_magenta('任务总数')} : {len(set(self.tasks))}\n"
                f"     {Color.thin_magenta('启动时间')} : {strf_time(TaskResultSet.start_time)}\n"
                f"{Color.yellow('====================================')}\n"
            )
        else:
            print(
                f"{Color.yellow('==========测试任务启动参数==========')}\n"
                f"   {Color.thin_magenta('并发线程数')} : {self.virtual_users}\n"
                f" {Color.thin_magenta('线程启动间隔')} : {self.user_add_interval}s\n"
                f"     {Color.thin_magenta('运行时间')} : {self.run_seconds}s\n"
                f"     {Color.thin_magenta('任务总数')} : {len(set(self.tasks))}\n"
                f"     {Color.thin_magenta('启动时间')} : {strf_time(TaskResultSet.start_time)}\n"
                f"{Color.yellow('====================================')}\n"
            )

    def start(self):
        TaskResultSet.start_time = time.time()
        self.print_start_info()
        try:
            self.run_task()
            self.pool.shutdown(wait=True)
        finally:
            end_time = TaskResultSet.end_time = time.time()
            table_data, rps, qps, response = init_result()
            table = ConsoleTable(table_data, caption="性能测试结果")
            print(table)
            print(f'\n\n{Color.green("测试完成！")}[完成时间：{Color.thin_cyan(strf_time(end_time))}]')
            file = self.save_to_html(table_data, rps, qps, response)
            print(f'结果保存至：{Color.thin_blue(file)}')

    def save_to_html(self, table, rps, qps, response):
        def get_avg(data):
            return int(sum(data) / len(data) * 1000)

        def init_table_data(table_data):
            return {"header": list(table_data[0].keys()), "rows": [list(row.values()) for row in table_data]}

        def init_data(data, func):
            tps = sorted(data.items())
            keys = set()
            for _, v in tps:
                for k in v:
                    keys.add(k)
            values = {k: [] for k in keys}
            for _, v in tps:
                for k in keys:
                    values[k].append(func(k, v))
            return json.dumps({
                'keys': [k.split(' ')[1] for k, v in tps],
                "values": [{"name": k, "data": v} for k, v in values.items()]
            }, ensure_ascii=False)

        if not os.path.exists(self.save_directory):
            os.makedirs(self.save_directory)
        filename = f'{self.name}_{datetime.datetime.now().strftime("%Y-%m-%d %H_%M_%S")}.html'
        file = os.path.join(self.save_directory, filename)
        source = Path(__file__).parent.parent / 'source'
        env = Environment(
            loader=FileSystemLoader(source),
            autoescape=select_autoescape(['html', 'xml'])
        )
        with open(source / 'echart.js', 'r', encoding='u8') as f:
            echart_js = f.read()
        template = env.get_template('result.html')
        run_arguments = [
            f"并发线程数 : {self.virtual_users}",
            f"线程启动间隔 : {self.user_add_interval}s"
        ]

        context = {
            "echartJs": echart_js,
            "title": self.name,
            'arguments': [
                *run_arguments,
                f"运行时间 : {self.run_seconds}s",
                f"任务总数 : {len(set(self.tasks))}",
                f"启动时间 : {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(TaskResultSet.start_time))}",
                f"结束时间 : {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(TaskResultSet.end_time))}"],
            "table": init_table_data(table),
            "rpsData": init_data(rps, lambda k, v: v.get(k, 0)),
            "qpsData": init_data(qps, lambda k, v: v.get(k, 0)),
            "responseData": init_data(response, lambda k, v: get_avg(v.get(k, [0])))
        }
        with open(file, 'w', encoding='u8') as f:
            f.write(template.render(context))
        return file


if __name__ == '__main__':
    pass
