"""Streamlit UI for Anticipatory Music Transformer accompaniment."""

from __future__ import annotations

import os
import sys
import warnings
from pathlib import Path

# Show tqdm-style download bars in the terminal where Streamlit runs
os.environ.setdefault("HF_HUB_ENABLE_PROGRESS_BARS", "1")

import streamlit as st

# transformers 4.29.x triggers huggingface_hub's deprecated resume_download warning
warnings.filterwarnings(
    "ignore",
    message=".*resume_download.*",
    category=FutureWarning,
    module="huggingface_hub.file_download",
)

_SRC = Path(__file__).resolve().parent
if str(_SRC) not in sys.path:
    sys.path.insert(0, str(_SRC))

from generation_service import (  # noqa: E402
    Mode,
    default_melody_programs,
    events_to_midi_bytes,
    events_to_musicxml_bytes,
    events_to_pdf_bytes,
    instrument_choices,
    parse_upload_to_events,
    prepare_events_window,
    run_accompaniment,
    validate_melody_vs_events,
)

# ~128M params — faster download & load (default). Medium is ~360M.
MODEL_FAST = "stanford-crfm/music-small-800k"
MODEL_QUALITY = "stanford-crfm/music-medium-800k"


def _safe_upload_stem(uploaded_file) -> str:
    if uploaded_file is None or not getattr(uploaded_file, "name", None):
        return "anticipation_output"
    stem = Path(str(uploaded_file.name)).name
    stem = Path(stem).stem
    if not stem or stem in (".", ".."):
        return "anticipation_output"
    return stem


@st.cache_resource
def load_model(model_id: str):
    import torch
    from transformers import AutoModelForCausalLM
    from transformers import logging as tr_logging

    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

    if torch.cuda.is_available():
        device = "cuda"
    elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        device = "mps"
    else:
        device = "cpu"

    extra: dict = {"torch_dtype": torch.float32}
    try:
        import accelerate  # noqa: F401

        extra["low_cpu_mem_usage"] = True
    except ImportError:
        pass
    # float32 everywhere matches Hugging Face CPU Spaces and avoids fp16 vs fp32 sampling drift.
    # CRFM checkpoints include unused keys (e.g. token_out_embeddings); harmless noise otherwise.
    _v = tr_logging.get_verbosity()
    tr_logging.set_verbosity_error()
    try:
        model = AutoModelForCausalLM.from_pretrained(model_id, **extra)
    finally:
        tr_logging.set_verbosity(_v)
    model = model.to(device)
    model.eval()
    return model, device


def _reset_outputs():
    for key in (
        "out_midi",
        "out_midi_err",
        "out_musicxml",
        "out_musicxml_err",
        "out_pdf",
        "out_pdf_err",
    ):
        st.session_state.pop(key, None)


st.set_page_config(page_title="Harmony from your melody", layout="wide")
st.title("Add harmony to your melody")
st.caption(
    "Upload a tune you wrote (sheet music or MIDI). Lock in your **melody**; the model **adds other parts**. "
    "Then download MIDI, sheet-music XML, or PDF."
)

st.subheader("1. Upload your file")
uploaded = st.file_uploader(
    "Choose one score file",
    type=["xml", "musicxml", "mxl", "mid", "midi"],
    help=(
        "From MuseScore, Finale, Dorico, etc.: export as MusicXML or MIDI. "
        "Select a **file**, not a project folder. Types: .xml, .musicxml, .mxl, .mid, .midi."
    ),
)

st.subheader("2. How much of the piece to use")
st.markdown(
    "The AI works on a **short slice** from the start of your file (long scores are trimmed). "
    "A smaller slice runs faster."
)
col_len_a, col_len_b = st.columns(2)
with col_len_a:
    clip_sec = st.slider(
        "Length of the excerpt (seconds)",
        min_value=5,
        max_value=120,
        value=20,
        step=1,
        help="Only this many seconds from the beginning of your file are loaded and harmonized.",
    )
with col_len_b:
    prompt_sec = st.slider(
        "How many seconds stay “exactly yours” first",
        min_value=0,
        max_value=int(clip_sec),
        value=min(5, int(clip_sec)),
        step=1,
        help=(
            "At the start of the excerpt, this many seconds are treated as fixed context; "
            "the model continues and adds harmony after that point within the same window."
        ),
    )

