import torch
from typing import Tuple, Dict
from joblib import Parallel, delayed
from resemble_enhance.enhancer.inference import denoise, enhance
import spaces

def denoise_audio(audio: torch.Tensor, sr: int, device: str) -> Tuple[torch.Tensor, int]:
    """
    Réduction de bruit sur un tenseur audio.
    """
    try:
        return denoise(audio, sr, device)
    except Exception as e:
        print(f"[denoise_audio] Erreur: {e}")
        return audio, sr


def enhance_audio(
    audio: torch.Tensor,
    sr: int,
    device: str,
    solver: str = "midpoint",
    nfe: int = 128,
    lambd: float = 0.1,
    tau: float = 0.01
) -> Tuple[torch.Tensor, int]:
    """
    Amélioration audio (résolution d'ODE).
    """
    try:
        return enhance(
            audio,
            sr,
            device,
            solver=solver.lower(),
            nfe=nfe,
            lambd=lambd,
            tau=tau
        )
    except Exception as e:
        print(f"[enhance_audio] Erreur: {e}")
        return audio, sr

@spaces.GPU(duration=5)
def process_audio(
    audio: torch.Tensor,
    sr: int,
    device: str = "cpu",
    solver: str = "midpoint",
    nfe: int = 128,
    tau: float = 0.01,
    denoise_before: bool = False,
) -> Dict[str, Tuple[torch.Tensor, int]]:
    """
    Pipeline complet de traitement audio.

    Args:
        audio: Tensor d'entrée de forme (1, T).
        sr: Fréquence d'échantillonnage.
        device: 'cpu' ou 'cuda'.
        solver: ODE solver.
        nfe: nombre d'évaluations de la fonction.
        tau: paramètre tau.
        denoise_before: si True, lambd=0.9, sinon lambd=0.1.
        parallel: si True, exécute denoise et enhance en parallèle (utile si indépendants).

    Returns:
        dict avec clés 'denoised' et 'enhanced', valeurs (tensor, sr).
    """
    # Choix de lambda
    lambd = 0.9 if denoise_before else 0.1
    # Exécution séquentielle (débruitage puis amélioration)
    denoised = denoise_audio(audio, sr, device)
    # Applique enhance sur le résultat du débruitage
    enhanced = enhance_audio(
        denoised[0],
        denoised[1],
        device,
        solver,
        nfe,
        lambd,
        tau
    )
    return {
        "denoised": denoised,
        "enhanced": enhanced
    }