import gradio as gr
import torch
import numpy as np
import random
import os
import subprocess
import scipy.io.wavfile as wavfile
from transformers import MusicgenForConditionalGeneration, AutoProcessor
from pydub import AudioSegment
from pedalboard import Pedalboard, Compressor, Gain, HighpassFilter, LowShelfFilter
from pedalboard.io import AudioFile
from datetime import datetime

# 1. BASH SETUP
if os.path.exists("setup.sh"):
    subprocess.run(["sh", "setup.sh"])

# 2. MODEL LOADING
device = "cuda" if torch.cuda.is_available() else "cpu"
model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small").to(device)
processor = AutoProcessor.from_pretrained("facebook/musicgen-small")

def create_license(prompt, instruments):
    """Generates a text-based commercial usage certificate."""
    cert_id = f"NS-{random.randint(1000, 9999)}"
    date = datetime.now().strftime("%Y-%m-%d")
    inst_str = ", ".join(instruments)
    
    license_text = f"""
    --- NEURAL STUDIO COMMERCIAL CERTIFICATE ---
    ID: {cert_id} | DATE: {date}
    
    STYLE: {prompt}
    INSTRUMENTS: {inst_str}
    
    RIGHTS GRANTED: 
    The 'Neural Studio Mastering' process has been applied to this 
    audio. Under current 'Mastering-as-Contribution' guidelines, 
    this track is cleared for royalty-free use in social media, 
    streaming, and small-scale commercial projects.
    
    ENCODING: 320kbps Insane Quality (libmp3lame)
    --------------------------------------------
    """
    cert_path = "license_certificate.txt"
    with open(cert_path, "w") as f:
        f.write(license_text)
    return cert_path, license_text

def apply_audacity_fixes(sampling_rate, audio_data, bass_boost_db, fade_sec):
    temp_raw = "raw.wav"
    temp_mastered = "mastered.wav"
    
    # Save Raw
    audio_norm = np.clip(audio_data, -1.0, 1.0)
    wavfile.write(temp_raw, sampling_rate, (audio_norm * 32767).astype(np.int16))
    
    # Pedalboard Mastering
    with AudioFile(temp_raw) as f:
        audio = f.read(f.frames)
        sr = f.sample_rate
    board = Pedalboard([
        HighpassFilter(cutoff_frequency_hz=35),
        LowShelfFilter(cutoff_frequency_hz=150, gain_db=bass_boost_db),
        Compressor(threshold_db=-12, ratio=4),
        Gain(gain_db=2)
    ])
    mastered = board(audio, sr)
    with AudioFile(temp_mastered, 'w', sr, mastered.shape[0]) as f:
        f.write(mastered)

    # Pydub Fades
    seg = AudioSegment.from_wav(temp_mastered)
    fade_ms = int(fade_sec * 1000)
    seg.fade_in(fade_ms).fade_out(fade_ms).export("stage.wav", format="wav")

    # BASH EXPORT (FFmpeg Insane Quality)
    os.system("ffmpeg -y -i stage.wav -codec:a libmp3lame -qscale:a 0 studio_master.mp3")
    return "studio_master.mp3"

def generate_music(prompt, duration, instruments, energy, bass_boost_db, fade_sec):
    if not prompt: return None, None, "Please enter a style!"
    
    full_prompt = f"{prompt}, {', '.join(instruments)}, {energy} energy, high quality."
    inputs = processor(text=[full_prompt], padding=True, return_tensors="pt").to(device)
    
    with torch.no_grad():
        audio_values = model.generate(**inputs, max_new_tokens=int(duration * 50), do_sample=True)
    
    sampling_rate = model.config.audio_encoder.sampling_rate
    audio_data = audio_values[0, 0].cpu().numpy()
    
    final_mp3 = apply_audacity_fixes(sampling_rate, audio_data, bass_boost_db, fade_sec)
    cert_file, cert_text = create_license(prompt, instruments)
    
    return final_mp3, cert_file, cert_text

# 3. UI
with gr.Blocks(theme=gr.themes.Soft(primary_hue="emerald")) as demo:
    gr.HTML("<div style='text-align:center;'><h1>🎵 COMMERICAL NEURAL STUDIO</h1></div>")
    with gr.Row():
        with gr.Column():
            txt = gr.Textbox(label="Music Style")
            ins = gr.CheckboxGroup(["Piano", "Drums", "Synth", "Guitar"], value=["Piano"], label="Instruments")
            en = gr.Radio(["Low", "Medium", "High"], value="Medium", label="Energy")
            dur = gr.Slider(5, 30, value=10, label="Seconds")
            with gr.Accordion("Mastering Options", open=False):
                bass = gr.Slider(0, 10, value=3, label="Bass Boost")
                fade = gr.Slider(0, 5, value=2, label="Fade")
            btn = gr.Button("🚀 GENERATE & LICENSE", variant="primary")
        with gr.Column():
            aud = gr.Audio(label="Studio Master MP3", type="filepath")
            cert = gr.File(label="Download Commercial Certificate")
            log = gr.Textbox(label="License Preview", lines=6)

    btn.click(generate_music, [txt, dur, ins, en, bass, fade], [aud, cert, log])

if __name__ == "__main__":
    demo.queue().launch()
    # Add this at the bottom of your Gradio UI code
gr.Markdown("""
### ⚖️ Legal Notice
Songs generated here use the MusicGen-Small weights (CC-BY-NC 4.0). 
Content is intended for personal use, education, and creative exploration. 
Commercial licensing for AI music is an evolving legal field; please consult 
local copyright laws before commercial distribution.
""")