# encoder.py
import numpy as np
import librosa
from PIL import Image
from pathlib import Path


def _find_high_cut_auto(mag_mid_db: np.ndarray, mag_side_db: np.ndarray, freqs: np.ndarray,
                        threshold_db: float = -80, freq_min: float = 8000, margin_db: float = 10):
    mean_mag = np.maximum(np.mean(mag_mid_db, axis=1), np.mean(mag_side_db, axis=1))
    
    from scipy.ndimage import uniform_filter1d
    smooth_size = max(3, len(mean_mag) // 20)
    mean_mag_smooth = uniform_filter1d(mean_mag, size=smooth_size)
    
    effective_threshold = threshold_db + margin_db
    below_threshold = mean_mag_smooth < effective_threshold
    min_idx = np.searchsorted(freqs, freq_min)
    
    cut_bin = len(freqs) - 1
    for i in range(min_idx, len(freqs) - 5):
        if np.all(below_threshold[i:i+5]):
            cut_bin = max(min_idx, i)
            break
            
    print(f"   Автосрез ВЧ: порог {effective_threshold:.0f} dB, частота среза {freqs[cut_bin]:.0f} Гц")
    return freqs[cut_bin], cut_bin


def _check_image_size_limit(img: Image.Image, config: dict) -> tuple:
    """Проверяет ограничения размера для разных форматов."""
    fmt = config.get("output_format", "png").lower()
    
    # Лимиты для разных форматов
    if fmt == "webp":
        MAX_SIZE = 16383
    elif fmt == "avif":
        MAX_SIZE = 65535  # AVIF поддерживает очень большие изображения
    elif fmt in ("jpeg", "jpg"):
        MAX_SIZE = 65535
    else:
        MAX_SIZE = float('inf')
    
    width, height = img.size
    scale = 1.0
    was_scaled = False
    
    if width > MAX_SIZE or height > MAX_SIZE:
        scale = min(MAX_SIZE / width, MAX_SIZE / height) if max(width, height) > 0 else 1.0
        new_width = max(1, int(width * scale))
        new_height = max(1, int(height * scale))
        
        print(f"\n   ⚠️  Размер изображения ({width}x{height}) превышает лимит формата {fmt.upper()} ({MAX_SIZE}px)")
        print(f"   Автоматическое масштабирование: {width}x{height} -> {new_width}x{new_height}")
        
        method = _get_scale_method(config.get("scale_method", "lanczos"))
        img = img.resize((new_width, new_height), method)
        was_scaled = True
    
    return img, was_scaled, scale


def _save_image_with_quality(img: Image.Image, image_path: str, config: dict):
    fmt = config.get("output_format", "png").lower()
    quality = max(1, min(100, config.get("output_quality", 95)))
    lossless = config.get("output_lossless", False)

    if fmt == "png":
        compress_level = max(1, min(9, 9 - int((quality / 100) * 8)))
        img.save(image_path, 'PNG', compress_level=compress_level)
        format_str = f"PNG (сжатие: {compress_level}/9)"
    
    elif fmt == "webp":
        width, height = img.size
        if width > 16383 or height > 16383:
            scale = min(16383 / width, 16383 / height)
            img = img.resize((int(width * scale), int(height * scale)), Image.LANCZOS)
        img.save(image_path, 'WEBP', lossless=lossless, quality=quality, method=6)
        format_str = f"WEBP {'(lossless)' if lossless else f'(качество: {quality}/100)'}"
    
    elif fmt in ("jpeg", "jpg"):
        if img.mode != 'RGB':
            background = Image.New('RGB', img.size, (255, 255, 255))
            background.paste(img, mask=img.split()[-1] if img.mode == 'RGBA' else None)
            img = background
        img.save(image_path, 'JPEG', quality=quality, optimize=True, progressive=True)
        format_str = f"JPEG (качество: {quality}/100)"
    
    elif fmt == "avif":
        # AVIF поддержка через pillow-avif-plugin или pillow-heif
        avif_speed = config.get("avif_speed", 6)
        avif_subsampling = config.get("avif_subsampling", "4:4:4")
        
        try:
            # Пробуем сохранить как AVIF
            img.save(
                image_path, 
                'AVIF', 
                quality=quality,
                lossless=lossless,
                speed=avif_speed,
                subsampling=avif_subsampling
            )
            format_str = f"AVIF {'(lossless)' if lossless else f'(качество: {quality}/100, скорость: {avif_speed})'}"
        except Exception as e:
            print(f"   ⚠️  Ошибка сохранения AVIF: {e}")
            print(f"   Попытка сохранения через pillow-heif...")
            try:
                # Альтернативный метод через pillow-heif
                img.save(image_path, 'AVIF', quality=quality)
                format_str = f"AVIF (качество: {quality}/100)"
            except Exception as e2:
                print(f"   ⚠️  Ошибка: {e2}")
                print(f"   Сохраняю как PNG вместо AVIF")
                # Fallback на PNG
                png_path = str(Path(image_path).with_suffix('.png'))
                img.save(png_path, 'PNG', compress_level=6)
                format_str = f"PNG (fallback, сжатие: 6/9)"
                image_path = png_path
    
    else:
        img.save(image_path, 'PNG', compress_level=9)
        format_str = "PNG (макс. сжатие)"
    
    size_mb = Path(image_path).stat().st_size / (1024 * 1024)
    print(f"   Формат: {format_str} ({size_mb:.2f} MB)")
    
    return image_path


def _get_scale_method(method_name: str) -> int:
    methods = {
        "lanczos": Image.LANCZOS,
        "bilinear": Image.BILINEAR,
        "bicubic": Image.BICUBIC,
        "nearest": Image.NEAREST,
    }
    return methods.get(method_name.lower(), Image.LANCZOS)


def audio_to_image(wav_path: str, image_path: str, n_fft: int, hop_length: int, config: dict):
    target_sr = 44100
    y, sr = librosa.load(wav_path, sr=target_sr, mono=False)
    
    if y.ndim == 1:
        y = np.vstack([y, np.zeros_like(y)])
    
    # --- ОБРЕЗКА ПО ДИАПАЗОНУ ---
    if config.get("trim_enabled", False):
        trim_start = config.get("trim_start", 0.0)
        trim_end = config.get("trim_end", 0.0)
        
        start_sample = int(trim_start * sr)
        end_sample = int(trim_end * sr) if trim_end > 0 else y.shape[1]
        
        start_sample = max(0, min(start_sample, y.shape[1]))
        end_sample = max(start_sample, min(end_sample, y.shape[1]))
        
        y = y[:, start_sample:end_sample]
        duration = y.shape[1] / sr
        
        print(f"   Аудио обрезано: {trim_start:.2f}s - {trim_end if trim_end > 0 else duration:.2f}s")
        print(f"   Длительность фрагмента: {duration:.2f} сек")
        
    print(f"   Аудио: {y.shape}, {sr} Гц")
    
    # Преобразование в Mid/Side
    mid = (y[0] + y[1]) * 0.5
    side = (y[0] - y[1]) * 0.5

    D_mid = librosa.stft(mid, n_fft=n_fft, hop_length=hop_length, window='hann')
    D_side = librosa.stft(side, n_fft=n_fft, hop_length=hop_length, window='hann')
    n_freqs_orig, n_frames = D_mid.shape
    freqs = librosa.fft_frequencies(sr=sr, n_fft=n_fft)

    low_cut_bin = 0
    if config.get("low_cut_enabled", True) and config.get("low_cut_freq", 0) > 0:
        low_cut_bin = np.searchsorted(freqs, config["low_cut_freq"])

    high_cut_bin = len(freqs)
    if config.get("high_cut_enabled", True):
        if config.get("high_cut_mode", "auto") == "auto":
            mag_mid_db = librosa.amplitude_to_db(np.abs(D_mid), ref=np.max)
            mag_side_db = librosa.amplitude_to_db(np.abs(D_side), ref=np.max)
            _, high_cut_bin = _find_high_cut_auto(
                mag_mid_db, mag_side_db, freqs,
                config.get("high_cut_auto_threshold_db", -80),
                config.get("high_cut_auto_freq_min", 8000),
                config.get("high_cut_auto_margin_db", 10)
            )
        else:
            high_cut_bin = np.searchsorted(freqs, config.get("high_cut_freq", 16000))

    D_mid = D_mid[low_cut_bin:high_cut_bin, :]
    D_side = D_side[low_cut_bin:high_cut_bin, :]
    n_freqs = D_mid.shape[0]
    high_cut_bins_removed = n_freqs_orig - high_cut_bin

    mag_mid, mag_side = np.abs(D_mid), np.abs(D_side)
    ref_mid, ref_side = np.max(mag_mid), np.max(mag_side)
    mag_min = -120

    mag_mid_db = librosa.amplitude_to_db(mag_mid, ref=ref_mid)
    mag_side_db = librosa.amplitude_to_db(mag_side, ref=ref_side)
    mag_mid_norm = np.clip((mag_mid_db - mag_min) / (-mag_min) * 255, 0, 255).astype(np.uint8)
    mag_side_norm = np.clip((mag_side_db - mag_min) / (-mag_min) * 255, 0, 255).astype(np.uint8)

    rgb = np.zeros((n_freqs, n_frames, 3), dtype=np.uint8)
    rgb[:, :, 1] = mag_mid_norm
    rgb[:, :, 2] = mag_side_norm
    rgb = np.flipud(rgb)
    
    img = Image.fromarray(rgb, 'RGB')
    
    # Проверка лимитов размера для lossy форматов
    fmt = config.get("output_format", "png").lower()
    if fmt in ("webp", "jpeg", "jpg", "avif"):
        img, was_auto_scaled, scale_x = _check_image_size_limit(img, config)
    else:
        was_auto_scaled = False
        scale_x = 1.0
    
    # Пользовательское масштабирование
    if config.get("scale_enabled", False):
        user_scale_x = config.get("scale_x", 1)
        user_scale_y = config.get("scale_y", 1)
        new_w = max(1, img.width // user_scale_x)
        new_h = max(1, img.height // user_scale_y)
        method = _get_scale_method(config.get("scale_method", "lanczos"))

        print(f"\n   Масштабирование изображения:")
        print(f"   Размер ДО: {img.width}x{img.height}")
        print(f"   Размер ПОСЛЕ: {new_w}x{new_h}")
        print(f"   Масштаб: 1/{user_scale_x} по ширине, 1/{user_scale_y} по высоте")
        print(f"   Метод интерполяции: {config.get('scale_method', 'lanczos')}")

        img = img.resize((new_w, new_h), method)
    
    # Сохраняем изображение и получаем актуальный путь (может измениться при fallback)
    saved_image_path = _save_image_with_quality(img, image_path, config)

    # Всегда используем генерацию фазы
    gl_iterations = config.get("phase_generate_iterations", 100)
    gl_random_seed = config.get("phase_generate_random_seed", 5)
    gl_mode = config.get("griffin_lim_mode", "standard")
    gl_parallel = config.get("griffin_lim_parallel", True)
    
    print(f"   Фаза будет сгенерирована при декодировании")
    print(f"   Режим Griffin-Lim: {gl_mode}")
    print(f"   Итераций: {gl_iterations}")
    print(f"   Random seed: {gl_random_seed}")
    
    if config.get("early_stop_enabled", True):
        print(f"   Ранняя остановка: ВКЛ. Порог: {config.get('early_stop_threshold', 0.0001)}")
    
    if gl_mode == "multi_scale":
        ms_factor = config.get("multi_scale_factor", 4)
        ms_coarse = config.get("multi_scale_coarse_iterations", -1)
        ms_fine = config.get("multi_scale_fine_iterations", -1)
        
        print(f"   Multi-scale factor: {ms_factor}")
        print(f"   Coarse iterations: {'all' if ms_coarse == -1 else ms_coarse}")
        print(f"   Fine iterations: {'half' if ms_fine == -1 else ms_fine}")
        
        if gl_parallel:
            print(f"   Параллельно: да")
    else:
        ms_factor = 4
        ms_coarse = -1
        ms_fine = -1
        
        if gl_parallel:
            print(f"   Параллельно: да")

    result_metadata = {
        'shape': (n_freqs, n_frames),
        'shape_original': (n_freqs_orig, n_frames),
        'low_cut_bin': low_cut_bin,
        'high_cut_bin': high_cut_bin,
        'high_cut_bins_removed': high_cut_bins_removed,
        'sr': sr,
        'n_fft': n_fft,
        'hop_length': hop_length,
        'original_length': len(mid),
        'n_channels': y.shape[0],
        'mag_min': mag_min,
        'ref_mid': float(ref_mid),
        'ref_side': float(ref_side),
        'phase_method': 'generate',
        'flipped': True,
        'auto_scaled': was_auto_scaled,
        'auto_scale_x': scale_x,
        'phase_generate_iterations': gl_iterations,
        'phase_generate_random_seed': gl_random_seed,
        'griffin_lim_mode': gl_mode,
        'griffin_lim_parallel': gl_parallel,
        'early_stop_threshold': config.get("early_stop_threshold", 0.0001),
        'early_stop_patience': config.get("early_stop_patience", 15),
        'multi_scale_factor': ms_factor,
        'multi_scale_coarse_iterations': ms_coarse,
        'multi_scale_fine_iterations': ms_fine,
    }
    
    return result_metadata, saved_image_path