"""CODA — finish the song you quit on.

Upload a short, unfinished music clip. CODA listens to it (key, tempo, meter),
then continues it into a longer, finished-sounding track in the same feel using
Stable Audio 3 Small Music — a single native audio-continuation call, 44.1 kHz
stereo — and splices the new part seamlessly onto your pristine original with a
level-matched crossfade and a clean closing fade.

Deliberately one job, done well. No lyrics generator, no cover-art printer — just
a real, listenable continuation of your clip.
"""
# stable-audio-tools hard-pins torch==2.7.1 in its pyproject.toml; ZeroGPU only
# accepts torch 2.8.0/2.9.1/2.10.0/2.11.0. We install the package without its
# dependency tree so the ZeroGPU-managed torch in the env is used instead.
import subprocess as _sp, sys as _sys
try:
    import stable_audio_tools as _  # noqa: F401
    del _
except ImportError:
    _sp.check_call([_sys.executable, "-m", "pip", "install",
                    "--no-cache-dir", "--no-deps", "-q",
                    "git+https://github.com/Stability-AI/stable-audio-tools.git"])
del _sp, _sys

import os
import tempfile

import gradio as gr
import librosa
import numpy as np
import soundfile as sf

import engine
import stitch
from analyze import fingerprint
from enhance import enhance_audio, enhance_to_tempfile, input_quality

# ZeroGPU shim: on a Space `spaces` exists and `@spaces.GPU` attaches a GPU for
# the duration of the call. Locally we no-op so the app still runs.
try:
    import spaces
except ImportError:
    class _FakeSpaces:
        def GPU(self, fn=None, **kw):
            return fn if fn else (lambda f: f)
    spaces = _FakeSpaces()

# bundled demo: a low-quality phone capture of an unfinished song by the band
# PUSHBACK (shared via TikTok). lo-fi in, finished-sounding out.
PUSHBACK_DEMO = os.path.join(os.path.dirname(__file__),
                             "examples", "pushback_demo.mp3")
PUSHBACK_CREDIT = ("Demo clip: **PUSHBACK** (via TikTok), used with thanks. "
                   "Bring your own clip to finish your own song.")

# total finished length. SA3 Small generates up to 120 s in one call, so this is
# a *total length* control (clip + continuation), not a "seconds to add" knob.
MIN_TOTAL, MAX_TOTAL, DEFAULT_TOTAL = 30, 120, 60

# SA3's weights are gated under the Stability Community License, so the download
# 401s unless the request is authenticated. Log in with the HF_TOKEN secret
# (which must come from an account that accepted the licence on the model page)
# BEFORE preload() pulls the weights. No-op when unset, so local runs that have
# already done `huggingface-cli login` still work.
_HF_TOKEN = os.environ.get("HF_TOKEN")
if _HF_TOKEN:
    try:
        from huggingface_hub import login as _hf_login
        _hf_login(token=_HF_TOKEN)
        print("[coda] authenticated to the HF Hub via HF_TOKEN", flush=True)
    except Exception as _e:
        print(f"[coda] HF login failed ({_e}); gated model download may 401",
              flush=True)
elif os.environ.get("SPACE_ID"):
    print("[coda] WARNING: no HF_TOKEN secret set — the gated SA3 weights will "
          "401. Add HF_TOKEN in the Space settings (Settings -> Variables and "
          "secrets) using a token from an account that accepted the licence.",
          flush=True)

# On a Space, pull weights into CPU RAM at boot so the GPU window is spent
# generating, not reading 3 GB off disk. `spaces` defers CUDA placement until
# the first @spaces.GPU call.
if os.environ.get("SPACE_ID"):
    try:
        engine.preload()
    except Exception as _e:
        print(f"[coda] preload failed ({_e}); will lazy-load", flush=True)


def _fmt_info(info, quality):
    """human-readable summary of what CODA heard."""
    lines = [
        f"**KEY**  `{info['key']}`",
        f"**TEMPO**  `{info['bpm']} BPM`",
        f"**METER**  `{info['time_signature']}`",
        f"**CLIP**  `{info['duration']}s`",
    ]
    if quality and quality.get("lofi"):
        lines.append(
            f"**SOURCE**  `lo-fi ~{quality['bandwidth_hz']/1000:.0f}kHz` "
            f"— CODA cleans a copy before it listens, so it follows the *song*, "
            f"not the hiss")
    return "  \n".join(lines)


