import librosa
import numpy as np
import subprocess
import os
import sys
import argparse
import json
import soundfile as sf
import tempfile

def fix_grid(input_file, output_file, target_bpm=None, verbose=False):
    temp_wav = None
    try:
        if verbose: print(f"[*] A carregar: {input_file}")
        # Carregar mantendo stereo e sample rate original
        y, sr = librosa.load(input_file, sr=None, mono=False)
        
        # Se for stereo, o librosa devolve (2, samples). 
        # Algumas funções esperam (samples, 2) ou precisam de tratamento especial.
        y_mono = librosa.to_mono(y) if y.ndim > 1 else y

        if verbose: print("[*] A detetar transientes (picos) e analisar beats...")
        # Analisamos na versão mono para maior precisão de deteção
        onset_env = librosa.onset.onset_strength(y=y_mono, sr=sr)
        onset_frames = librosa.onset.onset_detect(onset_envelope=onset_env, sr=sr)
        
        tempo, beat_frames = librosa.beat.beat_track(onset_envelope=onset_env, sr=sr)
        
        # Forçar o alinhamento (snap) do batimento com o pico de onda mais próximo
        snapped_beat_frames = []
        for beat in beat_frames:
            if len(onset_frames) > 0:
                nearest_onset = min(onset_frames, key=lambda x: abs(x - beat))
                # Distância máxima aceite (aprox. 150ms)
                if abs(nearest_onset - beat) < librosa.time_to_frames(0.15, sr=sr):
                    snapped_beat_frames.append(nearest_onset)
                else:
                    snapped_beat_frames.append(beat)
            else:
                snapped_beat_frames.append(beat)
                
        beat_frames = np.array(snapped_beat_frames)
        
        if isinstance(tempo, np.ndarray):
            tempo = tempo[0]
            
        if target_bpm is None:
            target_bpm = round(float(tempo))
            if verbose: print(f"[*] BPM alvo detetado: {tempo:.2f} -> Arredondado para: {target_bpm}")
        else:
            if verbose: print(f"[*] BPM alvo definido pelo utilizador: {target_bpm}")

        beat_times = librosa.frames_to_time(beat_frames, sr=sr)
        
        if len(beat_times) < 2:
            print("[!] Erro: Não foram detetados beats suficientes.")
            return

        # Criar WAV temporário para o Rubber Band (garante compatibilidade se o input for MP3)
        fd, temp_wav = tempfile.mkstemp(suffix=".wav")
        os.close(fd)
        if verbose: print(f"[*] A criar ficheiro temporário de trabalho...")
        # soundfile espera (samples, channels)
        sf.write(temp_wav, y.T if y.ndim > 1 else y, sr)

        # Calcular grid ideal
        t_ideal_interval = 60.0 / target_bpm
        first_beat_time = beat_times[0]
        
        timemap_path = "timemap.txt"
        valid_points = 0
        with open(timemap_path, "w") as f:
            last_target_frame = -1
            for t_real in beat_times:
                # O Librosa pode saltar beats ou detetar notas fantasma. 
                # Mapeamos o tempo real para o múltiplo ideal mais próximo, 
                # saltando possíveis detecções duplicadas e não desfazendo a grelha por beats em falta.
                delta = t_real - first_beat_time
                beats_elapsed = round(delta / t_ideal_interval)
                t_target = first_beat_time + (beats_elapsed * t_ideal_interval)
                
                source_frame = int(t_real * sr)
                target_frame = int(t_target * sr)
                
                if target_frame > last_target_frame:
                    f.write(f"{source_frame} {target_frame}\n")
                    last_target_frame = target_frame
                    valid_points += 1
                
        if verbose: print(f"[*] Timemap gerado com {valid_points} pontos de correção válidos.")

        if verbose: print("[*] A aplicar warping com Rubber Band (Motor R3)...")
        
        # O rácio global ajuda o rubberband a estimar a duração final
        ratio = float(tempo) / target_bpm
        
        cmd = [
            "rubberband",
            "--fine",
            f"-T{ratio}",
            "-M", timemap_path,
            temp_wav,
            output_file
        ]

        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode != 0:
            print(f"[!] Erro no Rubber Band: {result.stderr}")
        else:
            if verbose: print(f"[+] Sucesso! Ficheiro exportado: {output_file}")
            
        if os.path.exists(timemap_path):
            os.remove(timemap_path)

        # Metadados
        metadata = {
            "input": input_file,
            "original_bpm": float(tempo),
            "target_bpm": float(target_bpm),
            "beats_count": len(beat_times),
            "first_beat_time": float(first_beat_time)
        }
        meta_path = output_file.rsplit('.', 1)[0] + ".json"
        with open(meta_path, "w") as f:
            json.dump(metadata, f, indent=4)

    except Exception as e:
        print(f"[!] Ocorreu um erro: {str(e)}")
    finally:
        if temp_wav and os.path.exists(temp_wav):
            os.remove(temp_wav)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Forca-Grid: Corrige o drift de BPM em faixas de áudio.")
    parser.add_argument("input", help="Ficheiro de áudio de entrada")
    parser.add_argument("output", help="Ficheiro de áudio de saída")
    parser.add_argument("--bpm", type=float, help="BPM alvo")
    parser.add_argument("-v", "--verbose", action="store_true", help="Detalhes")
    
    args = parser.parse_args()
    fix_grid(args.input, args.output, args.bpm, args.verbose)