"""SoftChart — custom Gradio Server app for Hugging Face Spaces.

Generate a Taiko no Tatsujin chart from any audio file, using the full system:
  - SoftChartGenerator  (main model, plan-conditioned)
  - SoftChartPlanner     (auto song-level planning)
  - SoftChartBeat        (beat/downbeat for barline anchoring)
All models load from the Hub via from_pretrained. MIT licensed.
"""

import logging
import os
import re
import shutil
import tempfile
from pathlib import Path

import gradio as gr
import numpy as np
import torch
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

from softchart.generate import generate_song, generate_song_slot, load_hf
from softchart.fonts import cjk_font_path
from softchart.grid import debias_to_grid, fit_grid_fixed_bpm, fit_grid_piecewise
from softchart.hf import SoftChartPlanner
from softchart.preview_audio import synthesize_taiko_preview
from softchart.rhythm import snap_chart
from softchart.tja import append_measure_with_gogo, gogo_measure_mask, write_tja_slots
from softchart.tja_image import render_tja_image
from softchart.vocab import FPS, HOP, N_FFT, N_MELS, SR

LOGGER = logging.getLogger("softchart.space")
STATIC_DIR = Path(__file__).with_name("static")

PLAN_REPO = os.environ.get("SC_PLAN", "JacobLinCool/softchart-planner")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Selectable generators: three generations x three sizes (plus the v1.5 single
# model). All are dual (slot + time), beat- and plan-conditioned, MIT licensed.
# https://huggingface.co/collections/JacobLinCool/softchart-generators
MODELS = {
    "v1.7": "JacobLinCool/softchart-v17",
    "v1.7-small": "JacobLinCool/softchart-v17-small",
    "v1.7-tiny": "JacobLinCool/softchart-v17-tiny",
    "v1.6": "JacobLinCool/softchart-v16",
    "v1.6-small": "JacobLinCool/softchart-v16-small",
    "v1.6-tiny": "JacobLinCool/softchart-v16-tiny",
    "v1.5": "JacobLinCool/softchart-v15",
}
DEFAULT_MODEL = os.environ.get("SC_MODEL", "v1.7")

COURSE_DENS = {"easy": 1, "normal": 2, "hard": 4, "oni": 7}
CHAR = {"don": "1", "ka": "2", "don_big": "3", "ka_big": "4",
        "roll": "5", "roll_big": "6", "balloon": "7"}
SUB = 96

_MODELS = {}  # repo id -> {"gen","slot","beat"}, loaded lazily and cached
_PLANNER = {}  # the planner is model-agnostic and shared across generators


def get_models(model_choice=DEFAULT_MODEL, *, include_planner=False):
    repo = MODELS.get(model_choice)
    if repo is None:
        raise ValueError(
            f"unknown model {model_choice!r}; choose one of {list(MODELS)}"
        )
    if repo not in _MODELS:
        u = load_hf(repo, device=DEVICE)
        if (not getattr(u, "_dual", False) or u.beat is None
                or not getattr(u, "_has_plan", False)):
            raise RuntimeError(
                f"{repo} must provide dual generation, beat, and plan conditioning"
            )
        _MODELS[repo] = {"gen": u, "slot": u, "beat": u}
    models = dict(_MODELS[repo])
    if include_planner:
        if "plan" not in _PLANNER:
            _PLANNER["plan"] = SoftChartPlanner.from_pretrained(PLAN_REPO).to(DEVICE).eval()
        models["plan"] = _PLANNER["plan"]
    return models


def load_logmel(path):
    import librosa

    wav, _ = librosa.load(path, sr=SR, mono=True)
    fb = librosa.filters.mel(sr=SR, n_fft=N_FFT, n_mels=N_MELS, fmin=20.0, fmax=SR / 2)
    spec = torch.stft(torch.from_numpy(wav), N_FFT, hop_length=HOP,
                      window=torch.hann_window(N_FFT), center=True, return_complex=True)
    mel = np.log(fb @ spec.abs().pow(2).numpy() + 1e-5).astype(np.float32)
    return mel, wav


def auto_plan(mel, bpm, downbeats=None):
    T = mel.shape[1]
    dur = T / FPS
    flux = np.concatenate([[0], np.maximum(0, np.diff(mel, axis=1)).sum(0)])
    beat = 60.0 / bpm
    edges = (list(downbeats[::4]) + [dur]) if (downbeats is not None and len(downbeats) >= 2) \
        else list(np.arange(0, dur, 4 * beat)) + [dur]
    vals = [float(flux[int(a * FPS):int(b * FPS)].mean()) if int(b * FPS) > int(a * FPS) else 0.0
            for a, b in zip(edges, edges[1:])]
    if not vals:
        return None
    vals = np.array(vals)
    lo, hi = np.percentile(vals, 15), np.percentile(vals, 92)
    peak = int(np.argmax(vals))
    plan = []
    for i, (a, b) in enumerate(zip(edges, edges[1:])):
        frac = (vals[i] - lo) / max(hi - lo, 1e-6)
        d8 = int(np.clip(round(frac * 7), 0, 7))
        fl = 1 if (vals[i] <= lo and 0 < i < len(vals) - 1) else (2 if i == peak and vals[i] >= hi else 0)
        plan.append([round(a, 3), round(b, 3), d8, fl])
    return plan


def learned_plan(planner, mel, course, bpm, downbeats=None):
    dur = mel.shape[1] / FPS
    beat = 60.0 / bpm
    edges = (list(downbeats[::4]) + [dur]) if (downbeats is not None and len(downbeats) >= 2) \
        else list(np.arange(0, dur, 4 * beat)) + [dur]
    feats, spans = [], []
    for a, b in zip(edges, edges[1:]):
        seg = mel[:, int(a * FPS):int(b * FPS)]
        if seg.shape[1] < 2:
            continue
        fx = np.maximum(0, np.diff(seg, axis=1)).sum(0)
        feats.append(np.concatenate([seg.mean(1), seg.std(1), [fx.mean(), fx.std(), fx.max()]]))
        spans.append((round(float(a), 3), round(float(b), 3)))
    if not feats:
        return None
    cid = {"easy": 0, "normal": 1, "hard": 2, "oni": 3}[course]
    x = torch.tensor(np.array(feats), dtype=torch.float32)[None].to(DEVICE)
    with torch.no_grad():
        pd, pf = planner(x, torch.tensor([cid], device=DEVICE))
    d8 = pd[0].argmax(-1).cpu().numpy()
    fl = pf[0].argmax(-1).cpu().numpy()
    return [[a, b, int(d), int(f)] for (a, b), d, f in zip(spans, d8, fl)]


def group_quantize(times, phase, grid, min_run=3):
    n = len(times)
    slots = [0] * n
    i = 0
    while i < n:
        j = i
        while j + 1 < n:
            ioi = times[j + 1] - times[j]
            ref = (times[j] - times[i]) / (j - i) if j > i else ioi
            if 0.02 < ioi < 1.2 and abs(ioi - ref) < 0.22 * max(ref, 1e-6):
                j += 1
            else:
                break
        if j - i + 1 >= min_run:
            k = max(1, int(round((times[j] - times[i]) / (j - i) / grid)))
            anchor = int(round((times[i] - phase) / grid))
            for m in range(j - i + 1):
                slots[i + m] = anchor + m * k
        else:
            for m in range(i, j + 1):
                slots[m] = int(round((times[m] - phase) / grid))
        i = j + 1
    return slots


def upload_wave_name(audio_path):
    name = os.path.basename(str(audio_path).replace("\\", "/")).strip()
    name = re.sub(r"[\x00-\x1f\x7f]+", " ", name).strip()
    return name or "song.ogg"


def output_tja_path(wave_name, course, directory):
    stem = os.path.splitext(os.path.basename(wave_name))[0].strip() or "softchart"
    stem = re.sub(r"[<>:\"/\\|?*\x00-\x1f]+", "_", stem).strip(" ._") or "softchart"
    return os.path.join(directory, f"{stem}_{course}.tja")


def write_tja(gen, bpm, title, course, level, wave, downbeats=None, grid_fit=None,
              plan=None):
    hits = sorted((h["t"], CHAR[h["type"]]) for h in gen["hits"])
    beat = 60.0 / bpm
    grid = beat / (SUB / 4)
    bias = 0.0
    if grid_fit is not None and hits:
        # authoritative fitted grid: barlines ARE the fitted downbeats.
        # De-bias the generator's systematic latency (global shift only),
        # then anchor slot 0 on the last fitted barline at/before the first note.
        times, bias = debias_to_grid([t for t, _ in hits], grid_fit["phase"], grid)
        phase = grid_fit["phase"] + float(np.floor((times[0] - grid_fit["phase"]) / (4 * beat))) * 4 * beat
        q_times = list(times)
    else:
        times = np.array([t for t, _ in hits]) if hits else np.array([0.0])
        cands = np.arange(0, beat, grid / 4)
        phase = float(cands[int(np.argmin([np.mean(np.abs(((times - o) / grid) - np.round((times - o) / grid))) for o in cands]))])
        q_times = [t for t, _ in hits]
    slot_idx = group_quantize(q_times, phase, grid)
    if slot_idx and min(slot_idx) < 0:
        # note quantized just before the anchor barline: pull back whole bars
        # so nothing is dropped (barline alignment is preserved mod SUB)
        nb = int(np.ceil(-min(slot_idx) / SUB))
        slot_idx = [s + nb * SUB for s in slot_idx]
        phase -= nb * SUB * grid
    slots = {}
    for idx, (t, ch) in zip(slot_idx, hits):
        if idx >= 0 and idx not in slots:
            slots[idx] = ch
    for sp in gen["spans"]:
        i0 = int(round((sp["t0"] - bias - phase) / grid))
        i1 = int(round((sp["t1"] - bias - phase) / grid))
        while i0 in slots:
            i0 += 1
        while i1 in slots or i1 <= i0:
            i1 += 1
        if i0 >= 0:
            slots[i0] = CHAR[sp["type"]]
            slots[i1] = "8"
    if slots and grid_fit is None:
        # Gridless anchoring: shift so the first note sits on a detected
        # downbeat if one is nearby, otherwise on the first barline.
        first_t = min(slots) * grid + phase
        anchor_t = None
        if downbeats is not None and len(downbeats):
            near = downbeats[downbeats <= first_t + 0.12]
            if len(near) and first_t - near[-1] < 4 * beat:
                anchor_t = near[-1]
        shift = int(round((anchor_t - phase) / grid)) if anchor_t is not None else min(slots)
        if shift:
            slots = {k - shift: v for k, v in slots.items()}
            phase += shift * grid
    n_meas = (max(slots) // SUB + 1) if slots else 1
    measure_starts = phase + np.arange(n_meas + 1, dtype=float) * (4 * beat)
    gogo_mask = gogo_measure_mask(plan, measure_starts, n_meas)
    lines = []
    in_gogo = False
    for m in range(n_meas):
        in_gogo = append_measure_with_gogo(
            lines,
            "".join(slots.get(m * SUB + k, "0") for k in range(SUB)) + ",",
            m, gogo_mask, in_gogo)
    if in_gogo:
        lines.append("#GOGOEND")
    balloons = [10] * sum(1 for s in gen["spans"] if s["type"] == "balloon")
    return "\n".join([
        f"TITLE:{title} (SoftChart)", f"BPM:{bpm:g}", f"WAVE:{wave}",
        f"OFFSET:{-phase:.3f}", f"COURSE:{'Oni' if course == 'oni' else course.capitalize()}",
        f"LEVEL:{level}", f"BALLOON:{','.join(map(str, balloons))}" if balloons else "BALLOON:",
        "", "#START", *lines, "#END"]) + "\n"


def render_audio_plan(mel, title, course, out_path, plan=None):
    import matplotlib
    matplotlib.use("Agg")
    from matplotlib import font_manager
    import matplotlib.pyplot as plt

    font_path = cjk_font_path()
    if font_path is not None:
        font_manager.fontManager.addfont(font_path)
        matplotlib.rcParams["font.family"] = font_manager.FontProperties(fname=font_path).get_name()
    matplotlib.rcParams["axes.unicode_minus"] = False

    dur = mel.shape[1] / FPS
    fig = plt.figure(figsize=(13, 4.4 if plan else 3.2))
    gs = fig.add_gridspec(2 if plan else 1, 1,
                          height_ratios=[3.0, 1.0] if plan else [1],
                          hspace=0.14 if plan else 0.0)
    ax0 = fig.add_subplot(gs[0])
    ax0.imshow(mel, aspect="auto", origin="lower",
               cmap="magma", extent=[0, dur, 0, N_MELS])
    ax0.set_ylabel("mel")
    ax0.set_title(f"{title} — {course}  |  full-song mel spectrogram")
    ax0.set_xlim(0, dur)
    ax0.grid(axis="x", alpha=0.18)
    if plan:
        ax0.set_xticklabels([])
        ax1 = fig.add_subplot(gs[1], sharex=ax0)
        for a, b, d, f in plan:
            c = "#d64545" if f == 2 else ("#4a90d9" if f == 1 else "#999999")
            ax1.bar((a + b) / 2, max(d, 0.15), width=max((b - a) * 0.92, 0.01),
                    color=c, alpha=0.85)
        ax1.set_xlim(0, dur)
        ax1.set_ylim(0, 8)
        ax1.set_yticks([0, 4, 8])
        ax1.set_ylabel("plan", fontsize=8)
        ax1.set_xlabel("time (s)  —  plan: grey=density blue=gap red=climax")
        ax1.grid(axis="x", alpha=0.18)
    else:
        ax0.set_xlabel("time (s)")
    fig.savefig(out_path, dpi=130, bbox_inches="tight")
    plt.close(fig)
    return out_path


def _progress(stage, fraction, title, detail):
    return {
        "kind": "progress",
        "stage": stage,
        "progress": fraction,
        "title": title,
        "detail": detail,
    }


def _uploaded_file(value, original_name):
    if value is None:
        raise ValueError("Upload a music file to begin.")
    if not isinstance(original_name, str) or not original_name.strip():
        raise ValueError("The upload is missing its original filename. Please select it again.")
    data = value if isinstance(value, gr.FileData) else gr.FileData.model_validate(value)
    path = os.path.realpath(data.path)
    if not os.path.isfile(path):
        raise ValueError("The uploaded file is no longer available. Please select it again.")
    return path, upload_wave_name(original_name)


def _validate_controls(course, level, bpm, temperature, top_p, drum_volume):
    if course not in COURSE_DENS:
        raise ValueError(f"Unsupported difficulty: {course}")
    if not 1 <= int(level) <= 10:
        raise ValueError("Level must be between 1 and 10 stars.")
    if bpm != 0 and not 30 <= float(bpm) <= 400:
        raise ValueError("Manual BPM must be between 30 and 400; use 0 for auto-detection.")
    if not 0.2 <= float(temperature) <= 1.2:
        raise ValueError("Temperature must be between 0.2 and 1.2.")
    if not 0.5 <= float(top_p) <= 1.0:
        raise ValueError("Top-p must be between 0.5 and 1.0.")
    if not 0 <= float(drum_volume) <= 1.5:
        raise ValueError("Taiko volume must be between 0% and 150%.")


def _file_data(path, *, name=None, mime_type=None):
    return gr.FileData(
        path=os.path.realpath(path),
        orig_name=name or os.path.basename(path),
        mime_type=mime_type,
    ).model_dump()


app = gr.Server(
    title="SoftChart",
    description="Conditional Taiko chart generation with synchronized audio preview.",
)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")


@app.get("/", include_in_schema=False)
async def homepage():
    return FileResponse(
        STATIC_DIR / "index.html",
        headers={"Cache-Control": "no-cache"},
    )


@app.get("/health", include_in_schema=False)
async def health():
    return {"status": "ok", "device": DEVICE, "models_loaded": "gen" in _MODELS}


@app.api(
    name="generate_chart",
    description="Generate a TJA chart, visual previews, and a synchronized taiko audio mix.",
    concurrency_limit=1,
    concurrency_id="generation-gpu",
    queue=True,
    stream_every=0.1,
)
def generate_chart(
    audio: gr.FileData,
    audio_name: str,
    course: str,
    level: int,
    bpm_override: float,
    auto_plan_on: bool,
    use_beat: bool,
    use_planner: bool,
    sampling: bool,
    temperature: float,
    top_p: float,
    drum_volume: float,
    model_choice: str = DEFAULT_MODEL,
) -> dict[str, object]:
    """Stream the actual inference stages to the custom frontend."""
    workdir = tempfile.mkdtemp(prefix="softchart-request-")
    try:
        audio_path, wave_name = _uploaded_file(audio, audio_name)
        _validate_controls(course, level, bpm_override, temperature, top_p, drum_volume)
        level = int(level)
        bpm_override = float(bpm_override)
        temperature = float(temperature)
        top_p = float(top_p)
        drum_volume = float(drum_volume)

        if model_choice not in MODELS:
            raise ValueError(
                f"Unknown model {model_choice!r}. Choose one of: {', '.join(MODELS)}."
            )
        yield _progress(
            "loading", 0.03, "Preparing",
            f"Loading SoftChart {model_choice}.",
        )
        models = get_models(model_choice, include_planner=use_planner)

        yield _progress(
            "audio", 0.11, "Listening",
            "Mapping rhythm, melody, and timbre.",
        )
        mel, wav = load_logmel(audio_path)

        yield _progress(
            "beat", 0.21, "Finding the beat",
            "Finding beats and tempo.",
        )
        grid = dbs = None
        if use_beat:
            grid = fit_grid_piecewise(models["beat"], mel, device=DEVICE)
            if grid is not None:
                dbs = grid["downbeats"] if grid["ok"] else grid["db_peaks"]
                if not grid["ok"]:
                    grid = None
        if bpm_override > 0:
            bpm = bpm_override
            if use_beat and (grid is None or abs(grid["bpm"] - bpm) > 0.5):
                fixed_grid = fit_grid_fixed_bpm(models["beat"], mel, bpm, device=DEVICE)
                grid = fixed_grid if fixed_grid is not None and fixed_grid["ok"] else None
                if grid is not None:
                    dbs = grid["downbeats"]
        elif grid is not None:
            bpm = float(grid["bpm"])
        elif dbs is not None and len(dbs) > 4:
            period = float(np.median(np.diff(dbs)))
            if period <= 0:
                raise RuntimeError("The beat model returned an invalid downbeat interval.")
            bpm = 240.0 / period
            while bpm >= 210:
                bpm /= 2.0
            while bpm < 70:
                bpm *= 2.0
        else:
            import librosa
            bpm = float(np.atleast_1d(librosa.beat.beat_track(y=wav, sr=SR)[0])[0])
        if not np.isfinite(bpm) or bpm <= 0:
            raise RuntimeError("Could not determine a reliable BPM. Enter it in Advanced settings.")
        if grid is None and abs(bpm - round(bpm)) < 0.06:
            bpm = float(round(bpm))

        yield _progress(
            "plan", 0.34, "Shaping the arc",
            "Planning density and climaxes.",
        )
        plan = None
        if getattr(models["gen"], "_has_plan", False):
            if use_planner:
                plan = learned_plan(models["plan"], mel, course, bpm, dbs)
            elif auto_plan_on:
                plan = auto_plan(mel, bpm, dbs)

        yield _progress(
            "generate", 0.47, "Writing the chart",
            "Writing playable Taiko patterns.",
        )
        title = os.path.splitext(wave_name)[0]
        slot_used = grid is not None
        if slot_used:
            generated = generate_song_slot(
                models["slot"], mel, grid, course, level=level,
                density_bucket=COURSE_DENS[course], greedy=not sampling, seed=0,
                temperature=temperature, top_p=top_p, device=DEVICE, plan=plan,
            )
            tja = write_tja_slots(
                generated, grid, title, course, level, wave_name, plan=plan,
            )
        else:
            generated = generate_song(
                models["gen"], mel, course, level=level,
                density_bucket=COURSE_DENS[course], greedy=not sampling,
                temperature=temperature, top_p=top_p, seed=0,
                device=DEVICE, plan=plan,
            )
            generated = snap_chart(generated, bpm)
            tja = write_tja(
                generated, bpm, title, course, level, wave_name, dbs,
                grid_fit=grid, plan=plan,
            )

        yield _progress(
            "export", 0.81, "Rendering",
            "Building the TJA and previews.",
        )
        tja_path = output_tja_path(wave_name, course, workdir)
        with open(tja_path, "w", encoding="utf-8") as output_file:
            output_file.write(tja)
        chart_image_path = os.path.join(workdir, "chart.png")
        plan_image_path = os.path.join(workdir, "song-plan.png")
        render_tja_image(tja, out_path=chart_image_path)
        render_audio_plan(mel, title, course, plan=plan, out_path=plan_image_path)

        yield _progress(
            "mix", 0.91, "Mixing",
            "Mixing Taiko with your track.",
        )
        stem = Path(tja_path).stem
        preview_path = os.path.join(workdir, f"{stem}_taiko-preview.wav")
        mix_stats = synthesize_taiko_preview(
            audio_path,
            tja,
            preview_path,
            drum_gain=drum_volume,
            sample_rate=44100,
        )

        grid_rms = round(float(grid["rms_ms"]), 1) if grid is not None else None
        metrics = {
            "bpm": round(float(bpm), 1),
            "notes": len(generated["hits"]),
            "spans": len(generated["spans"]),
            "timing": "slot-exact" if slot_used else "time-quantized",
            "grid_rms_ms": grid_rms,
            "preview_hits": int(mix_stats["rendered_hit_count"]),
        }
        yield {
            "kind": "complete",
            "stage": "complete",
            "progress": 1.0,
            "title": "Ready",
            "detail": "Play it or download it.",
            "metrics": metrics,
            "files": {
                "tja": _file_data(tja_path, mime_type="text/plain"),
                "audio": _file_data(preview_path, mime_type="audio/wav"),
                "chart_image": _file_data(chart_image_path, mime_type="image/png"),
                "plan_image": _file_data(plan_image_path, mime_type="image/png"),
            },
        }
    except Exception as exc:
        LOGGER.exception("SoftChart generation failed")
        detail = (str(exc) if isinstance(exc, (ValueError, RuntimeError))
                  else "The server could not complete this chart. Please try again shortly.")
        yield {
            "kind": "error",
            "stage": "error",
            "progress": 0.0,
            "title": "Generation did not complete",
            "detail": detail,
        }
    finally:
        shutil.rmtree(workdir, ignore_errors=True)


if __name__ == "__main__":
    app.launch(max_file_size="200mb")