#!/usr/bin/env python3
import os
import re
import time
import subprocess
import string
import logging
import argparse
import redis
import json
import tempfile
import psycopg2
from collections import namedtuple
from random import choices
from typing import Literal
import psutil

logger = logging.getLogger(__name__)


def performance_server_test_(cli_benchmark=False, cache=None, db=None,
                             redis_params={}, pg_params={}, folders=""):
    class PG:
        """Простой класс, мимикрирующий под поведение драйвера PG в CMF"""
        def __init__(self, pg_params={}):
            self.pg_params = pg_params
        def connect(self):
            conn = psycopg2.connect(**pg_params)
            return conn.cursor()

    ConditionType = Literal['greater', 'less']
    
    def _test_skipped(caption):
        return f'<p><span><b>{caption}</b>: </span><span>Skipped</span></p>'
    
    def _test_format(caption, actual_val, cutoff_val, fmt_func, cond: ConditionType):
        if cond == 'greater':
            success = actual_val >= cutoff_val
        else:
            success = actual_val <= cutoff_val
        color = 'green' if success else 'red'
        return f'<p><span><b>{caption}</b>: </span><span style="color:{color}">' \
               f'{fmt_func.format(x=actual_val)}</span></p>'
    
    # 20_000_000
    # AMD Ryzen 9 7950X + Pycharm debug = 0.89, 0.84, 0.9
    # bcrm = 1.3, 1.2, 1.3
    # 25_000_000
    # AMD Ryzen 9 7950X + Pycharm debug = 0.92, 1.02, 1.03
    # bcrm = 1.59, 1.7, 1.78

    tests = []
    folders = folders.split(':') if folders else ['/mnt/tmp', '/mnt/shared']

    Result = namedtuple(
        'Result',
        ['caption', 'actual_val', 'cutoff_val', 'fmt_func', 'cond'],
        defaults=(None, None, None, None)
    )

    reference_values = [
        # Test-1: Скорость цикла
        (2.0, '{x:.2f} sec.', 'less'),
        # Test-2: Производительность Redis (запись/чтение)
        (0.5, '{x:.2f} sec.', 'less'),
        # Test-3: Производительность PostgreSQL (запись/чтение)
        (1.0, '{x:.1f} sec.', 'less'),
        # Test-4: Объем памяти на сервере
        (8, '{x} GB.', 'greater'),
        # Test-5: Частота процессора
        (3, '{x:.1f} GHz', 'greater'),
        # Test-6: Частота процессора Turbo
        (3.8, '{x:.1f} GHz', 'greater'),
        # Test-7: Количество ядер'
        (8, '{x: d}', 'greater'),
        # Test-8: Производительность запись/чтение на диск
        (0.5, '{x:.1f} sec.', 'less'),
        # Test-9: Тест виртуализации (Redis cli)
        (0.03, '{x:.3f} msec.', 'less'),
        # Test-10: Latency до Redis
        (1.0, '{x:.1f} msec.', 'less'),
        # Test-11:  Latency до PostgreSQL
        (5.0, '{x:.1f} msec.', 'less'),
    ]

    if not cli_benchmark:
        logger.debug('Тест 1 cкорость цикла')
        # TEST 1
        cutoff_value, fmt_func, cond = reference_values[0]
        start = time.time()
        i = 0
        while i < 25_000_000:
            i += 1
        test1_time = time.time() - start
        test1 = Result(
            caption=f'Скорость цикла (не более {cutoff_value} sec.)', 
            actual_val=test1_time,
            cutoff_val=cutoff_value,
            fmt_func=fmt_func,
            cond=cond 
        )
        tests.append(test1)

    logger.debug('Тест 2 Redis чтение/запись')
    # TEST 2
    cutoff_value, fmt_func, cond = reference_values[1]
    key_count = 1000

    # redis_db = APP.REDIS_DB.redis
    if cache:
        redis_db = cache
    else:
        redis_db = redis.Redis(**redis_params, decode_responses=False)

    if db:
        db_conn = db
    else:
        db_conn = PG(pg_params)

    if not cli_benchmark:
        db_key = "_performance_server_test2_{}"
        new_captcha = ''.join(choices(string.digits, k=200))

        test2_start_time = time.time()
        try:
            # Падаем, если слишком долго
            for i in range(key_count):  # set
                if i % 100 == 0 and time.time() - test2_start_time > 60: 
                    raise RuntimeError('Тест выполнялся слишком долго')
                redis_db.set(db_key.format(i), new_captcha)
            for i in range(key_count):  # get
                if i % 100 == 0 and time.time() - test2_start_time > 60: 
                    raise RuntimeError('Тест выполнялся слишком долго')
                redis_db.get(db_key.format(i))
            for i in range(key_count):  # delete
                if i % 100 == 0 and time.time() - test2_start_time > 60: 
                    raise RuntimeError('Тест выполнялся слишком долго')
                redis_db.delete(db_key.format(i))
            test2_time = time.time() - test2_start_time
            test2 = Result(
                caption=f'Производительность Redis (запись/чтение, {key_count} записей)',
                actual_val=test2_time,
                cutoff_val=cutoff_value,
                fmt_func=fmt_func,
                cond=cond 
            )
        except Exception as e:
            test2 = Result(f'Производительность Redis (запись/чтение, {key_count} записей)')
            logger.debug(f'TEST 2 Exception: {e}')

        tests.append(test2)

    if not cli_benchmark:
        logger.debug('Тест 3 PostgreSQL чтение/запись')
        # TEST 3
        cutoff_value, fmt_func, cond = reference_values[2]
        row_count = 1000
        test_table_name = 'public.performance_server_test3'

        begin = 'BEGIN'
        commit = 'COMMIT'

        s1 = f'DROP TABLE IF EXISTS {test_table_name};'

        s2 = f'CREATE TABLE {test_table_name} ('
        s2 += 'id integer, '
        s2 += ', '.join(f'name{i} integer' for i in range(15))
        s2 += ', '
        s2 += ', '.join(f'name_s{i} varchar' for i in range(15))
        s2 += ');'

        s3 = f'INSERT INTO {test_table_name} VALUES ('
        s3 += '{id}, '
        s3 += ', '.join(f'{i}' for i in range(15))
        s3 += ', '
        s3 += ', '.join([f"'{''.join(choices(string.ascii_letters, k=128))}'" for _ in range(15)])
        s3 += ');'
        print([f"'{''.join(choices(string.ascii_letters, k=128))}'" for _ in range(15)])

        s4 = f'DELETE FROM {test_table_name} WHERE id = {{id}};'

        with db_conn.connect() as db_connect:
            db_connect.execute(s1)
            db_connect.execute(s2)

            test3_start_time = time.time()
            try:
                # Падаем, если слишком долго
                for i in range(row_count):  # insert, commit
                    if i % 100 == 0 and time.time() - test3_start_time > 60:
                        raise RuntimeError('Тест выполнялся слишком долго')
                    db_connect.execute(begin)
                    db_connect.execute(s3.format(id=i))
                    db_connect.execute(commit)

                for i in range(row_count):  # delete, commit
                    if i % 100 == 0 and time.time() - test3_start_time > 60:
                        raise RuntimeError('Тест выполнялся слишком долго')
                    db_connect.execute(begin)
                    db_connect.execute(s4.format(id=i))
                    db_connect.execute(commit)

                test3_time = time.time() - test3_start_time

                test3 = Result(
                    caption=f'Производительность PostgreSQL (запись/чтение, {row_count} записей)',
                    actual_val=test3_time,
                    cutoff_val=cutoff_value,
                    fmt_func=fmt_func,
                    cond=cond
                )
            except Exception as err:
                test3 = Result(f'Производительность PostgreSQL (запись/чтение, {key_count} записей)')
                logger.debug(f'TEST 2 Exception: {e}')

            db_connect.execute(s1)
        tests.append(test3)

    if not cli_benchmark:
        logger.debug('Тест 4 Объём памяти')
        # TEST 4
        cutoff_value, fmt_func, cond = reference_values[3]
        test4_val = int(psutil.virtual_memory()[0]/1024/1024/1024)
        test4 = Result(
            caption=f'Объем памяти на сервере (не менее {cutoff_value} Гб.)',
            actual_val=test4_val,
            cutoff_val=cutoff_value,
            fmt_func=fmt_func,
            cond=cond
        )
        tests.append(test4)

    if not cli_benchmark:
        logger.debug('Тест 5 Частота процессора')
        # TEST 5
        cutoff_value, fmt_func, cond = reference_values[4]
        test5_val = psutil.cpu_freq().max / 1000
        test5 = Result(
            caption=f'Частота процессора (не менее  {cutoff_value} GHz)',
            actual_val=test5_val,
            cutoff_val=cutoff_value,
            fmt_func=fmt_func,
            cond=cond
        )
        tests.append(test5)

    if not cli_benchmark:
        logger.debug('Тест 6 Частота процессора Turbo')
        # TEST 6
        cutoff_value, fmt_func, cond = reference_values[5]
        cpufreq_file = '/sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_max_freq'
        if not os.path.exists(cpufreq_file):
            test6 = Result(f'Частота процессора (не менее {cutoff_value} GHz)')
        else:
            with open(cpufreq_file, 'r') as f:
                test6_val = int(f.read()) / 1000 / 1000
            test6 = Result(
                caption=f'Частота процессора Turbo (не менее {cutoff_value} GHz)',
                actual_val=test6_val,
                cutoff_val=cutoff_value,
                fmt_func=fmt_func,
                cond=cond
            )
        tests.append(test6)

    if not cli_benchmark:
        logger.debug('Тест 7 Количество ядер')
        # TEST 7
        cutoff_value, fmt_func, cond = reference_values[6]
        #??? psutil.cpu_count(logical=True)
        test7_val = psutil.cpu_count()
        test7 = Result(
            caption=f'Количество ядер (не менее {cutoff_value})',
            actual_val=test7_val,
            cutoff_val=cutoff_value,
            fmt_func=fmt_func,
            cond=cond
        )
        tests.append(test7)

    logger.debug('Тест 8 redis cli')
    # TEST 8
    cutoff_value, fmt_func, cond = reference_values[8]
    test8_res = subprocess.run(['redis-cli', '--intrinsic-latency',  '5'], check=True, timeout=20, capture_output=True, text=True)
    test8_stdout = test8_res.stdout.strip()
    test8_res = re.findall(r'avg latency: ([\d.]* \w*)', test8_stdout)
    if test8_res:
        test8_res = test8_res[0]
    else:
        test8_res = None

    test8_val = None
    if test8_res:
        float_str = "".join(char for char in test8_res if char.isdigit() or char =='.')
        test8_val = float(float_str)
    
    if test8_val:
        test8 = Result(
            caption=f'Тест виртуализации Redis cli (не более {cutoff_value} msec.)',
            actual_val=test8_val,
            cutoff_val=cutoff_value,
            fmt_func=fmt_func,
            cond=cond 
        )
    else:
        test8 = Result('Тест виртуализации (Redis cli)')
    tests.append(test8)

    # TEST 9
    cutoff_value, fmt_func, cond = reference_values[7]
    file_count = 1000
    s_8kb = b'x' * 8 * (10**3)
    for folder in folders:
        logger.debug('Тест 9 Чтение/запись %s', folder)
        with tempfile.TemporaryDirectory(dir=folder) as tmpdir:
            test9_start_time = time.time()
            try:
                for i in range(file_count):
                    f = os.open(os.path.join(tmpdir, str(i)), os.O_RDWR | os.O_CREAT | os.O_SYNC)
                    os.write(f, s_8kb)
                    os.close(f)
            except Exception as err:
                raise
            test9_time = time.time() - test9_start_time
            
        test9 = Result(
            caption=f'Производительность запись/чтение на диск, {file_count} файлов, {folder} (не более {cutoff_value} sec.)',
            actual_val=test9_time,
            cutoff_val=cutoff_value,
            fmt_func=fmt_func,
            cond=cond
        )
        tests.append(test9)

    if not cli_benchmark:
        logger.debug('Тест 10 Redis Latency')
        # TEST 10
        cutoff_value, fmt_func, cond = reference_values[9]
        try:
            test10_res = subprocess.run([
                                            'redis-cli', '--latency',
                                            '-h', redis_params['host'],
                                            '-p', f"{redis_params['port']}",
                                            '-i', '5', '--raw'
                                        ],
                                        check=True, timeout=20, capture_output=True, text=True)
            # В stdout будет что-то вроде этого: "0 1 0.62 469"
            test10_stdout = test10_res.stdout.strip()
            test10_min, test10_max, test10_avg, *_ = map(int, map(float, test10_stdout.split()))
            test10 = Result(
                caption=f'Latency до Redis (не более {cutoff_value} мс) min {test10_min}мс. / max {test10_max}мс. / avg {test10_avg}мс.',
                actual_val=test10_avg,
                cutoff_val=cutoff_value,
                fmt_func=fmt_func,
                cond=cond
            )
        except Exception as e:
            # Упадем если redis локальный или необычные настройки
            test10 = Result(f'Latency до Redis (не более 1мс)')
            logger.debug(f'TEST 10 Exception: {e}')
        tests.append(test10)

    if not cli_benchmark:
        logger.debug('Тест 11 PostgreSQL Latency')
        # TEST 11
        cutoff_value, fmt_func, cond = reference_values[10]
        try_count = 1000
        with db_conn.connect() as db_connect:
            timings = []
            for i in range(try_count):
                st = time.time()
                db_connect.execute('select 1;')
                timings.append(time.time()-st)
        test11_min = min(timings) * 1000
        test11_max = max(timings) * 1000
        test11_avg = (sum(timings) / len(timings)) * 1000
        test11_98 = sorted(timings)[int(try_count/100*98)] * 1000
        test11 = Result(
            caption=f'Latency до PostgreSQL (не более {cutoff_value} мс): min {test11_min:.3f}мс. / max {test11_max:.3f}мс. / avg {test11_avg:.3f}мс. / 98p {test11_98:.3f}мс.',
            actual_val=test11_avg,
            cutoff_val=cutoff_value,
            fmt_func=fmt_func,
            cond=cond 
        )
        tests.append(test11)

    logger.debug('Tests DONE')

    if cli_benchmark:
        return tests

    results = []
    for test in tests:
        if test.actual_val:
            results.append(
                _test_format(
                    caption=test.caption,
                    actual_val=test.actual_val,
                    cutoff_val=test.cutoff_val,
                    fmt_func=test.fmt_func,
                    cond=test.cond,
                )
            )
        else:
            results.append(_test_skipped(test.caption))

    return ''.join(results)


