import gradio as gr
import os
import tempfile
import numpy as np
import soundfile as sf
from analyze import fingerprint, find_key, get_tempo, get_duration

try:
    import spaces
except ImportError:
    class _FakeSpaces:
        def GPU(self, fn=None, **kw):
            if fn: return fn
            return lambda f: f
    spaces = _FakeSpaces()


CUSTOM_CSS = """
/* ---- base ---- */
.gradio-container {
    background: #1a1714 !important;
    font-family: 'Georgia', 'Times New Roman', serif !important;
    max-width: 760px !important;
    margin: 0 auto !important;
}
.app-header {
    text-align: center;
    padding: 40px 20px 10px 20px;
}
.app-header h1 {
    color: #f0ece4 !important;
    font-size: 3.2em !important;
    font-weight: 400 !important;
    letter-spacing: 0.15em !important;
    margin-bottom: 4px !important;
    font-family: 'Georgia', serif !important;
}
.app-header p {
    color: #8a7e6e !important;
    font-size: 1em !important;
    font-style: italic !important;
    margin-top: 0 !important;
}
.upload-panel, .results-panel, .continue-panel {
    background: #211e19 !important;
    border: 1px solid #3a332a !important;
    border-radius: 12px !important;
    padding: 24px !important;
    margin-bottom: 16px !important;
}
.gradio-container label, .gradio-container .label-wrap span {
    color: #c8956c !important;
    font-family: 'Georgia', serif !important;
    font-size: 0.95em !important;
}
.gradio-container .prose, .gradio-container p, .gradio-container span {
    color: #f0ece4 !important;
}
.gradio-container .upload-button {
    background: #2a2520 !important;
    border: 2px dashed #4a4035 !important;
    color: #c8956c !important;
    border-radius: 8px !important;
}
.gradio-container .upload-button:hover {
    border-color: #c8956c !important;
    background: #302a22 !important;
}
.stat-card {
    background: #2a2520 !important;
    border: 1px solid #3a332a !important;
    border-radius: 10px !important;
    padding: 20px !important;
    text-align: center !important;
}
.stat-label {
    color: #8a7e6e !important;
    font-size: 0.8em !important;
    text-transform: uppercase !important;
    letter-spacing: 0.12em !important;
    margin-bottom: 6px !important;
}
.stat-value {
    color: #c8956c !important;
    font-size: 1.8em !important;
    font-weight: 400 !important;
    font-family: 'Georgia', serif !important;
}
.gradio-container button.primary {
    background: #c8956c !important;
    color: #1a1714 !important;
    border: none !important;
    border-radius: 8px !important;
    font-family: 'Georgia', serif !important;
    font-size: 1em !important;
    padding: 12px 32px !important;
    letter-spacing: 0.06em !important;
    transition: all 0.2s ease !important;
}
.gradio-container button.primary:hover {
    background: #d4a57c !important;
}
.gradio-container button.secondary {
    background: transparent !important;
    color: #c8956c !important;
    border: 1px solid #4a4035 !important;
    border-radius: 8px !important;
    font-family: 'Georgia', serif !important;
}
.gradio-container button.secondary:hover {
    border-color: #c8956c !important;
}
footer { display: none !important; }
.gradio-container .tab-nav button {
    color: #8a7e6e !important;
    border: none !important;
    background: transparent !important;
    font-family: 'Georgia', serif !important;
}
.gradio-container .tab-nav button.selected {
    color: #c8956c !important;
    border-bottom: 2px solid #c8956c !important;
}
.gradio-container input, .gradio-container textarea {
    background: #2a2520 !important;
    color: #f0ece4 !important;
    border-color: #3a332a !important;
}
.gradio-container .audio-player {
    background: #2a2520 !important;
    border-radius: 8px !important;
}
.gradio-container .progress-bar {
    background: #c8956c !important;
}
.tape-deco {
    text-align: center;
    padding: 8px 0;
    color: #3a332a;
    font-size: 0.85em;
    letter-spacing: 0.3em;
}
"""


