#!/usr/bin/env python3
"""
notenames.py — turn piano sheet music into "note-name notation":
every notehead becomes a boxed letter with an explicit accidental
(Eb, F#, B, ...), so the key signature never has to be computed.

Pipeline:  image/PDF  --oemer-->  MusicXML  --musicxml2ly-->  .ly  --lilypond-->  PDF + PNG

Usage:
  notenames.py INPUT [-o OUTDIR] [--scale N] [--dpi N] [--no-omr]

  INPUT  .png/.jpg/.webp/.pdf  (goes through OMR)
         .musicxml/.xml/.mxl   (skips OMR, just re-renders)
         .ly                   (skips OMR and conversion, just re-renders)
"""
import argparse, os, re, shutil, subprocess, sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
INCLUDE = HERE / "notename-heads.ily"


def run(cmd, **kw):
    print("+", " ".join(str(c) for c in cmd), flush=True)
    return subprocess.run([str(c) for c in cmd], check=True, **kw)


def rasterize(src: Path, outdir: Path, dpi: int) -> list[Path]:
    """PDF -> one PNG per page; images -> single PNG (upscaled if small)."""
    from PIL import Image
    pages = []
    if src.suffix.lower() == ".pdf":
        import fitz  # pymupdf
        doc = fitz.open(src)
        for i, page in enumerate(doc, 1):
            pix = page.get_pixmap(dpi=dpi)
            p = outdir / f"{src.stem}_p{i}.png"
            pix.save(p)
            pages.append(p)
    else:
        im = Image.open(src).convert("RGB")
        # OMR wants roughly 300 dpi; upscale small screenshots.
        if im.width < 2800:
            f = max(1, round(3000 / im.width))
            im = im.resize((im.width * f, im.height * f), Image.LANCZOS)
        p = outdir / f"{src.stem}_p1.png"
        im.save(p)
        pages.append(p)
    return pages


def omr(png: Path, outdir: Path) -> Path:
    oemer = HERE / ".venv" / "bin" / "oemer"
    run([oemer, "--save-cache", "-o", outdir, png])
    xml = outdir / (png.stem + ".musicxml")
    if not xml.exists():
        sys.exit(f"oemer produced no MusicXML for {png}")
    return xml


def to_ly(xml: Path, outdir: Path) -> Path:
    ly = outdir / (xml.stem + ".ly")
    run(["musicxml2ly", "--no-beaming", "-o", ly, xml])
    return ly


def inject(ly: Path, outdir: Path | None = None) -> Path:
    """Add the include right after \\version so note-name heads become the default."""
    text = ly.read_text()
    inc = f'\\include "{INCLUDE}"\n'
    if inc in text:
        return ly
    if re.search(r'^\\version', text, re.M):
        text = re.sub(r'(^\\version[^\n]*\n)', r'\1' + inc.replace("\\", "\\\\"), text, count=1, flags=re.M)
    else:
        text = '\\version "2.24.0"\n' + inc + text
    out = (outdir or ly.parent) / (ly.stem + "_names.ly")
    out.write_text(text)
    return out


def engrave(ly: Path, outdir: Path):
    run(["lilypond", "-dno-point-and-click", "-o", outdir / ly.stem, ly])
    run(["lilypond", "-dno-point-and-click", "--png", "-dresolution=150", "-o", outdir / ly.stem, ly])


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("input")
    ap.add_argument("-o", "--outdir", default=None)
    ap.add_argument("--dpi", type=int, default=300, help="PDF rasterization dpi")
    args = ap.parse_args()

    src = Path(args.input).expanduser().resolve()
    outdir = Path(args.outdir or (HERE / "out" / src.stem)).resolve()
    outdir.mkdir(parents=True, exist_ok=True)
    suf = src.suffix.lower()

    if suf == ".ly":
        ly_files = [inject(src, outdir)]
    elif suf in (".musicxml", ".xml", ".mxl"):
        ly_files = [inject(to_ly(src, outdir))]
    else:
        ly_files = []
        for png in rasterize(src, outdir, args.dpi):
            ly_files.append(inject(to_ly(omr(png, outdir), outdir)))

    for ly in ly_files:
        engrave(ly, outdir)
    print("\nDone. Outputs in", outdir)
    for f in sorted(outdir.iterdir()):
        if f.suffix in (".pdf", ".png") and "_names" in f.name:
            print("  ", f)


if __name__ == "__main__":
    main()