def print_item(item):
    """Выведет значения в терминал."""
    if isinstance(item, dict):
        caption = item['title']
        actual_val = item['value']
        cutoff_val = item['reference']
        fmt_func = item['format']
        cond = item['condition']
    else:
        (
            caption,     # Строка с описанием теста
            actual_val,  # Значение или None
            cutoff_val,  # максимальное приемлемое значение
            fmt_func,    # функция для форматирования значения (добавить unit, округлить и т.д.)
            cond,        # 'greater' - больше = лучше: 'less' - меньше = лучше
        ) = item

    if actual_val is None:
        print(f"{caption}: Skipped")
        return

    val = fmt_func(actual_val) if callable(fmt_func) else fmt_func.format(x=actual_val)
    text = f"{caption}: {val}"

    print(text)


def configure_logging(debug=False):
    level = logging.DEBUG if debug else logging.INFO
    logging.basicConfig(
        level=level,
        format="%(asctime)s - [%(levelname)s] - %(module)s.%(funcName)s - %(message)s",
    )


def parse_args():
    p = argparse.ArgumentParser("Redis benchmark")
    p.add_argument("--json", action="store_true", default=False,
                   help='Вывод результатов в json')
    p.add_argument("--debug", action="store_true", default=False,
                   help='Вывод дебага (при --json всегда отключено)')
    p.add_argument("--folders", type=str, default="",
                   help="Директории для теста дисков")
    p.add_argument("--all", action="store_true",
                   help="Запустить все web бенчмарки, даже если они дублируются отдельными скриптами")
    p.add_argument("--r-db", type=int, default=7, help="Redis, Номер БД, по-умолчанию 7")
    p.add_argument("--r-host", type=str, help="Redis, хост (ip, домен), имеет приоритет над сокетом")
    p.add_argument("--r-port", type=int, default=6379, help='Порт, по-умолчанию 6379')
    p.add_argument("--r-unix-socket-path", type=str, default="/var/run/redis/redis-server.sock",
                   help="Redis, путь к сокету при использовании такого типа подключения")
    p.add_argument("--r-password", type=str, default=None,
                   help='Redis, пароль для подключения к Redis, необязательный параметр')
    p.add_argument("--r-username", type=str, default=None,
                   help='Redis, пользователь для подключения к Redis, необязательный параметр')
    p.add_argument("--r-ssl", action="store_true", default=False,
                   help='Redis, использовать SSL/TLS для подключения')
    p.add_argument("--r-ssl-ca-certs", type=str, default="/opt/eva-app/custom/redis-ca.crt",
                   help="Redis, путь к сертификату CA")
    p.add_argument("--r-ssl-certfile", type=str, default="/opt/eva-app/custom/redis.crt",
                   help="Redis, путь к сертификату клиента")
    p.add_argument("--r-ssl-keyfile", type=str, default="/opt/eva-app/custom/redis.key",
                   help="Redis, путь к ключи сертификата клиента")
    p.add_argument("--r-type", type=str, default="redis",
                   help=argparse.SUPPRESS)
                   # help='Тип кеша Евы, пока поддерживается только Redis'
    p.add_argument('-d', '--p-db', type=str, default='evadb',
                   help='PSQL, имя БД для подключения (default: evadb)')
    p.add_argument('-u', '--p-user', type=str, default='postgres',
                   help='PSQL, имя пользователя БД для подключения (default: postgres)')
    p.add_argument('-p', '--p-password', type=str, default=None,
                   help='PSQL, пароль пользователя БД для подключения')
    p.add_argument('-H', '--p-host', type=str, default=None,
                   help='PSQL, адрес сервера БД (default: localhost)')
    p.add_argument('-P', '--p-port', type=int, default=5432,
                   help='PSQL, порт сервера БД для подключения (default: 5432)')
    return p.parse_args()