def analyze_on_upload(audio_path):
    """Fast CPU-only pass the instant a clip loads, so the user sees what CODA
    heard immediately instead of a dead screen."""
    if not audio_path:
        return gr.update(value="", visible=False), gr.update(interactive=False)
    try:
        listen_path = enhance_to_tempfile(audio_path)
        info = fingerprint(listen_path)
        quality = input_quality(audio_path)
        # SA3 continues up to a 120s total, so a clip needs headroom for at
        # least MIN_NEW seconds of new audio. Block over-long clips here, with a
        # clear message, instead of failing at generation time.
        if info["duration"] > engine.MAX_SOURCE_SECONDS:
            msg = (f"### Clip too long\nThat clip is **{info['duration']:.0f}s**. "
                   f"CODA continues clips up to **{engine.MAX_SOURCE_SECONDS:.0f}s** "
                   f"(Stable Audio 3's {engine.MAX_TOTAL_SECONDS:.0f}s total cap). "
                   f"Trim it shorter and re-upload.")
            return gr.update(value=msg, visible=True), gr.update(interactive=False)
        md = "### CODA heard\n" + _fmt_info(info, quality)
        return gr.update(value=md, visible=True), gr.update(interactive=True)
    except Exception as e:
        print(f"[coda] analysis failed ({e})", flush=True)
        return (gr.update(value=f"Couldn't read that file: {e}", visible=True),
                gr.update(interactive=False))


@spaces.GPU(duration=120)
def _continue_on_gpu(listen_path, total_seconds, vibe):
    """ONLY the SA3 diffusion call runs inside the GPU window. All CPU work —
    enhancement, key/tempo analysis, decode, splice — happens OUTSIDE it (in
    `finish_song`), so the scarce ZeroGPU allocation is spent generating instead
    of decoding/analyzing audio. That's the stream-abort fix: the GPU task now
    starts and finishes fast instead of sitting through a slow analysis preamble
    until the browser aborts the stream.

    Pin the seed to the lab's known-good draw. SA3 is generative: a random seed
    gives a different (often weaker) continuation every run. seed=7 produced the
    verified lab_out/sa3 takes, so the app reproduces that instead of re-rolling.
    """
    return engine.continue_audio(
        listen_path, total_seconds=int(total_seconds),
        prompt=(vibe or "").strip(), seed=7)


def finish_song(audio_path, total_seconds, vibe, remaster,
                progress=gr.Progress()):
    """Orchestrate the job: CPU prep -> GPU continuation -> CPU splice.
    Returns (output_wav_path, summary_markdown)."""
    if not audio_path:
        raise gr.Error("Upload a clip (or load the PUSHBACK demo) first.")

    total_seconds = int(total_seconds)

    # --- CPU prep (outside the GPU window) ---
    progress(0.05, desc="Listening to your clip…")
    listen_path = enhance_to_tempfile(audio_path)
    info = fingerprint(listen_path)
    # the pristine original — what the listener hears for the first stretch
    original, sr = librosa.load(audio_path, sr=None, mono=False)
    if remaster:
        progress(0.15, desc="Remastering your part…")
        original = enhance_audio(original, sr)

    # --- GPU continuation (the ONLY @spaces.GPU call) ---
    progress(0.35, desc="Composing the continuation…")
    try:
        new_tail, source_seconds, SR = _continue_on_gpu(
            listen_path, total_seconds, vibe)
    except ValueError as e:
        # e.g. the clip is a full-length track, not a clip to continue
        raise gr.Error(str(e))

    # --- CPU splice + write (outside the GPU window) ---
    progress(0.9, desc="Splicing onto your original…")
    out = stitch.stitch(original, sr, new_tail, source_seconds)

    out_path = os.path.join(tempfile.mkdtemp(), "coda_finished.wav")
    sf.write(out_path, out.T, SR, subtype="PCM_16")

    progress(1.0, desc="Done.")
    total = out.shape[-1] / SR
    added = total - source_seconds
    vibe_note = f" guided by *“{vibe.strip()}”*" if (vibe or "").strip() else ""
    summary = (
        f"### Finished — {total:.0f}s\n"
        f"Your **{source_seconds:.0f}s** clip in **{info['key']}** at "
        f"**{info['bpm']} BPM** continued for **~{added:.0f}s** more{vibe_note}, "
        f"then crossfaded onto your original and faded to a clean close.\n\n"
        f"*Stable Audio 3 generated the continuation as 44.1 kHz stereo in a "
        f"single pass; your original recording plays untouched up to the seam.*"
    )
    return out_path, summary


def load_demo():
    """Load the PUSHBACK demo clip into the uploader."""
    return PUSHBACK_DEMO


# --- dark "DAW console" theme -------------------------------------------------
THEME = gr.themes.Base(
    primary_hue=gr.themes.colors.cyan,
    secondary_hue=gr.themes.colors.orange,
    neutral_hue=gr.themes.colors.slate,
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui"],
    font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
).set(
    body_background_fill="#0b0d10",
    body_background_fill_dark="#0b0d10",
    background_fill_primary="#14171c",
    background_fill_secondary="#14171c",
    block_background_fill="#14171c",
    block_border_color="#222831",
    block_label_background_fill="#14171c",
    block_title_text_color="#9aa4b2",
    body_text_color="#e6e9ef",
    body_text_color_subdued="#9aa4b2",
    button_primary_background_fill="#39d0d8",
    button_primary_background_fill_hover="#4fe0e8",
    button_primary_text_color="#06181a",
    color_accent_soft="#1a2026",
    border_color_accent="#39d0d8",
    input_background_fill="#0e1115",
    slider_color="#ffb347",
)

