#!/usr/bin/env python3
"""Crop licensed demo clips and prepare notes with the fixed browser ONNX model."""
from __future__ import annotations

import argparse
import csv
import hashlib
import json
from pathlib import Path
import re
import shutil
import subprocess
import time

import numpy as np
import onnxruntime as ort
import soundfile as sf
import torch

ROOT = Path(__file__).resolve().parents[1]


def digest(path: Path) -> str:
    with path.open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def generate(sessions, features):
    encoder, cross_cache, decoder = sessions
    encoded = encoder.run(None, {"spectrogram": features})[0]
    cross = cross_cache.run(None, {"encoder_hidden_states": encoded})
    cache = [np.zeros((13, 6, 0, 64), np.float32) for _ in range(16)]
    token = np.zeros((1, 13, 1), np.int64)
    unfinished = np.ones_like(token)
    windows = []
    for position in range(256):
        feeds = {"token_ids": token, "position_ids": np.asarray([position], np.int64)}
        for layer in range(8):
            for suffix, values in (("self_key", cache[2*layer]), ("self_value", cache[2*layer+1]),
                                   ("cross_key", cross[2*layer]), ("cross_value", cross[2*layer+1])):
                feeds[f"layer_{layer}_{suffix}"] = values
        result = decoder.run(None, feeds)
        token = result[0].argmax(-1).astype(np.int64)
        visible = token * unfinished
        windows.append(visible)
        cache = result[1:]
        unfinished *= visible != 1
        if not unfinished.any():
            break
    return np.concatenate(windows, axis=-1)[0].tolist(), int(unfinished.sum())


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset-root", type=Path, default=Path("/home/jack/workspace/datasets"))
    parser.add_argument("--graph-dir", type=Path, default=ROOT / "progress/evaluations/yptf_web_fp32_d26")
    parser.add_argument("--soundfont", type=Path, default=Path("/home/jack/workspace/models/soundfonts/GeneralUser_GS.sf2"))
    parser.add_argument("--output-dir", type=Path, default=ROOT / "demo/web/public/examples")
    args = parser.parse_args()
    config = json.loads((ROOT / "configs/demo/public_examples.json").read_text())
    for row in config["examples"]:
        if not re.fullmatch(r"[a-z0-9-]+", row["id"]):
            raise SystemExit(f"Example ID must be safe in a URL: {row['id']}")
    manifest = json.loads((args.graph_dir / "manifest.json").read_text())
    if manifest["precision"] != "fp32":
        raise SystemExit("Public examples require the checked FP32 browser export")
    options = ort.SessionOptions()
    options.intra_op_num_threads = 4
    options.inter_op_num_threads = 1
    torch.set_num_threads(4)
    sessions = []
    for name in ("encoder", "cross_cache", "decoder_step"):
        path = args.graph_dir / f"{name}.onnx"
        if digest(path) != manifest["graphs"][name]["sha256"]:
            raise SystemExit(f"Graph hash mismatch: {path}")
        sessions.append(ort.InferenceSession(str(path), sess_options=options, providers=["CPUExecutionProvider"]))

    # Check each selected FMA license against the frozen dataset metadata.
    with (args.dataset_root / "fma/fma_metadata/tracks.csv").open() as stream:
        reader = csv.reader(stream)
        groups, fields = next(reader), next(reader)
        license_index = list(zip(groups, fields)).index(("track", "license"))
        licenses = {int(row[0]): row[license_index] for row in reader if row[0].isdigit()}
    for row in config["examples"]:
        if row["collection"] == "Free Music Archive":
            assert licenses[row["track_id"]] == row["license_metadata"], row["id"]
            assert row["license_metadata"] in {"Attribution 3.0 United States", "Attribution 3.0 International"}

    args.output_dir.mkdir(parents=True, exist_ok=True)
    catalog = {"schema_version": 1, "default_example": config["default_example"], "examples": []}
    credits = ["# Demo audio credits", "", "All clips are cropped to at most 12 seconds and converted to mono 16 kHz PCM.", "Prepared notes are Hendrix predictions, not dataset annotations.", ""]
    hann = torch.hann_window(2048)
    started = time.monotonic()
    for row in config["examples"]:
        audio_path = args.output_dir / f"{row['id']}.wav"
        source = args.dataset_root / row["source"]
        subprocess.run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-ss", str(row["offset_seconds"]),
                        "-i", str(source), "-t", str(row["duration_seconds"]), "-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", str(audio_path)], check=True)
        audio, rate = sf.read(audio_path, dtype="float32")
        assert rate == 16000 and audio.ndim == 1
        duration = len(audio) / rate
        tokens, truncations = [], 0
        for start in range(0, len(audio), 32767):
            window = np.zeros(32767, np.float32)
            segment = audio[start:start+32767]
            window[:len(segment)] = segment
            magnitudes = torch.stft(torch.from_numpy(window), n_fft=2048, hop_length=300,
                                    window=hann, return_complex=True).abs()[1:].transpose(0, 1)
            features = torch.log(magnitudes + 1e-5)
            features = (features - features.min()) / (features.max() - features.min() + 0.008)
            predicted, unfinished = generate(sessions, features.unsqueeze(0).numpy())
            tokens.append(predicted)
            truncations += bool(unfinished)
        decoded = subprocess.run(["node", "--experimental-strip-types", str(ROOT / "scripts/decode_demo_tokens.mjs")],
                                 input=json.dumps({"windows": tokens, "segmentSeconds": 32767/16000, "duration": duration}),
                                 text=True, capture_output=True, check=True)
        notes = json.loads(decoded.stdout)
        payload = {"schema_version": 1, "duration": duration, "notes": notes,
                   "checkpoint_sha256": manifest["checkpoint"]["sha256"], "audio_sha256": digest(audio_path),
                   "runtime": "Fixed FP32 ONNX browser graphs, CPU, shared browser note decoder", "truncated_windows": truncations}
        note_path = args.output_dir / f"{row['id']}.json"
        note_path.write_text(json.dumps(payload, separators=(",", ":")) + "\n")
        peaks = [float(np.sqrt(np.mean(chunk**2))) for chunk in np.array_split(audio, 56)]
        maximum = max(peaks) or 1
        peaks = [round((value/maximum)**0.65, 3) for value in peaks]
        public = {**row, "audio_url": f"/examples/{audio_path.name}", "notes_url": f"/examples/{note_path.name}",
                  "duration": duration, "peaks": peaks, "note_count": len(notes), "audio_sha256": digest(audio_path),
                  "notes_sha256": digest(note_path), "source_sha256": digest(source)}
        catalog["examples"].append(public)
        credits += [f"## {row['title']}", f"Artist/creator: {row['artist']}",
                    f"Source: [{row['collection']}]({row['source_url']}), `{row['source']}`.",
                    f"License: [{row['license_name']}]({row['license_url']}).",
                    f"Change: {row['offset_seconds']}–{row['offset_seconds']+duration:g} seconds, mono 16 kHz excerpt.", ""]
        print(f"{row['id']}: {duration:g}s, {len(notes)} notes, {truncations} truncated windows", flush=True)
    (args.output_dir / "catalog.json").write_text(json.dumps(catalog, indent=2) + "\n")
    (args.output_dir / "CREDITS.md").write_text("\n".join(credits))
    fonts = args.output_dir.parent / "soundfonts"
    fonts.mkdir(parents=True, exist_ok=True)
    shutil.copyfile(args.soundfont, fonts / "GeneralUser_GS.sf2")
    shutil.copyfile(ROOT / "deploy/huggingface/GeneralUser-LICENSE.txt", fonts / "LICENSE.txt")
    print(f"Prepared {len(catalog['examples'])} examples in {time.monotonic()-started:.1f}s: {args.output_dir}", flush=True)


if __name__ == "__main__":
    main()