"""One isolated SheetSage2 CUDA transcription per request."""

from __future__ import annotations

import importlib
import json
from pathlib import Path
import shutil
import sys

import numpy as np
import soundfile as sf
import torch
from transformers import AutoModel


SHEETSAGE_REV = "eab522a8168e8b8b8c4856bf8609cd86198f01fe"


def run(source: Path, output: Path) -> None:
    if not torch.cuda.is_available():
        raise RuntimeError("SheetSage2 requires the configured NVIDIA GPU")
    waveform, rate = sf.read(source, dtype="float32", always_2d=True)
    if rate != 24000 or waveform.shape[1] != 1 or not np.isfinite(waveform).all():
        raise ValueError("Expected finite mono 24 kHz PCM")
    torch.manual_seed(1024)
    np.random.seed(1024)
    model = AutoModel.from_pretrained(
        "m-a-p/SheetSage2", revision=SHEETSAGE_REV, trust_remote_code=True,
        torch_dtype=torch.float32, attn_implementation="sdpa",
    ).eval().to("cuda")
    native = output / "native"
    result = model.transcribe(waveform[:, 0], sampling_rate=24000,
                              output_dir=native, dtype="fp32", preset="default",
                              melody_only=False)
    if result.get("abc_error") or not (native / "score.abc").is_file():
        raise RuntimeError(f"SheetSage2 could not export ABC: {result.get('abc_error')}")
    package = model.__class__.__module__.rsplit(".", 1)[0]
    exporter = importlib.import_module(package + ".exports_sheetsage2").export_result
    events = json.loads((native / "events.json").read_text())
    melody = exporter(events, model.tokenizer, output / "melody-only",
                      duration=result["duration_seconds"], melody_only=True)
    if melody.get("abc_error") or not (output / "melody-only/score.abc").is_file():
        raise RuntimeError(f"SheetSage2 could not export melody ABC: {melody.get('abc_error')}")
    shutil.copy2(native / "score.abc", output / "full.abc")
    shutil.copy2(output / "melody-only/score.abc", output / "melody.abc")
    shutil.copy2(native / "melody_vocal.mid", output / "notes.mid")
    shutil.copy2(native / "beat.lab", output / "beat.lab")
    (output / "sheet.json").write_text(json.dumps({
        "model": "m-a-p/SheetSage2", "revision": SHEETSAGE_REV,
        "duration_s": result["duration_seconds"],
        "source_timeline": True,
    }, indent=2) + "\n")


if __name__ == "__main__":
    run(Path(sys.argv[1]), Path(sys.argv[2]))