CSS = """
:root { --coda-accent:#39d0d8; --coda-amber:#ffb347; }
.gradio-container { max-width: 1060px !important; margin: 0 auto !important; }
#coda-head { text-align:center; padding: 8px 0 2px; }
#coda-title { font-weight:800; letter-spacing:.5px; font-size:2.5rem; margin:0;
  background:linear-gradient(90deg,#39d0d8,#ffb347);
  -webkit-background-clip:text; background-clip:text; -webkit-text-fill-color:transparent; }
#coda-tagline { color:#9aa4b2; margin:.1rem 0 0; font-size:1.02rem; letter-spacing:.3px; }
#coda-rule { height:2px; border:0; margin:10px auto 16px; max-width:240px;
  background:linear-gradient(90deg,transparent,#39d0d8,#ffb347,transparent); opacity:.8; }
#coda-eq { display:flex; gap:4px; justify-content:center; align-items:flex-end; height:26px; margin-top:8px; }
#coda-eq span { width:5px; background:linear-gradient(180deg,#39d0d8,#1b6e72); border-radius:2px;
  animation: codaeq 1.1s ease-in-out infinite; }
#coda-eq span:nth-child(2){animation-delay:.15s} #coda-eq span:nth-child(3){animation-delay:.30s}
#coda-eq span:nth-child(4){animation-delay:.45s} #coda-eq span:nth-child(5){animation-delay:.6s}
#coda-eq span:nth-child(6){animation-delay:.3s} #coda-eq span:nth-child(7){animation-delay:.1s}
@keyframes codaeq { 0%,100%{height:7px;opacity:.55} 50%{height:24px;opacity:1} }
.coda-card { border:1px solid #222831; border-radius:14px; padding:16px 18px;
  background:#14171c; box-shadow: inset 0 1px 0 #1c222b; }
.coda-card h3 { color:var(--coda-accent); text-transform:uppercase; letter-spacing:1.5px;
  font-size:.8rem; margin:.1rem 0 .7rem; }
#coda-foot { text-align:center; color:#6b7480; font-size:.85rem; margin-top:14px; }
"""

EQ_BARS = "<div id='coda-eq'>" + "".join("<span></span>" for _ in range(7)) + "</div>"

with gr.Blocks(title="CODA") as app:
    with gr.Column(elem_id="coda-head"):
        gr.HTML("<h1 id='coda-title'>🎵 CODA</h1>"
                "<p id='coda-tagline'>the songs you quit on, finished.</p>"
                + EQ_BARS)
    gr.HTML("<hr id='coda-rule'/>")
    gr.Markdown(
        "Upload a short, unfinished music clip. CODA reads its key, tempo and "
        "groove, then **Stable Audio 3** continues it into a finished-sounding "
        "track in the same feel — 44.1 kHz stereo, spliced seamlessly onto your "
        "original.")

    with gr.Row(equal_height=False):
        with gr.Column(scale=1):
            audio_input = gr.Audio(
                label="Your unfinished clip", type="filepath", sources=["upload"])
            demo_btn = gr.Button("🎧 Try the demo — PUSHBACK (via TikTok)",
                                 size="sm")
            total_slider = gr.Slider(
                MIN_TOTAL, MAX_TOTAL, value=DEFAULT_TOTAL, step=1,
                label="Finished length (seconds)",
                info="Total length of the finished track. Longer = a bit slower.")
            vibe = gr.Textbox(
                label="Describe the vibe (optional)", lines=1,
                placeholder="e.g. warm lo-fi, vinyl crackle, mellow piano",
                info="Leave empty for pure audio-led continuation.")
            remaster = gr.Checkbox(
                value=False, label="Remaster my part too",
                info="Apply the same lo-fi cleanup to your original section so "
                     "the whole track sits at one level.")
            finish_btn = gr.Button("Finish this song", variant="primary",
                                   interactive=False)

        with gr.Column(scale=1):
            info_md = gr.Markdown(visible=False, elem_classes="coda-card")
            # built visible inside the result group: Gradio drops visible=False
            # audio players at build time, so we never toggle player visibility.
            with gr.Group(elem_classes="coda-card"):
                gr.Markdown("### Your finished song")
                output_audio = gr.Audio(label="", type="filepath",
                                        interactive=False)
                summary_md = gr.Markdown(visible=False)

    gr.Markdown(PUSHBACK_CREDIT, elem_id="coda-foot")

    audio_input.change(fn=analyze_on_upload, inputs=[audio_input],
                       outputs=[info_md, finish_btn])
    demo_btn.click(fn=load_demo, inputs=[], outputs=[audio_input])

    def _show_summary():
        return gr.update(visible=True)

    finish_btn.click(
        fn=finish_song,
        inputs=[audio_input, total_slider, vibe, remaster],
        outputs=[output_audio, summary_md]).then(
        fn=_show_summary, inputs=[], outputs=[summary_md])


if __name__ == "__main__":
    # Gradio 6 moved theme/css to launch(); pass them here (and they remain on
    # Blocks above) so the dark DAW theme applies however the Space serves it.
    app.queue().launch(theme=THEME, css=CSS)