import os
os.environ["PIP_ONLY_BINARY"] = ":all:"

import gradio as gr
from audiocraft.models import MusicGen
from huggingface_hub import hf_hub_download
import torch
import numpy as np
from scipy.io.wavfile import write

# Load the model once at startup to save time during inference
def load_model():
    # 1. Load the base architecture
    model = MusicGen.get_pretrained('facebook/musicgen-small')
    
    # 2. Download and load YOUR weights from your repo
    try:
        weights_path = hf_hub_download(repo_id="Vclord/anime-musicgen-weights", filename="pytorch_model.bin")
        state_dict = torch.load(weights_path, map_location="cpu")
        model.lm.load_state_dict(state_dict)
    except Exception as e:
        print(f"Loading custom weights failed, using base model. Error: {e}")
        
    return model

model = load_model()

def generate_music(description, duration, temperature, history):
    # Limit duration for CPU performance (10 seconds max)
    model.set_generation_params(duration=10,
                               temperature=temperature,
                               top_k=250) 
    
    # Generate on CPU
    wav = model.generate([description])
    
    # Convert to expected Gradio format: (sample_rate, numpy_array)
    audio_data = wav[0, 0].cpu().numpy().astype(np.float32)
    sr = 32000

    # Save to a temporary file for the Download button
    file_path = "generated_anime_theme.wav"
    write(file_path, sr, audio_data)

    # 5. Update History (Prompt, Audio Array)
    new_entry = [description, (sr, audio_data)]
    history.insert(0, new_entry)
    
    return (sr, audio_data), file_path, history

# --- Define Example Prompts ---
example_prompts = [
    ["High-energy Shonen opening, fast electric guitars, driving drums, 180 BPM"],
    ["Melancholic Shojo ending theme, soft piano, emotional strings, lo-fi aesthetic"],
    ["Epic orchestral battle theme, heavy brass, cinematic percussion, dark choir"],
    ["Cyberpunk synthwave anime style, neon atmosphere, pulsating bass, 120 BPM"],
    ["Slice of life background music, acoustic guitar, cheerful and lighthearted"]
]

# --- Midnight Neon CSS ---
css = """
/* 1. Force the browser and container to stay Dark regardless of Theme */
html { color-scheme: dark !important; }

.gradio-container { 
    background-color: #0b0f19 !important; 
    color: #ffffff !important;
}

/* 2. EXAMPLES TEXT VISIBILITY FIX */
/* Target every possible container for the example prompts */
.gr-examples, .gr-examples table, .gr-examples table td, .gr-examples button {
    background-color: #1e293b !important;
    color: #ffffff !important;
    opacity: 1 !important;
    border-color: #334155 !important;
}

/* Force the text inside the cells to be white and visible */
.gr-examples table td div, 
.gr-examples table td span,
.gr-examples .table-cell {
    color: #ffffff !important;
    background-color: transparent !important;
}

/* 3. HOVER EFFECT */
.gr-examples table td:hover, .gr-examples button:hover {
    background-color: #334155 !important;
    color: #2dd4bf !important;
}

/* 4. MARKDOWN & LABELS */
.prose h1, .prose h2, .prose h3, .prose p, .prose span, .prose strong { 
    color: #ffffff !important; 
}

.markdown-notice {
    background: rgba(45, 212, 191, 0.1) !important;
    border-left: 4px solid #2dd4bf !important;
    padding: 10px;
    color: #ffffff !important;
}

/* 5. BUTTONS & INPUTS */
.main-btn { 
    background: linear-gradient(135deg, #2dd4bf, #8b5cf6) !important; 
    color: white !important; 
}
.style-btn { 
    background-color: #1e293b !important; 
    border: 1px solid #334155 !important; 
    color: #2dd4bf !important; 
}
label span { color: #2dd4bf !important; font-weight: 600 !important; }

/* 6. SESSION HISTORY VISIBILITY */
.gr-dataset { background-color: #1e293b !important; border: 1px solid #334155 !important; }
.gr-dataset table td { color: #ffffff !important; }
"""

with gr.Blocks(css=css, theme=gr.themes.Default(primary_hue="teal", secondary_hue="slate")) as demo:
    gr.Markdown("# 🌌 Anime MusicGen: Midnight Studio")
    gr.Markdown("Fine-tuned for anime openings, background tracks, and character themes.")
    gr.Markdown("### Note: The track generation takes about 1 - 2 minutes, please do not close the window before that.", elem_classes="markdown-notice")
    
    history_state = gr.State([])

    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="Musical Prompt", 
                placeholder="Describe the vibe...",
                lines=3
            )

            gr.Examples(
                examples=example_prompts,
                inputs=[prompt_input],
                label="Quick Start Examples (Click to select)"
            )
            
            with gr.Accordion("🛠️ Advanced Settings", open=False):
                duration_slider = gr.Slider(2, 10, value=10, step=1, label="Duration (sec)")
                temp_slider = gr.Slider(0.1, 1.5, value=1.0, step=0.1, label="Creativity (Temp)")
            
            gr.Markdown("### ✨ Quick Style Tags")
            with gr.Row():
                s1 = gr.Button("🎹 J-Pop", size="sm", elem_classes="style-btn")
                s2 = gr.Button("🎸 Rock", size="sm", elem_classes="style-btn")
                s3 = gr.Button("🎻 Epic", size="sm", elem_classes="style-btn")
            
            generate_btn = gr.Button("Compose Theme", variant="primary", elem_classes="main-btn")
        
        with gr.Column(scale=1):
            output_audio = gr.Audio(label="Latest Preview", type="numpy", interactive=False)
            download_file = gr.File(label="Download High-Quality .WAV", interactive=False)

    # Style Tag Logic
    def add_style(current_prompt, style):
        return f"{current_prompt} {style}".strip()
    
    s1.click(fn=add_style, inputs=[prompt_input, gr.State("upbeat J-pop")], outputs=prompt_input)
    s2.click(fn=add_style, inputs=[prompt_input, gr.State("distorted electric guitar rock opening")], outputs=prompt_input)
    s3.click(fn=add_style, inputs=[prompt_input, gr.State("cinematic orchestral strings")], outputs=prompt_input)

    # History Display
    gr.Markdown("### 🕒 Session History")
    history_display = gr.Dataset(
        components=[gr.Textbox(visible=False), gr.Audio(visible=False)],
        label="Select a previous track to reload",
        samples=[],
        type="values"
    )

    # Generation Trigger
    generate_btn.click(
        fn=generate_music,
        inputs=[prompt_input, duration_slider, temp_slider, history_state],
        outputs=[output_audio, download_file, history_state]
    ).then(
        fn=lambda h: gr.update(samples=h),
        inputs=[history_state],
        outputs=[history_display]
    )

    # Clicking history reloads the audio preview
    def reload_audio(selection):
        return selection[1]

    history_display.click(
        fn=reload_audio,
        inputs=[history_display],
        outputs=[output_audio]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)