#!/usr/bin/env python3
# coding: utf-8
"""Скрипт для запуска тестов в консоли."""
import os
import sys
import subprocess
import argparse
import logging
import json
from urllib.parse import urlparse
from pathlib import Path

BM_DIR = Path(__file__).resolve().parent
BM_CONFIG = BM_DIR / 'config.py'

for func in (os.chdir, sys.path.append):
    _ = func('/opt/eva-app') if Path('/opt/eva-app').exists() else None

os.chdir(BM_DIR)
logger = logging.getLogger(__name__)


def benchmark_(color=True, config=None, folders="", timeout=0, dry=False):
    def admin_web_test():
        """Выполнит тест сервера из веб-интерфейса"""
        return []
        # # init_views_and_ds = True if ds_registry else False
        # with cmf_context(init_views_and_ds=True):
        #     admin = models.CmfPerson.get(
        #         id='CmfPerson:00000000-0000-0000-0000-000000000001',
        #         fields=['is_admin', 'is_support'])
        #     cmf.app.set_current_person(admin)
        #     g.acl_admin_mode = True
        #     g.is_admin = True
        #     res = models.CmfGlobalSettings.performance_server_test(cli_benchmark=True)
        # return res

    def colorize(value, condition_type, cutoff, text):
        """Покрасит результат теста в красный или зелёный цвет.
        Зависит от референсного значения."""
        # ANSI colors
        reset = "\033[0m"
        green = "\033[32m"
        red = "\033[31m"

        color = reset

        if condition_type == "less":
            if value <= cutoff:
                color = green
            else:
                color = red

        elif condition_type == "greater":
            if value >= cutoff:
                color = green
            else:
                color = red

        return f"{color}{text}{reset}"

    def print_item(item):
        """Выведет значения в терминал."""
        do_colorize = True  # для tuple, старая схема
        if isinstance(item, dict):
            # mandatory
            caption = item['title']
            actual_val = item['value']
            cutoff_val = item['reference']
            fmt_func = item['format']
            cond = item['condition']
            # optional
            do_colorize = item.get('colorize', True)
        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)
        if do_colorize:
            val = colorize(actual_val, cond, cutoff_val, val)
        text = f"{caption}: {val}"

        print(text)

    def get_redis_params():
        """Вернёт параметры подключения к redis в нужном формате"""
        params = config.cache_settings['default'].copy()
        params['params'] = []
        params['web_bm_params'] = []
        for k, v in config.cache_settings['default'].items():
            params['params'].append(f"--{k.replace('_','-')}={v}")
            params['web_bm_params'].append(f"--r-{k.replace('_','-')}={v}")
        return params

    def get_pg_params():
        # TODO: получить из параметров подключения
        parsed = urlparse(config.data_sources['default']['sqlalchemy.url'])
        pg_username = parsed.username or "postgres"
        pg_password = parsed.password or ""
        pg_host = parsed.hostname or ""
        pg_port = parsed.port or 5432
        pg_dbname = parsed.path.lstrip('/').split('/', 1)[0] or "evadb"
        pg_web_params = [
            f'--p-host={pg_host}', f'--p-port={pg_port}',
            f'--p-user={pg_username}', f'--p-password={pg_password}',
            f'--p-db={pg_dbname}'
        ]
        return [pg_host, pg_port, pg_username, pg_password, pg_dbname, pg_web_params]

    # def is_rotational(path):
    #     """
    #     Пытаемся понять, что за девайс (hdd, ssd), т.к. референсы будут разные.
    #     """
    #     try:
    #         # Получаем девайс, на котором живёт директория
    #         result = subprocess.run(
    #             ["findmnt", "-no", "SOURCE", "--target", path],
    #             stdout=subprocess.PIPE,
    #             stderr=subprocess.PIPE,
    #             text=True,
    #             check=True
    #         )

    #         source = result.stdout.strip()
    #         logger.debug('path %s source %s', path, source)

    #         if not source.startswith("/dev/"):
    #             msg = f"Не получилось определить тип устройства для {source}, возможно сетевое хранилище"
    #             logger.debug(msg)
    #             raise RuntimeError(msg)

    #         device = os.path.basename(source)

    #         # Нужно получить имя девайса, без партиций, чтобы обратиться к /sys/block/<dev>
    #         if device.startswith("nvme"):
    #             # dev = nvme0n1; partition = nvme0n1p1
    #             device = device.split("p")[0]
    #         elif device.startswith("md"):
    #             # dev, partition = md1, md150
    #             device = device
    #         else:
    #             # dev = sda, vda; partition = sda1, vda3
    #             device = device.rstrip("0123456789")

    #         rotational_path = f"/sys/block/{device}/queue/rotational"

    #         if not os.path.exists(rotational_path):
    #             msg = f"Не получилось определить тип устройства {device}"
    #             logger.debug(msg)
    #             raise RuntimeError(msg)

    #         with open(rotational_path) as f:
    #             return f.read().strip() == "1"

    #     except RuntimeError:
    #         return False

    def run_test(test, args, timeout=60):
        """Запустит тест с указанными параметрами и лимитом по времени"""
        cmd = [str(test)] + args
        logging.debug(' '.join(cmd))
        out = subprocess.run(cmd, capture_output=True,
                             check=True, text=True, timeout=timeout)
        return out.stdout

    # BM_DIR = Path("/opt/fox_acrm/tests")

    # logger = logging.getLogger()
    # logger.setLevel('INFO')

    redis_params = get_redis_params()
    redis_host = redis_params.get('host')
    redis_port = redis_params.get('port')
    logger.debug("%s %s", f"{redis_host=}", f"{redis_port=}")
    pg_host, pg_port, pg_username, pg_password, pg_dbname, pg_web_params = get_pg_params()
    pg_table = 'eva_performance_test'
    logger.debug("%s %s", f"{pg_host=}", f"{pg_port=}")

    web_bm_params = redis_params['web_bm_params'] + pg_web_params
    logger.debug("%s", f"{web_bm_params=}")

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

    tests = {
        'cpu_count': {
            'file': BM_DIR / 'cpu_count.py',
            'args': ['--json']
        },
        'cpu_cycle': {
            'file': BM_DIR / 'cpu_cycle.py',
            'args': ['--json']
        },
        'cpu_max_frequency': {
            'file': BM_DIR / 'cpu_max_frequency.py',
            'args': ['--json']
        },
        'cpu_temperature': {
            'file': BM_DIR / 'cpu_temperature.py',
            'args': ['--json']
        },
        'cpu_throttle': {
            'file': BM_DIR / 'cpu_throttle.py',
            'args': ['--json']
        },
        'ram_size': {
            'file': BM_DIR / 'ram_size.py',
            'args': ['--json']
        },
        'db_postgres': {
            'file': BM_DIR / 'db_postgres.py',
            'args': ['--json', f'--host={pg_host}', f'--port={pg_port}',
                     f'--username={pg_username}', f'--password={pg_password}',
                     f'--dbname={pg_dbname}', f'--table={pg_table}']
        },
        'db_redis_throughput': {
            'file': BM_DIR / 'db_redis_throughput.py',
            'args': ['--json'] + redis_params['params']
        },
        # TODO: не обсудили референсы для теста, и нужен ли он вообще для задачи. Обсудить потом.
        # 'web_benchmark': {
        #     'file': BM_DIR / 'web_benchmark.py',
        #     'args': ['--json'] + web_bm_params + [f"--folders={':'.join(folders)}"],
        #     'timeout': 120  # долгий тест
        # },
    }

    for folder in folders:
        ## dd sequential, запись больших файлов, в документ загрузили видео, в задачу приложили pdf
        args =  ['--json', f'--path={folder}', '--count=1000']
        # лишнее, т.к. должно быть ssd
        # if is_rotational(folder):
        #     fio_params.append('--hdd')
        tests[f'disk_dd.py ({folder})'] = {
            'file': BM_DIR / 'disk_dd.py',
            'args': args,
            'timeout': 120,
        }
        ## fio 4k random r/w, обычная работа с диском, запись в БД, мелкие вложения и тд,
        ## понемногу пишем логи и тд
        args = ['--json', f'--target={folder}', '--loops=1', '--numjobs=1',
                      '--size=5120', '--runtime=20']
        # лишнее, т.к. должно быть ssd
        # if is_rotational(folder):
        #     fio_params.append('--hdd')
        tests[f'disk_fio.py ({folder})'] = {
            'file': BM_DIR / 'disk_fio.py',
            'args': args + ['--prepare=skip'],
            'prepare': [str(BM_DIR / 'disk_fio.py')] + args + ['--prepare=only']
        }
        ## Дозапись в файл, т.е. логи, проверяем что Еве не подсунули s3 для всего на свете
        args = ['--json', f'--path={folder}', '--duration=10', '--threads=2']
        tests[f'disk_append.py ({folder})'] = {
            'file': BM_DIR / 'disk_append.py',
            'args': args + ['--prepare=skip'],
            'prepare': [str(BM_DIR / 'disk_append.py')] + args + ['--prepare=only']
        }

    hosts = { 'PostgreSQL': pg_host, 'Redis': redis_host }
    for alias, host in hosts.items():
        tests[f'network_flood_ping ({alias})'] = {
            'file': BM_DIR / 'network_flood_ping.py',
            'args': ['--json', f'--host={host}', f'--alias={alias}']
        }
        tests[f'network_latency_ping ({alias})'] = {
            'file': BM_DIR / 'network_ping.py',
            'args': ['--json', f'--host={host}', f'--alias={alias}']
        }

    admin_web_results = admin_web_test()
    for item in admin_web_results:
        print_item(item)

    for fname, cfg in tests.items():
        if dry:
            if cfg.get('prepare'):
                cmd = cfg['prepare']
                logger.info('%s prepare: %s', fname, ' '.join([str(i) for i in cfg['prepare']]))
            logger.info('%s cmd: %s', fname, ' '.join([str(i) for i in [cfg['file']] + cfg['args']]))
            continue
        try:
            # По приоритету: 1) пользователь указал вручную;
            #                2) указано в параметрах теста (tests['test']['timeout'])
            #                3) стандартный дефолт
            timeout = timeout if timeout else cfg.get('timeout') if cfg.get('timeout') else 60
            # Это для fio, ему нужно подготовить файлы, может занять какое-то время,
            # оно может растянуть тест дольше дефолтного таймаута
            if cfg.get('prepare'):
                cmd = cfg['prepare']
                logging.debug('prepare: %s', ' '.join(cmd))
                subprocess.run(cmd, check=True, timeout=timeout, capture_output=True)
                logging.debug('Выполнили %s prepare', fname)
            metrics = run_test(cfg['file'], cfg['args'])
            metrics = json.loads(metrics)

            for item in metrics:
                try:
                    print_item(item)
                except Exception as e:
                    logging.debug("Malformed result in %s: %s, %s", fname, str(item), e)
                    print(f"{fname}: Skipped")

        except subprocess.TimeoutExpired as e:
            logging.debug("Бенчмарк %s не уложился отведённое время %d секунд", fname, 60)
            print(f"{fname}: Timeout, skipping")
        except Exception as e:
            logging.debug("Benchmark failed: %s, %s", fname, e)
            print(f"{fname}: Skipped")


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("--dry", action="store_true", default=False,
                   help='Не запускать тесты. Полезно, чтобы посмотреть команды запуска.')
    p.add_argument("--folders", type=str, default="",
                   help="Директории для теста дисков")
    p.add_argument("--timeout", type=int, default=0,
                   help="Максимальное время выполнения каждого теста")
    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-socket", 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-user", 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, адрес сервера БД')
    p.add_argument('-P', '--p-port', type=int, default=None,
                   help='PSQL, порт сервера БД для подключения')
    return p.parse_args()


