"""SampleFlip — Upload a sample, get a beat."""

import os
import sys
import json
import tempfile

# Download drum kits from HF dataset repo on startup
SPACE_DIR = os.path.dirname(os.path.abspath(__file__))
KITS_DIR = os.path.join(tempfile.gettempdir(), 'sampleflip_kits')

if not os.path.exists(KITS_DIR) or len(os.listdir(KITS_DIR)) < 5:
    print('Downloading drum kits from HF dataset...')
    from huggingface_hub import snapshot_download
    KITS_DIR = snapshot_download(
        'ronantakizawa/sampleflip-kits',
        repo_type='dataset',
        local_dir=KITS_DIR,
    )
    print(f'Drum kits ready: {KITS_DIR}')

os.environ['SAMPLEFLIP_KIT_DIR'] = KITS_DIR
os.environ['SAMPLEFLIP_OUTPUT_DIR'] = tempfile.mkdtemp(prefix='sampleflip_out_')

# Add source to path
sys.path.insert(0, os.path.join(SPACE_DIR, 'src'))

import gradio as gr

GENRES = [
    'trap', 'boombap', 'jazzhouse', 'progressive_house', 'rnb',
    'drill', 'melodic_trap', '2hollis', 'techno', 'breakcore',
]


def generate_beat(audio_file, prompt, genre_override, bpm_override):
    """Generate a beat from an uploaded sample."""
    if audio_file is None:
        raise gr.Error("Please upload an audio file (WAV or MP3).")

    from sampleflip.agent import plan_beat, generate_drums, pick_bass_pattern
    from sampleflip.core.render_beat import GENRE_CONFIGS

    logs = []
    def log(msg):
        logs.append(msg)
        print(msg)

    # Step 1: LLM plans genre/bpm/name from prompt (or use defaults)
    if prompt and prompt.strip():
        log("Planning beat from description...")
        genre_arg = genre_override if genre_override != "auto" else None
        bpm_arg = bpm_override if bpm_override and bpm_override > 0 else None
        try:
            plan = plan_beat(prompt, genre_override=genre_arg, bpm_override=bpm_arg)
            g = plan['genre']
            b = plan['bpm']
            name = plan['name']
        except Exception as e:
            log(f"  Planning failed ({e}), using defaults")
            g = genre_override if genre_override != "auto" else "trap"
            b = bpm_override if bpm_override and bpm_override > 0 else None
            name = "Beat"
    else:
        g = genre_override if genre_override != "auto" else "trap"
        b = bpm_override if bpm_override and bpm_override > 0 else None
        name = "Beat"

    log(f"  Genre: {g} | BPM: {b or 'auto'} | Name: {name}")

    # Step 2: Bass pattern
    log("Designing bass pattern...")
    try:
        bass_pat = pick_bass_pattern(prompt or f"{g} beat", g)
        os.environ['SAMPLEFLIP_BASS_PATTERN'] = bass_pat
        log(f"  Bass: {bass_pat}")
    except Exception:
        pass

    # Step 3: Drum pattern
    cfg = GENRE_CONFIGS[g]
    nbars = cfg['bars']
    arrangement = cfg['arrangement']

    log("Generating drum pattern...")
    drums_json = None
    try:
        drum_data = generate_drums(prompt or f"{g} beat", g, b or cfg['bpm_default'], nbars, arrangement)
        n_pats = len([k for k in drum_data['patterns'] if k != 'silent'])
        log(f"  Created {n_pats} patterns for {nbars} bars")
        drums_json = os.path.join(tempfile.gettempdir(), f'sf_drums_{name}.json')
        with open(drums_json, 'w') as f:
            json.dump(drum_data, f)
    except Exception as e:
        log(f"  Drum generation failed ({e}), using defaults")

    # Step 4: Render
    log(f"Generating {g} beat...")
    output_dir = os.environ['SAMPLEFLIP_OUTPUT_DIR']

    core_dir = os.path.join(SPACE_DIR, 'src', 'sampleflip', 'core')
    if core_dir not in sys.path:
        sys.path.insert(0, core_dir)

    from sampleflip.core.render_beat import render
    render(
        audio_file, name, genre=g,
        bpm_hint=b, nbars=None,
        loop_start=None, loop_end=None,
        lofi=None, no_bass=False,
        vinyl_slow=False, drums_json=drums_json,
    )

    # Find output MP3
    mp3_files = [f for f in os.listdir(output_dir) if f.endswith('.mp3') and name in f]
    if not mp3_files:
        raise gr.Error("Render completed but no MP3 found.")

    mp3_path = os.path.join(output_dir, sorted(mp3_files)[-1])
    log(f"Done! {mp3_path}")

    return mp3_path, "\n".join(logs)


# Gradio UI
with gr.Blocks(title="SampleFlip", theme=gr.themes.Base()) as demo:
    gr.Markdown("# SampleFlip")
    gr.Markdown("Upload a sample, get a full produced beat with drums, bass, arrangement, and mastering.")

    with gr.Row():
        with gr.Column(scale=2):
            audio_input = gr.Audio(
                label="Upload your sample (WAV or MP3)",
                type="filepath",
            )
            prompt = gr.Textbox(
                label="Describe the beat (optional)",
                placeholder='e.g. "dark trap beat" or "jazzy house groove"',
                lines=1,
            )
        with gr.Column(scale=1):
            genre = gr.Dropdown(
                choices=["auto"] + GENRES,
                value="auto",
                label="Genre",
            )
            bpm = gr.Slider(
                minimum=0, maximum=200, step=1, value=0,
                label="BPM (0 = auto)",
            )

    generate_btn = gr.Button("Generate Beat", variant="primary", size="lg")

    with gr.Row():
        audio_out = gr.Audio(label="Your Beat", type="filepath")

    logs_out = gr.Textbox(label="Generation Log", lines=10, interactive=False)

    generate_btn.click(
        fn=generate_beat,
        inputs=[audio_input, prompt, genre, bpm],
        outputs=[audio_out, logs_out],
        api_name="generate_beat",
    )

    gr.Markdown("""
### How to use
1. **Upload** a melody loop, chord progression, or any audio sample (WAV/MP3)
2. **Describe** what kind of beat you want (optional — helps the AI pick drums and bass)
3. **Select** a genre and BPM (or leave on auto)
4. **Click Generate** — AI adds drums, bass, arrangement, FX, and masters to -14 LUFS

**Find samples:** Search YouTube for "free melody loop" or "sample pack preview", download with [yt-dlp](https://github.com/yt-dlp/yt-dlp), and upload here.
""")

demo.launch(server_name="0.0.0.0", server_port=7860)