#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import argparse
import json
import shutil
import re
import subprocess
import logging
from pathlib import Path


logger = logging.getLogger(__name__)

TITLE_TARGET = ''


class TestRunner():
    def __init__(self, check_path='/mnt/shared', count=0, keep_testfile=False, hdd=False):
        if os.path.isdir(check_path):
            self.test_path = Path(check_path)
            self.dd_of = self.test_path / "eva_disk_benchmark_of.tmp"
        else:
            self.dd_of = Path(check_path)
            self.test_path = self.dd_of.parent

        self.dd_if = self.test_path / "eva_disk_benchmark_if.tmp"

        self.count = count
        self.keep_testfile = keep_testfile
        self.hdd = hdd
        self.metrics = []
        self.results = []

    def run(self, cmd, *, capture=False):
        logger.debug(' '.join(cmd))
        return subprocess.run(
            cmd,
            check=True,
            text=True,
            stdout=subprocess.PIPE if capture else subprocess.DEVNULL,
            stderr=subprocess.PIPE if capture else subprocess.DEVNULL,
        )

    def check_space(self):
        disk_usage = shutil.disk_usage(self.test_path)
        # 4G input + some output + 10G reserve, в гигабайтах
        required = (4 + self.count/1024 + 10)
        if disk_usage.free <= required * 1024**3:
            raise RuntimeError(f'В директории {self.test_path} доступно менее {required}Гб')

    def prepare_files(self):
        logger.info("Подготовка исходного файла")
        count = 4096 if self.count > 4096 else self.count
        self.run([
            "dd",
            "if=/dev/urandom",
            f"of={self.dd_if}",
            f"count={count}",
            "bs=1M",
        ])
        logger.info("Подгружаем файл в кеш")
        self.run(["cat", str(self.dd_if)], capture=False)
        logger.debug("Повторнй cat")
        self.run(["cat", str(self.dd_if)], capture=False)

    def run_dd(self, count: int, bs: str, flags: str):
        args = [
            "dd",
            f"if={self.dd_if}",
            f"of={self.dd_of}",
            f"count={count}",
            f"bs={bs}",
        ]
        if flags:
            args.append(f"oflag={flags}")

        logger.info("DD %s count %s size %s %s",
                    self.dd_of.parent, count, bs, flags or "noflags BUFFERED")

        res = self.run(args, capture=True)

        dd_re = re.compile(
            r"(?P<bytes>\d+)\s+bytes.*copied,\s*"
            r"(?P<seconds>[\d.]+)\s*s,\s*"
            r"(?P<speed>[\d.]+)\s*(?P<unit>[MG]B)/s"
        )

        match = dd_re.search(res.stderr)
        if not match:
            raise RuntimeError(f"Unable to parse dd output:\n{res.stderr}")

        speed = float(match.group("speed"))
        if match.group("unit") == "GB":
            speed *= 1024

        result = {
            "count": count,
            "bs": bs,
            "flags": flags,
            "bytes": int(match.group("bytes")),
            "seconds": float(match.group("seconds")),
            "mbps": speed,
        }
        result["human"] = \
            f"count={result['count']} bs=1M bytes={result['bytes']}"\
            f" time={result['seconds']:.3f}s speed={result['mbps']} MB/s"
        return result

    def get_metrics(self):
        if not self.results:
            return
        for item in self.results:
            count = item['count']
            if count >= 1024:
                count = f"{count / 1024:.0f}GB"
            else:
                count = f"{count:.0f}MB"
            # ref = 110 if self.hdd else 280
            ref = 70
            title = f"Последовательная запись {count}, {item['flags']}"
            title = f'{title}, {self.test_path}' if self.test_path else title
            title += f' (не менее {ref:.0f} MB/sec)'
            self.metrics.append({
                'title': title,
                'value': item['mbps'],
                'reference': ref,  # cutoff value
                'format': '{x:.2f} MB/sec',
                'condition': 'greater',  # condition type
            })

    def run_test(self):

        counts = []
        if self.count:
            counts.append(self.count)
        else:
            counts += [512, 1024, 4096]

        # Нужно очистить до проверки места, могло остаться после прошлого теста.
        self.dd_if.unlink(missing_ok=True)
        self.dd_of.unlink(missing_ok=True)

        self.check_space()
        self.prepare_files()

        try:
            for count in counts:
                result = self.run_dd(
                    count=count,
                    bs="1M",
                    flags="direct",
                )
                self.results.append(result)
                self.get_metrics()
        finally:
            self.dd_if.unlink(missing_ok=True)
            if not self.keep_testfile:
                self.dd_of.unlink(missing_ok=True)


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():
    parser = argparse.ArgumentParser(
        description="""Тест стореджа с помощью dd.
 WARNING Тест временно займёт 8Гб.
 На разделе должно быть свободно по крайней мере 18Гб.""",
    )
    parser.add_argument("--path", type=str, default="/mnt/shared",
                        help="Путь к стореджу для проверки")
    parser.add_argument("--json", action="store_true",
                        help="Напечатать вывод в json")
    parser.add_argument("--count", type=int,
                        help="Размер файла для теста, в мегабайтах")
    parser.add_argument("-k", "--keep", action="store_true",
                        help="Оставить тестовый файл, не удалять")
    parser.add_argument("--hdd", action="store_true",
                        help=argparse.SUPPRESS)
                        # help="Тест проводится на hdd - флаг влияет на референсы (по-умолчанию, считаем, что ssd)")
    parser.add_argument("--debug", action="store_true", default=False,
                        help='Вывод дебага (при --json всегда отключено)')

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    if not args.json:
        configure_logging(debug=args.debug)

    if args.json:
        TITLE_TARGET = args.path

    runner = TestRunner(check_path=args.path, count=args.count,
                        keep_testfile=args.keep, hdd=args.hdd)
    runner.run_test()

    if args.json:
        print(json.dumps(runner.metrics, ensure_ascii=False))
    else:
        for result in runner.results:
            logger.info(result['human'])