import gradio as gr
import librosa
import numpy as np
import soundfile as sf
import tempfile
import subprocess
import os

def extract_audio_if_needed(audio_path):
    if audio_path.endswith(".mp4"):
        temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
        # Use ffmpeg to extract audio
        subprocess.run([
            "ffmpeg", "-y", "-i", audio_path,
            "-vn", "-acodec", "pcm_s16le", "-ar", "44100", temp_wav.name
        ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        return temp_wav.name
    return audio_path

def summarize_audio(audio):
    if isinstance(audio, tuple):
        audio_path = audio[0]
    else:
        audio_path = audio

    audio_path = extract_audio_if_needed(audio_path)

    y, sr = librosa.load(audio_path, sr=None)
    chroma = librosa.feature.chroma_cqt(y=y, sr=sr)
    similarity = librosa.segment.recurrence_matrix(chroma, mode='affinity', sym=True)
    segments = librosa.segment.agglomerative(similarity, k=4)
    boundaries = librosa.frames_to_time(np.where(np.diff(segments))[0], hop_length=512, sr=sr)

    summary_audio = []
    for i in range(min(3, len(boundaries))):
        start = int(boundaries[i] * sr)
        end = start + int(15 * sr)  # 15 seconds per segment
        summary_audio.append(y[start:end])

    combined = np.concatenate(summary_audio)
    tmp_out = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    sf.write(tmp_out.name, combined, sr)
    return tmp_out.name

demo = gr.Interface(
    fn=summarize_audio,
    inputs=gr.Audio(type="filepath"),
    outputs=gr.Audio(type="filepath"),
    title="Jam Summarizer",
    description="Upload a long jam session or song. This tool finds key musical sections and outputs a ~1 min highlight reel."
)

if __name__ == "__main__":
    demo.launch()