if __name__ == "__main__":
    args = parse_args()

    # Параметры Redis
    if args.r_host:
        redis_params = {"host": args.r_host, "port": args.r_port, "db": args.r_db,
                        "password": args.r_password, "username": args.r_username}
        if args.r_ssl:
            redis_params['ssl_ca_certs'] = args.r_ssl_ca_certs
            redis_params['ssl_certfile'] = args.r_ssl_certfile
            redis_params['ssl_keyfile'] = args.r_ssl_keyfile
    else:
        redis_params = {"unix_socket_path": args.r_unix_socket_path,
                        "password": args.r_password, "db": args.r_db}

    # Параметры PostgreSQL
    pg_params = {
        'dbname': args.p_db,
        'user': args.p_user,
    }
    if args.p_host:
        pg_params['password'] = args.p_password
        pg_params['host'] = args.p_host
        pg_params['port'] = args.p_port

    # Настройка логов; только если не json
    if not args.json:
        configure_logging(debug=args.debug)

    # Запуск бенчмарка
    res = performance_server_test_(cli_benchmark=not args.all, folders=args.folders,
                                   redis_params=redis_params, pg_params=pg_params)

    # Вывод результатов
    if args.json:
        print(json.dumps(res, ensure_ascii=False))
    else:
        for r in res:
            print_item(r)