import os
import random
import sys
import time
from pathlib import Path

from loguru import logger as llogger

logger = llogger.bind(module="FSBench", prefix="test")

b = 1
kb = 1024 * b
mb = 1024 * kb
rsize = [
    # 100 * mb, 90 * mb, 80 * mb, 70 * mb, 60 * mb, 50 * mb, 40 * mb, 30 * mb, 20 * mb, 10 * mb,  # 100mb -> 10mb
    5 * mb, 4.5 * mb, 4 * mb, 3.5 * mb, 3 * mb, 2.5 * mb, 2 * mb, 1.5 * mb, 1 * mb,  # 5mb -> 1mb
    500 * kb, 450 * kb, 400 * kb, 350 * kb, 300 * kb, 250 * kb, 200 * kb, 150 * kb, 100 * kb,  # 500kb -> 100kb
    # 50 * kb, 45 * kb, 40 * kb, 35 * kb, 30 * kb, 25 * kb, 20 * kb, 15 * kb, 10 * kb  # 50kb -> 10kb
]

min_float = sys.float_info.min * sys.float_info.epsilon

def _test(path: Path, size: int):
    size = int(size)
    path = path / ".tests" / "speed"
    logger.debug(f"Speed testing. Path: {path}, Size: {size / mb:.4f}MB")
    os.makedirs(path, exist_ok=True)
    _test_file = path / ".test_speed"
    if os.path.exists(_test_file):
        os.remove(_test_file)

    # Запись файла
    t1 = time.perf_counter()
    with open(_test_file, 'wb') as f:
        f.write(b'\0' * size)
    t2 = time.perf_counter()

    # Чтение файла
    t3 = time.perf_counter()
    with open(_test_file, 'rb') as f:
        while f.read(1024 * 1024):  # Читаем блоками по 1MB
            pass
    t4 = time.perf_counter()

    os.remove(_test_file)

    # Защита от деления на ноль
    write_time = max(t2 - t1, 1e-9)
    read_time = max(t4 - t3, 1e-9)

    return size / write_time / mb, size / read_time / mb  # MB/s upload, MB/s download

def test_speed(path: Path, iterations: int = 5):
    data_upload = []
    data_download = []
    for _ in range(iterations):
        u, d = _test(path, random.choice(rsize))
        data_upload.append(u)
        data_download.append(d)
    avg_uspeed = sum(data_upload) / len(data_upload)
    avg_dspeed = sum(data_download) / len(data_download)
    return avg_uspeed, avg_dspeed