events = None
choices = []
parse_error = None
if uploaded is not None:
    try:
        raw = uploaded.getvalue()
        events = parse_upload_to_events(raw, uploaded.name)
        events = prepare_events_window(events, float(clip_sec))
        choices = instrument_choices(events)
    except Exception as e:
        parse_error = str(e)
        events = None

if parse_error:
    st.error(f"We could not read that file: {parse_error}")

if uploaded is None:
    st.info("Upload a file in step 1 to continue.")
    st.stop()

if not events:
    st.warning("Once the file loads correctly, you will choose your melody part below.")
    st.stop()

if not choices:
    st.error("This file has no notes we could read. Try exporting again as MusicXML or MIDI.")
    st.stop()

st.subheader("3. Which part is your melody?")
st.markdown(
    "Each line is an **instrument sound** (MIDI program) and how many notes it has. "
    "**Select the part(s) you want to keep exactly as written** (usually your melody). "
    "The AI adds **other instruments** you did not write. "
    "If the file is only one instrument, select that one—the model still creates backing parts."
)
options = [f"Instrument {prog} — {count} notes" for prog, count in choices]
prog_by_label = {options[i]: choices[i][0] for i in range(len(options))}
count_map = {p: c for p, c in choices}
defaults = default_melody_programs(choices)
default_labels = [f"Instrument {p} — {count_map[p]} notes" for p in defaults if p in count_map]

selected_labels = st.multiselect(
    "Melody (locked parts)",
    options=options,
    default=[lbl for lbl in default_labels if lbl in options],
    help="Locked parts stay note-for-note. New harmony uses other MIDI instruments (fine for a single-line melody).",
)

st.subheader("4. Model size")
st.caption(
    "Inference uses **float32** everywhere (same as CPU Spaces) so results line up across machines. "
    "First time only: the checkpoint downloads from Hugging Face (watch the **terminal** for a progress bar). "
    "After that, the same model stays cached until you restart the app or pick a different model."
)
model_tier = st.radio(
    "Choose one",
    options=["fast", "quality"],
    index=0,
    horizontal=True,
    format_func=lambda x: (
        "Faster — smaller download (~128M), good for trying the app"
        if x == "fast"
        else "Higher quality — larger download (~360M), slower on CPU"
    ),
)
preset_model_id = MODEL_FAST if model_tier == "fast" else MODEL_QUALITY

with st.expander("Advanced options (optional)"):
    st.markdown("You can usually leave these as they are.")
    custom_model_id = st.text_input(
        "Custom Hugging Face model id (optional)",
        value="",
        key="hf_model_override",
        placeholder="Leave empty to use the model size you chose above",
        help="Only if you use another compatible stanford-crfm checkpoint.",
    )
    mode: Mode = st.radio(
        "Generation style",
        options=["anticipatory", "autoregressive"],
        index=0,
        format_func=lambda x: "Standard — anticipatory (recommended)" if x == "anticipatory" else "Research baseline — plain autoregressive",
        help="Anticipatory mode matches the paper’s infilling setup. Autoregressive is a simpler baseline.",
    )
    top_p = st.slider(
        "How surprising / varied the harmony is",
        min_value=0.5,
        max_value=1.0,
        value=0.95,
        step=0.01,
        help="Higher values (e.g. 0.95–1.0) allow more diverse notes; lower values are more conservative.",
    )
    generation_seed = st.number_input(
        "Random seed",
        min_value=0,
        max_value=2**31 - 1,
        value=42,
        step=1,
        help=(
            "The model samples notes at random; this fixes the RNG so **the same file, settings, model, "
            "and seed** match between your laptop and Hugging Face. Change the seed for a new variation."
        ),
    )

model_id = custom_model_id.strip() if custom_model_id.strip() else preset_model_id

st.subheader("5. Generate")
gen = st.button("Generate harmony", type="primary")