def config_from_args(args):
    class Config():
        """Простой класс, мимикрирующий под config.py cmf/eva"""
        def __init__(self, args):
            # Redis params
            cs = {}
            if args.r_host:
                cs = {"host": args.r_host, "port": args.r_port, "db": args.r_db}
                if args.r_ssl:
                    cs['ssl_ca_certs'] = args.r_ssl_ca_certs
                    cs['ssl_certfile'] = args.r_ssl_certfile
                    cs['ssl_keyfile'] = args.r_ssl_keyfile
            else:
                cs = {"unix_socket_path": args.r_socket, "db": args.r_db}

            if args.r_password:
                cs['password'] = args.r_password
            if args.r_user:
                cs['user'] = args.r_user
            self.cache_settings = {"default": cs}

            # PostgreSQL params

            ds = {'sqlalchemy.url': ''}
            ds['sqlalchemy.url'] = f"postgresql+psycopg2://" \
                f"{args.p_user}{':' + args.p_password if args.p_password else ''}" \
                f"@{args.p_host if args.p_host else ''}{':' + str(args.p_port) if args.p_port else ''}/{args.p_db}"
            self.data_sources = {"default": ds}
            self.DEBUG = args.debug

    return Config(args)


if __name__ == "__main__":
    args = parse_args()
    # Если есть config.py - берём параметры из него
    # Если запустили на Еве, берём параметры из конфигов Евы
    # Иначе можем взять из аргументов
    if BM_CONFIG.exists():
        import config
    elif os.path.exists('/opt/eva-app/custom/config.py'):
        import eva_config_load as config
    else:
        config = config_from_args(args)

    # Logging
    debug = args.debug if args.debug else getattr(config, 'DEBUG', False)
    configure_logging(debug=debug)

    benchmark_(config=config, folders=args.folders, timeout=args.timeout, dry=args.dry)