def make_stat_html(label, value):
    return (
        '<div class="stat-card">'
        '<div class="stat-label">' + label + '</div>'
        '<div class="stat-value">' + str(value) + '</div>'
        '</div>'
    )


def analyze_track(audio):
    if audio is None:
        return (
            make_stat_html("key", "---"),
            make_stat_html("tempo", "---"),
            make_stat_html("duration", "---"),
            gr.update(visible=False),
            None,
            None
        )

    try:
        info = fingerprint(audio)
        key_html = make_stat_html("key", info["key"])
        bpm_html = make_stat_html("tempo", str(info["bpm"]) + " bpm")
        dur_html = make_stat_html("duration", str(info["duration"]) + "s")
        return (key_html, bpm_html, dur_html, gr.update(visible=True), audio, info)
    except Exception as e:
        err = make_stat_html("error", str(e)[:50])
        return (err, make_stat_html("tempo", "---"), make_stat_html("duration", "---"), gr.update(visible=False), None, None)


@spaces.GPU(duration=180)
def run_continuation(audio_path, info):
    if not audio_path:
        return None, "upload a track first."

    try:
        # lazy import so the UI boots fast and torch only loads on demand
        import librosa
        from continue_music import continue_track, stitch_with_crossfade

        key = info.get("key") if info else None
        bpm = info.get("bpm") if info else None

        continuation, sr = continue_track(audio_path, key=key, bpm=bpm)
        original, _ = librosa.load(audio_path, sr=sr, mono=True)
        full = stitch_with_crossfade(original, continuation, sr)

        out_path = os.path.join(tempfile.mkdtemp(), "coda_continuation.wav")
        sf.write(out_path, full, sr)

        added = len(continuation) / sr
        return out_path, f"added {added:.1f}s. seam crossfaded at {len(original)/sr:.1f}s."
    except Exception as e:
        return None, f"continuation failed: {str(e)[:200]}"


HEADER_HTML = (
    '<div class="app-header">'
    '<h1>CODA</h1>'
    '<p>upload an unfinished song. it picks up where you left off.</p>'
    '</div>'
    '<div class="tape-deco">- - - - - - - - - - - -</div>'
)

FOOTER_HTML = (
    '<div class="tape-deco" style="margin-top: 20px;">- - - - - - - - - - - -</div>'
    '<div style="text-align:center; padding:16px 0; color:#4a4035; font-size:0.8em; font-family:Georgia,serif;">'
    'CODA // 100% local // no cloud APIs // built for Build Small 2025'
    '</div>'
)


with gr.Blocks(title="CODA", css=CUSTOM_CSS, theme=gr.themes.Base()) as app:

    current_track = gr.State(None)
    track_info = gr.State(None)

    gr.HTML(HEADER_HTML)

    with gr.Group(elem_classes="upload-panel"):
        audio_input = gr.Audio(
            label="drop a track",
            type="filepath",
            sources=["upload", "microphone"],
        )

    with gr.Group(elem_classes="results-panel"):
        gr.HTML('<div class="tape-deco" style="margin-bottom:12px;">analysis</div>')
        with gr.Row():
            key_display = gr.HTML(make_stat_html("key", "---"))
            bpm_display = gr.HTML(make_stat_html("tempo", "---"))
            dur_display = gr.HTML(make_stat_html("duration", "---"))

    with gr.Group(elem_classes="continue-panel", visible=False) as continue_section:
        gr.HTML('<div class="tape-deco" style="margin-bottom:12px;">continue</div>')
        continue_btn = gr.Button("continue this track", variant="primary")
        result_audio = gr.Audio(label="finished track", type="filepath", interactive=False)
        continue_output = gr.Textbox(label="status", interactive=False, lines=2)

    audio_input.change(
        fn=analyze_track,
        inputs=[audio_input],
        outputs=[key_display, bpm_display, dur_display, continue_section, current_track, track_info]
    )

    continue_btn.click(
        fn=run_continuation,
        inputs=[current_track, track_info],
        outputs=[result_audio, continue_output]
    )

    gr.HTML(FOOTER_HTML)


if __name__ == "__main__":
    app.launch()