if gen:
    _reset_outputs()
    melody_programs = [prog_by_label[lbl] for lbl in selected_labels]
    if not melody_programs:
        st.error("Choose at least one instrument as your melody.")
    else:
        try:
            validate_melody_vs_events(events, melody_programs)
        except ValueError as e:
            st.error(str(e))
        else:
            with st.status("Loading the AI model…", expanded=True) as load_status:
                load_status.write(
                    "**First run:** downloading can take **several minutes** on a slow connection "
                    "(hundreds of MB). Check the terminal for a download bar. "
                    "**Next runs** with the same model reuse the cache and are much quicker."
                )
                load_status.write(f"Using: `{model_id}`")
                try:
                    model, _device = load_model(model_id.strip())
                except Exception as e:
                    load_status.update(label="Model load failed", state="error")
                    st.error(f"Could not load the model: {e}")
                    model = None
                else:
                    load_status.update(label="Model ready", state="complete")
            if model is not None:
                with st.spinner("Writing harmony… Often **1–5+ minutes** on CPU; faster on GPU."):
                    try:
                        out = run_accompaniment(
                            model,
                            events,
                            melody_programs,
                            float(prompt_sec),
                            float(clip_sec),
                            top_p,
                            mode,
                            seed=int(generation_seed),
                        )
                    except Exception as e:
                        st.error(f"Generation stopped: {e}")
                    else:
                        st.success("Here is your new score. Download below.")
                        try:
                            st.session_state["out_midi"] = events_to_midi_bytes(out)
                            st.session_state.pop("out_midi_err", None)
                        except Exception as e:
                            st.session_state.pop("out_midi", None)
                            st.session_state["out_midi_err"] = str(e)
                        try:
                            st.session_state["out_musicxml"] = events_to_musicxml_bytes(out)
                            st.session_state.pop("out_musicxml_err", None)
                        except Exception as e:
                            st.session_state.pop("out_musicxml", None)
                            st.session_state["out_musicxml_err"] = str(e)
                        try:
                            st.session_state["out_pdf"] = events_to_pdf_bytes(out)
                            st.session_state.pop("out_pdf_err", None)
                        except Exception as e:
                            st.session_state["out_pdf"] = None
                            st.session_state["out_pdf_err"] = str(e)

st.divider()
st.subheader("6. Download")
base = _safe_upload_stem(uploaded)

if st.session_state.get("out_midi"):
    st.download_button(
        label="Download as MIDI (for DAWs, players)",
        data=st.session_state["out_midi"],
        file_name=f"{base}_with_harmony.mid",
        mime="audio/midi",
    )
if st.session_state.get("out_musicxml"):
    st.download_button(
        label="Download as MusicXML (open in notation software)",
        data=st.session_state["out_musicxml"],
        file_name=f"{base}_with_harmony.musicxml",
        mime="application/vnd.recordare.musicxml+xml",
    )
if st.session_state.get("out_pdf"):
    st.download_button(
        label="Download as PDF (printable score)",
        data=st.session_state["out_pdf"],
        file_name=f"{base}_with_harmony.pdf",
        mime="application/pdf",
    )
elif st.session_state.get("out_pdf_err"):
    st.caption(
        f"**PDF export:** {st.session_state['out_pdf_err']} "
        "On macOS, install LilyPond with: `brew install lilypond`, then restart Streamlit."
    )
if st.session_state.get("out_midi_err"):
    st.caption(f"MIDI export issue: {st.session_state['out_midi_err']}")
if st.session_state.get("out_musicxml_err"):
    st.caption(f"MusicXML export issue: {st.session_state['out_musicxml_err']}")

if not any(
    st.session_state.get(k)
    for k in ("out_midi", "out_musicxml", "out_pdf")
):
    st.caption("Run step 4 to create files you can download.")

st.divider()
with st.expander("About limits & credits"):
    st.markdown(
        """
- Very long or very dense scores may time out; use a shorter excerpt in step 2.
- Sheet music is converted through MIDI internally, so fine engraving may change slightly.
- **Anticipatory Music Transformer** (Apache-2.0): Thickstun et al.; inference code from
  [anticipation](https://github.com/jthickstun/anticipation). Weights from
  [Stanford CRFM on Hugging Face](https://huggingface.co/stanford-crfm).
"""
    )