# test_runner.py
import os
import sys
import time
from pathlib import Path
from datetime import datetime
import numpy as np
import librosa
from pydub import AudioSegment
from config import PRESETS, DEFAULT_CONFIG
from encoder import audio_to_image
from decoder import image_to_audio
from utils import metadata_to_json, file_size_kb, file_size_mb, clean_temp_files, ensure_dir


# ==================== ТЕСТОВЫЕ КОНФИГУРАЦИИ ====================
# Все используют одинаковый active_preset: 75p_n4096
TEST_CONFIGS = {
    # PNG (максимальное сжатие)
    "png_max": {
        "output_format": "png",
        "output_quality": 9,
        "output_lossless": True,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    
    # JPEG
    "jpeg_q100": {
        "output_format": "jpeg",
        "output_quality": 100,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    "jpeg_q75": {
        "output_format": "jpeg",
        "output_quality": 75,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    "jpeg_q50": {
        "output_format": "jpeg",
        "output_quality": 50,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    "jpeg_q25": {
        "output_format": "jpeg",
        "output_quality": 25,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    "jpeg_q5": {
        "output_format": "jpeg",
        "output_quality": 5,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    
    # WebP
    "webp_q100": {
        "output_format": "webp",
        "output_quality": 100,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    "webp_q75": {
        "output_format": "webp",
        "output_quality": 75,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    "webp_q50": {
        "output_format": "webp",
        "output_quality": 50,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    "webp_q25": {
        "output_format": "webp",
        "output_quality": 25,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    "webp_q5": {
        "output_format": "webp",
        "output_quality": 5,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    "webp_lossless": {
        "output_format": "webp",
        "output_quality": 100,
        "output_lossless": True,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    
    # AVIF
    "avif_q100": {
        "output_format": "avif",
        "output_quality": 100,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    "avif_q75": {
        "output_format": "avif",
        "output_quality": 75,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    "avif_q50": {
        "output_format": "avif",
        "output_quality": 50,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    "avif_q25": {
        "output_format": "avif",
        "output_quality": 25,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    "avif_q5": {
        "output_format": "avif",
        "output_quality": 5,
        "output_lossless": False,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
    "avif_lossless": {
        "output_format": "avif",
        "output_quality": 100,
        "output_lossless": True,
        "scale_enabled": False,
        "scale_x": 1,
        "scale_y": 1,
        "scale_method": "lanczos"
    },
}


class TeeOutput:
    """Перенаправляет вывод одновременно в консоль и файл."""
    def __init__(self, file_path: str):
        self.file = open(file_path, 'w', encoding='utf-8')
        self.stdout = sys.stdout
        
    def write(self, message):
        self.stdout.write(message)
        self.file.write(message)
        
    def flush(self):
        self.stdout.flush()
        self.file.flush()
        
    def close(self):
        self.file.close()


def mp3_to_wav(mp3_path: str, wav_path: str) -> str:
    """Конвертирует MP3 в WAV."""
    print(f"   Конвертация MP3: {Path(mp3_path).name}")
    audio = AudioSegment.from_mp3(mp3_path)
    audio.export(wav_path, format="wav")
    size_mb = Path(wav_path).stat().st_size / (1024 * 1024)
    print(f"   Готово ({size_mb:.1f} MB)")
    return wav_path


def calculate_snr(original: np.ndarray, recovered: np.ndarray) -> float:
    """Вычисляет SNR между оригиналом и восстановленным сигналом."""
    min_len = min(original.shape[1] if original.ndim > 1 else len(original),
                  recovered.shape[1] if recovered.ndim > 1 else len(recovered))
    
    if original.ndim == 1:
        orig = original[:min_len]
    else:
        orig = original[0, :min_len]
        
    if recovered.ndim == 1:
        rec = recovered[:min_len]
    else:
        rec = recovered[0, :min_len]
    
    if np.max(np.abs(rec)) > 0:
        rec = rec / np.max(np.abs(rec)) * np.max(np.abs(orig))
    
    noise = orig - rec
    signal_power = np.mean(orig ** 2)
    noise_power = np.mean(noise ** 2)
    
    if noise_power == 0:
        return float('inf')
    
    snr = 10 * np.log10(signal_power / noise_power)
    return float(snr)


def run_single_test(config_name: str, test_config: dict, base_config: dict, 
                    mp3_file: str, data_dir: str) -> dict:
    """Запускает один тест и возвращает результаты."""
    print(f"\n{'='*70}")
    print(f"ТЕСТ: {config_name}")
    print(f"{'='*70}")
    
    # Объединяем с базовыми настройками и фиксируем пресет
    full_config = {**base_config, **test_config}
    full_config["active_preset"] = "75p_n4096"
    
    # Создаем отдельную папку для теста: data/[config_name]/
    test_dir = Path(data_dir) / config_name
    ensure_dir(str(test_dir))
    
    # Пути к файлам внутри папки теста
    wav_temp = str(test_dir / "temp_stereo.wav")
    ext = "jpg" if full_config["output_format"] == "jpeg" else full_config["output_format"]
    image_path = str(test_dir / f"spectrogram.{ext}")
    recovered_path = str(test_dir / "recovered.wav")
    metadata_path = str(test_dir / "metadata.json")
    report_path = str(test_dir / "report.txt")
    
    # Перенаправляем вывод в файл отчета
    tee = TeeOutput(report_path)
    sys.stdout = tee
    
    try:
        preset = PRESETS[full_config["active_preset"]]
        n_fft = preset["N_FFT"]
        hop_length = preset["HOP_LENGTH"]
        
        print(f"Параметры: FFT={n_fft}, HOP={hop_length}, SR=44100")
        print(f"Формат: {full_config['output_format']}, Качество: {full_config['output_quality']}")
        print(f"Lossless: {full_config['output_lossless']}")
        
        # Шаг 1: Конвертация MP3 -> WAV
        print("\n[1/3] Конвертация MP3 -> WAV")
        t_start = time.time()
        mp3_to_wav(mp3_file, wav_temp)
        y_orig, sr_orig = librosa.load(wav_temp, sr=None, mono=False)
        duration = len(y_orig.T if y_orig.ndim > 1 else y_orig) / sr_orig
        print(f"   Длительность: {duration:.2f} сек")
        conv_time = time.time() - t_start
        
        # Шаг 2: Кодирование в изображение
        print("\n[2/3] Кодирование в изображение")
        t_encode = time.time()
        metadata, saved_image_path = audio_to_image(wav_temp, image_path, n_fft, hop_length, full_config)
        encode_time = time.time() - t_encode
        
        # Сохраняем метаданные
        metadata_to_json(metadata, metadata_path)
        
        # Шаг 3: Декодирование
        print("\n[3/3] Восстановление аудио")
        t_decode = time.time()
        audio_recovered = image_to_audio(saved_image_path, recovered_path, metadata)
        decode_time = time.time() - t_decode
        
        # Вычисляем SNR
        snr = calculate_snr(y_orig, audio_recovered)
        
        # Итоговая статистика
        print("\n" + "="*60)
        print("ИТОГОВАЯ СТАТИСТИКА")
        print("="*60)
        
        image_size_mb = file_size_mb(saved_image_path)
        
        print(f"   Выходное изображение: {image_size_mb:.3f} MB")
        print(f"   SNR: {snr:.2f} dB")
        print(f"   Время конвертации: {conv_time:.1f} сек")
        print(f"   Время кодирования: {encode_time:.1f} сек")
        print(f"   Время декодирования: {decode_time:.1f} сек")
        print(f"   Общее время: {conv_time + encode_time + decode_time:.1f} сек")
        
        # Очищаем временный WAV
        clean_temp_files(wav_temp)
        
        result = {
            "test_name": config_name,
            "format": full_config["output_format"],
            "quality": full_config["output_quality"],
            "lossless": full_config["output_lossless"],
            "image_size_mb": round(image_size_mb, 4),
            "snr_db": round(snr, 2),
            "encode_time_sec": round(encode_time, 1),
            "decode_time_sec": round(decode_time, 1),
            "total_time_sec": round(conv_time + encode_time + decode_time, 1),
            "status": "success"
        }
        
    except Exception as e:
        print(f"\n   ❌ ОШИБКА: {e}")
        result = {
            "test_name": config_name,
            "format": full_config.get("output_format", "unknown"),
            "quality": full_config.get("output_quality", 0),
            "status": "error",
            "error_message": str(e)
        }
    
    finally:
        # Восстанавливаем стандартный вывод
        sys.stdout = tee.stdout
        tee.close()
        
        # Выводим результат в консоль
        print(f"   Результат сохранен: {report_path}")
        if result.get("status") == "success":
            print(f"   Размер изображения: {result['image_size_mb']:.3f} MB")
            print(f"   SNR: {result['snr_db']:.2f} dB")
    
    return result


def generate_summary_report(all_results: list, data_dir: str):
    """Генерирует сводный отчет по всем тестам."""
    summary_path = Path(data_dir) / "summary_report.txt"
    
    # Генерируем текстовый отчет
    with open(summary_path, 'w', encoding='utf-8') as f:
        f.write("="*90 + "\n")
        f.write("СВОДНЫЙ ОТЧЕТ ПО ТЕСТИРОВАНИЮ ФОРМАТОВ СЖАТИЯ\n")
        f.write(f"Дата: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write("="*90 + "\n\n")
        
        # Группируем по форматам
        formats = {}
        for r in all_results:
            fmt = r.get("format", "unknown")
            if fmt not in formats:
                formats[fmt] = []
            formats[fmt].append(r)
        
        for fmt, tests in formats.items():
            f.write(f"\n{'─'*90}\n")
            f.write(f"ФОРМАТ: {fmt.upper()}\n")
            f.write(f"{'─'*90}\n")
            f.write(f"{'Тест':<20} {'Качество':>8} {'Lossless':>8} {'Размер(MB)':>12} {'SNR(dB)':>10} {'Время(сек)':>10} {'Статус':>10}\n")
            f.write(f"{'─'*90}\n")
            
            for t in tests:
                name = t["test_name"]
                quality = t.get("quality", "N/A")
                lossless = "Да" if t.get("lossless") else "Нет"
                size = f"{t['image_size_mb']:.4f}" if t.get("image_size_mb") else "N/A"
                snr = f"{t['snr_db']:.2f}" if t.get("snr_db") else "N/A"
                total_time = f"{t['total_time_sec']:.1f}" if t.get("total_time_sec") else "N/A"
                status = t.get("status", "unknown")
                
                f.write(f"{name:<20} {quality:>8} {lossless:>8} {size:>12} {snr:>10} {total_time:>10} {status:>10}\n")
        
        # Топ-3 по размеру
        successful = [r for r in all_results if r.get("status") == "success" and r.get("image_size_mb")]
        if successful:
            f.write(f"\n{'='*90}\n")
            f.write("ТОП-3 ПО МИНИМАЛЬНОМУ РАЗМЕРУ ИЗОБРАЖЕНИЯ\n")
            f.write(f"{'='*90}\n")
            sorted_by_size = sorted(successful, key=lambda x: x["image_size_mb"])
            for i, r in enumerate(sorted_by_size[:3], 1):
                f.write(f"{i}. {r['test_name']}: {r['image_size_mb']:.4f} MB (SNR: {r['snr_db']:.2f} dB)\n")
            
            f.write(f"\n{'='*90}\n")
            f.write("ТОП-3 ПО МАКСИМАЛЬНОМУ SNR\n")
            f.write(f"{'='*90}\n")
            sorted_by_snr = sorted(successful, key=lambda x: x["snr_db"], reverse=True)
            for i, r in enumerate(sorted_by_snr[:3], 1):
                f.write(f"{i}. {r['test_name']}: {r['snr_db']:.2f} dB (Размер: {r['image_size_mb']:.4f} MB)\n")
    
    print(f"\n{'='*70}")
    print(f"Сводный отчет сохранен:")
    print(f"   {summary_path}")


def run_all_tests():
    """Запускает все тесты."""
    print("="*70)
    print("ЗАПУСК ТЕСТИРОВАНИЯ ФОРМАТОВ СЖАТИЯ")
    print(f"Пресет: 75p_n4096 (N_FFT=4096, HOP_LENGTH=1024)")
    print(f"Всего тестов: {len(TEST_CONFIGS)}")
    print("="*70)
    
    # Создаем корневую папку для результатов
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    data_dir = f"test_results_{timestamp}"
    ensure_dir(data_dir)
    
    # Базовые настройки из DEFAULT_CONFIG
    base_config = {
        "mp3_file": DEFAULT_CONFIG["mp3_file"],
        "trim_enabled": DEFAULT_CONFIG.get("trim_enabled", False),
        "trim_start": DEFAULT_CONFIG.get("trim_start", 0.0),
        "trim_end": DEFAULT_CONFIG.get("trim_end", 0.0),
        "phase_generate_iterations": DEFAULT_CONFIG.get("phase_generate_iterations", 100),
        "phase_generate_random_seed": DEFAULT_CONFIG.get("phase_generate_random_seed", -1),
        "griffin_lim_mode": DEFAULT_CONFIG.get("griffin_lim_mode", "fast"),
        "griffin_lim_parallel": DEFAULT_CONFIG.get("griffin_lim_parallel", True),
        "early_stop_enabled": DEFAULT_CONFIG.get("early_stop_enabled", True),
        "early_stop_threshold": DEFAULT_CONFIG.get("early_stop_threshold", 0.0001),
        "early_stop_patience": DEFAULT_CONFIG.get("early_stop_patience", 10),
        "low_cut_enabled": DEFAULT_CONFIG.get("low_cut_enabled", False),
        "high_cut_enabled": DEFAULT_CONFIG.get("high_cut_enabled", False),
        "avif_speed": DEFAULT_CONFIG.get("avif_speed", 0),
        "avif_subsampling": DEFAULT_CONFIG.get("avif_subsampling", "4:2:0"),
        "data_dir": data_dir,
    }
    
    mp3_file = base_config["mp3_file"]
    if not os.path.exists(mp3_file):
        print(f"\n❌ Ошибка: Файл {mp3_file} не найден!")
        print("Поместите MP3 файл 'track.mp3' в папку с программой или укажите путь в config.py")
        return
    
    all_results = []
    total_tests = len(TEST_CONFIGS)
    
    for i, (config_name, test_config) in enumerate(TEST_CONFIGS.items(), 1):
        print(f"\n{'▀'*70}")
        print(f"ТЕСТ {i}/{total_tests}: {config_name}")
        print(f"{'▀'*70}")
        
        result = run_single_test(config_name, test_config, base_config, mp3_file, data_dir)
        all_results.append(result)
    
    # Генерируем сводный отчет
    generate_summary_report(all_results, data_dir)
    
    # Подсчитываем статистику
    successful = [r for r in all_results if r.get("status") == "success"]
    failed = [r for r in all_results if r.get("status") == "error"]
    
    print(f"\n{'='*70}")
    print(f"ТЕСТИРОВАНИЕ ЗАВЕРШЕНО")
    print(f"{'='*70}")
    print(f"   Всего тестов: {total_tests}")
    print(f"   Успешно: {len(successful)}")
    print(f"   С ошибками: {len(failed)}")
    print(f"   Результаты сохранены в: {data_dir}/")
    
    if failed:
        print(f"\n   Тесты с ошибками:")
        for r in failed:
            print(f"   - {r['test_name']}: {r.get('error_message', 'Неизвестная ошибка')}")


if __name__ == "__main__":
    run_all_tests()