#!/usr/bin/env python3
"""
Slim the embedded cover art across the music library.

WHY. iOS AVPlayer fetches the head of a streamed file in ~32 KB range requests, one round trip
each, and it cannot reach audio until it has read past the whole ID3 tag. A 955 KB embedded JPEG
(Ben Harper, 2026-08-29) therefore cost ~30 sequential round trips - a two-second silence on the
iPhone that Android (ExoPlayer, one connection) never feels. 28% of the library rips carry art
this size.

THE LIMIT, AND WHY. Every track KEEPS an embedded front cover - the folder cover.jpg is this
plugin's primary art, but embedded art is what any OTHER player/user of the files sees, so it
stays. It is recompressed to fit the transport: 600x600 max (the iTunes standard, plenty for any
lock screen), JPEG, quality stepped down from 85 until it fits in EMBED_TARGET = 80 KB. 80 KB
plus the text frames keeps the whole tag within ~3 range-request chunks - imperceptible - while
staying visually indistinguishable at the sizes any UI renders it.

Only files whose picture frames total more than PROCESS_OVER = 100 KB are touched at all; a file
already carrying reasonable art is left byte-identical.

SAFETY. Work happens on a temp copy; the post-tag byte region (the audio, verbatim) is hashed
before and after and must match or the original is kept. If an album folder has NO cover.jpg or
folder.jpg, the original full-resolution embedded art is rescued to cover.jpg before compression.

Usage:
  slim-embedded-art.py                                 dry-run over ~/Music
  slim-embedded-art.py "Albums/Ben Harper/[1994] ..."  dry-run, one directory
  slim-embedded-art.py [subpath] --apply               actually rewrite
"""
import hashlib
import io
import os
import shutil
import sys

from mutagen.id3 import ID3, APIC, ID3NoHeaderError
from PIL import Image

ROOT = os.path.expanduser("~/Music")
EMBED_TARGET = 80_000
PROCESS_OVER = 100_000


def audio_region_hash(path):
    """SHA-256 of everything after the ID3v2 tag - the audio, byte-for-byte."""
    with open(path, "rb") as f:
        head = f.read(10)
        off = 0
        if head[:3] == b"ID3":
            off = 10 + ((head[6] << 21) | (head[7] << 14) | (head[8] << 7) | head[9])
            if head[5] & 0x10:  # v2.4 footer
                off += 10
        f.seek(off)
        h = hashlib.sha256()
        while chunk := f.read(1 << 20):
            h.update(chunk)
        return h.hexdigest()


def pictures(tag):
    return tag.getall("APIC") + tag.getall("PIC")


def compress_cover(data):
    """(jpeg_bytes, note) at <= EMBED_TARGET, stepping quality then size down."""
    img = Image.open(io.BytesIO(data))
    if img.mode != "RGB":
        img = img.convert("RGB")
    for side in (600, 500, 400):
        work = img.copy()
        work.thumbnail((side, side))
        for q in (85, 80, 70, 60, 50):
            buf = io.BytesIO()
            work.save(buf, "JPEG", quality=q, optimize=True)
            if buf.tell() <= EMBED_TARGET:
                return buf.getvalue(), f"{side}px q{q}"
    return buf.getvalue(), f"{side}px q{q} OVER"  # pathological; still vastly smaller


def rescue_cover(dirpath, data):
    """Write the original art to cover.jpg when the folder has no art file at all."""
    for name in ("cover.jpg", "folder.jpg"):
        if os.path.exists(os.path.join(dirpath, name)):
            return False
    img = Image.open(io.BytesIO(data))
    if img.mode != "RGB":
        img = img.convert("RGB")
    img.save(os.path.join(dirpath, "cover.jpg"), "JPEG", quality=90)
    return True


def process(path, apply):
    try:
        tag = ID3(path)
    except ID3NoHeaderError:
        return None
    pics = pictures(tag)
    total = sum(len(p.data) for p in pics)
    if total <= PROCESS_OVER:
        return None
    largest = max(pics, key=lambda p: len(p.data))
    if not apply:
        return f"would slim {total // 1024} KB of art"

    before = audio_region_hash(path)
    tmp = path + ".slimtmp"
    shutil.copy2(path, tmp)
    try:
        rescued = rescue_cover(os.path.dirname(path), largest.data)
        jpeg, note = compress_cover(largest.data)
        t = ID3(tmp)
        t.delall("APIC")
        t.delall("PIC")
        t.add(APIC(encoding=0, mime="image/jpeg", type=3, desc="", data=jpeg))
        v2 = 4 if tag.version >= (2, 4, 0) else 3
        t.save(v2_version=v2, padding=lambda info: 0)
        if audio_region_hash(tmp) != before:
            raise RuntimeError("audio bytes changed - keeping original")
        os.replace(tmp, path)
        note_out = f"{total // 1024} KB -> {len(jpeg) // 1024} KB ({note})"
        return note_out + (" +rescued cover.jpg" if rescued else "")
    except Exception as e:
        if os.path.exists(tmp):
            os.remove(tmp)
        return f"FAILED: {e}"


def main():
    args = [a for a in sys.argv[1:] if a != "--apply"]
    apply = "--apply" in sys.argv
    base = os.path.join(ROOT, args[0]) if args else ROOT
    done = failed = skipped = 0
    for dirpath, _, files in os.walk(base):
        for name in sorted(files):
            if not name.lower().endswith(".mp3"):
                continue
            result = process(os.path.join(dirpath, name), apply)
            rel = os.path.join(dirpath, name)[len(ROOT) + 1 :]
            if result is None:
                skipped += 1
            elif result.startswith("FAILED"):
                failed += 1
                print(f"  !! {rel}  {result}", flush=True)
            else:
                done += 1
                if done % 250 == 0 or not apply:
                    print(f"  {rel}  {result}", flush=True)
    print(f"{'applied' if apply else 'dry-run'}: {done} slimmed, {skipped} already fine, {failed} failed", flush=True)


if __name__ == "__main__":
    main()