# main.py
import os
import time
import numpy as np
import librosa
from pathlib import Path
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


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 full_pipeline(config: dict):
    """Полный цикл: аудио -> изображение -> аудио (с генерацией фазы)."""
    preset = PRESETS[config["active_preset"]]
    n_fft = preset["N_FFT"]
    hop_length = preset["HOP_LENGTH"]
    data_dir = config.get("data_dir", "data")
    ensure_dir(data_dir)

    # --- Информация о параметрах ---
    print(f"Параметры: FFT={n_fft}, HOP={hop_length}, SR=44100")
    
    if config.get("trim_enabled", False):
        trim_end_str = f"{config.get('trim_end', 0)}s" if config.get('trim_end', 0) > 0 else "конец"
        print(f"Диапазон обработки: {config.get('trim_start', 0)}s - {trim_end_str}")

    wav_temp = Path(data_dir) / "temp_stereo.wav"

    # Шаг 1: Конвертация
    print("\n[1/3] Конвертация MP3 -> WAV")
    mp3_to_wav(config["mp3_file"], str(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} сек")

    # Шаг 2: Кодирование
    print("\n[2/3] Кодирование в изображение")
    
    fmt = config.get("output_format", "png").lower()
    ext = "jpg" if fmt == "jpeg" else fmt
    image_path = str(Path(data_dir) / f"spectrogram.{ext}")

    t_encode = time.time()
    metadata, _ = audio_to_image(wav_temp, image_path, n_fft, hop_length, config)
    encode_time = time.time() - t_encode

    metadata_path = str(Path(data_dir) / "metadata.json")
    metadata_to_json(metadata, metadata_path)

    # Шаг 3: Декодирование
    print("\n[3/3] Восстановление аудио")
    recovered_path = str(Path(data_dir) / "recovered.wav")
    t_decode = time.time()
    audio_recovered = image_to_audio(image_path, recovered_path, metadata)
    decode_time = time.time() - t_decode

    # Итоговая статистика
    print("\n" + "="*60)
    print("ИТОГОВАЯ СТАТИСТИКА")
    print("="*60)
    
    input_size_mb = file_size_mb(config["mp3_file"])
    image_size_kb = file_size_kb(image_path)
    wav_size_mb = file_size_mb(recovered_path)
    
    print(f"   Входной MP3: {input_size_mb:.1f} MB")
    print(f"   Выходное изображение: {image_size_kb:.0f} KB")
    
    if config.get("trim_enabled", False):
        trim_end = config.get("trim_end", 0) if config.get("trim_end", 0) > 0 else duration
        fragment_duration = trim_end - config.get("trim_start", 0)
        original_total = input_size_mb
        fragment_ratio = fragment_duration / duration
        effective_input = original_total * fragment_ratio
        print(f"   Обработанный фрагмент: {fragment_duration:.1f} сек из {duration:.1f} сек")
    else:
        effective_input = input_size_mb
    
    image_size_mb = image_size_kb / 1024
    compression_ratio = effective_input / image_size_mb if image_size_mb > 0 else 0
    print(f"   Степень сжатия: {compression_ratio:.1f}:1 (MP3:{effective_input:.1f}MB -> Изобр:{image_size_mb:.2f}MB)")
    print(f"   Восстановленный WAV: {wav_size_mb:.1f} MB")
    print(f"   Время кодирования: {encode_time:.1f} сек")
    print(f"   Время декодирования: {decode_time:.1f} сек")
    
    clean_temp_files(str(wav_temp))
    return audio_recovered


if __name__ == "__main__":
    full_pipeline(DEFAULT_CONFIG)