import os
import gradio as gr
import shutil
from faster_whisper import WhisperModel
import demucs.separate
from matchering import process as match_process
from mutagen.mp3 import MP3
from mutagen.id3 import ID3, TIT2, TPE1, TALB, APIC

os.makedirs("uploads", exist_ok=True)
os.makedirs("outputs", exist_ok=True)

# 1. Stem Splitter Engine (Demucs CPU)
def split_stems(audio_file):
    if audio_file is None: 
        return [None]*4
    demucs.separate.main(["-n", "htdemucs", "-o", "outputs/stems", audio_file])
    track_name = os.path.splitext(os.path.basename(audio_file))[0]
    stem_dir = f"outputs/stems/htdemucs/{track_name}"
    
    stems = []
    for stem in ["vocals.wav", "drums.wav", "bass.wav", "other.wav"]:
        path = os.path.join(stem_dir, stem)
        stems.append(path if os.path.exists(path) else None)
    return stems

# 2. Synced Lyrics Engine (Whisper CPU)
def generate_lrc(audio_file):
    if audio_file is None: 
        return "Please upload an audio file."
    model = WhisperModel("tiny", device="cpu", compute_type="float32")
    segments, _ = model.transcribe(audio_file, beam_size=1)
    
    lrc_lines = []
    for segment in segments:
        start_min = int(segment.start // 60)
        start_sec = segment.start % 60
        time_str = f"[{start_min:02d}:{start_sec:05.2f}]"
        lrc_lines.append(f"{time_str} {segment.text.strip()}")
    
    return "\n".join(lrc_lines)

# 3. Reference Mastering Engine (Matchering)
def master_audio(target, reference):
    if not target or not reference: 
        return None
    out_master = "outputs/mastered_track.wav"
    if os.path.exists(out_master): 
        os.remove(out_master)
    match_process(target=target, reference=reference, results=[out_master])
    return out_master

# 4. Metadata Tagging Engine (Mutagen)
def tag_audio(audio, title, artist, album, cover):
    if not audio: 
        return None
    out_tagged = "outputs/tagged_track.mp3"
    shutil.copy(audio, out_tagged)
    
    audio_obj = MP3(out_tagged, ID3=ID3)
    try:
        audio_obj.add_tags()
    except Exception:
        pass
    
    audio_obj.tags.add(TIT2(encoding=3, text=title))
    audio_obj.tags.add(TPE1(encoding=3, text=artist))
    audio_obj.tags.add(TALB(encoding=3, text=album))
    
    if cover:
        with open(cover, "rb") as f:
            art_data = f.read()
        audio_obj.tags.add(APIC(encoding=3, mime='image/jpeg', type=3, desc=u'Cover', data=art_data))
        
    audio_obj.save()
    return out_tagged

# UI Setup
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Free Cloud Audio AI Suite\nRun stem splitting, lyric syncing, mastering, and tagging 24/7 without billing limits.")
    
    with gr.Tab("Stem Splitter"):
        audio_input = gr.Audio(type="filepath", label="Upload Song")
        split_btn = gr.Button("Split Stems")
        vocals_out = gr.Audio(label="Vocals")
        drums_out = gr.Audio(label="Drums")
        bass_out = gr.Audio(label="Bass")
        other_out = gr.Audio(label="Instrumental/Other")
        split_btn.click(split_stems, inputs=audio_input, outputs=[vocals_out, drums_out, bass_out, other_out])
        
    with gr.Tab("Synced Lyrics"):
        lyric_input = gr.Audio(type="filepath", label="Upload Vocal Track")
        lyric_btn = gr.Button("Generate Timed Lyrics")
        lrc_output = gr.Textbox(label="LRC File", show_copy_button=True)
        lyric_btn.click(generate_lrc, inputs=lyric_input, outputs=lrc_output)
        
    with gr.Tab("AI Mastering"):
        target_input = gr.Audio(type="filepath", label="Upload Your Mix (WAV)")
        ref_input = gr.Audio(type="filepath", label="Upload Reference Track (WAV)")
        master_btn = gr.Button("Apply Mastering")
        master_output = gr.Audio(label="Mastered Output")
        master_btn.click(master_audio, inputs=[target_input, ref_input], outputs=master_output)
        
    with gr.Tab("Metadata Tagging"):
        tag_audio_in = gr.Audio(type="filepath", label="Upload MP3")
        in_title = gr.Textbox(label="Song Title")
        in_artist = gr.Textbox(label="Artist Name")
        in_album = gr.Textbox(label="Album Name")
        in_cover = gr.Image(type="filepath", label="Cover Art (JPEG)")
        tag_btn = gr.Button("Tag Track")
        tag_output = gr.File(label="Ready for Distribution")
        tag_btn.click(tag_audio, inputs=[tag_audio_in, in_title, in_artist, in_album, in_cover], outputs=tag_output)

demo.queue().launch()