#!/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
# No song has 50 KB of lyrics. The frame that set this limit held the complete text of Frank
# McCourt's "Angela's Ashes" - 1.5 MB of memoir in an Eels track's USLT, courtesy of some tagger
# malfunction - and AVPlayer pages through ALL of it before the first audio frame, art or not.
LYRICS_SANE = 50_000


def audio_region_hash(path, offset=None):
    """SHA-256 of the audio: after the ID3v2 tag, before any trailing ID3v1 tag.

    The v2 offset comes from mutagen's own parse. The trailing 128-byte ID3v1 block (starts
    "TAG") is EXCLUDED: mutagen updates an existing v1 tag on save - measured on Shamrocks and
    Shenanigans (2026-08-29), it normalised the v1 genre byte 0xFF to 0x0F (Rap) and the guard
    refused a whole album of perfectly good rewrites over that one metadata byte.
    """
    if offset is None:
        try:
            offset = ID3(path).size
        except Exception:
            offset = 0
        if not offset:
            with open(path, "rb") as f:
                head = f.read(10)
                if head[:3] == b"ID3":
                    offset = 10 + ((head[6] << 21) | (head[7] << 14) | (head[8] << 7) | head[9])
                    if head[5] & 0x10:  # v2.4 footer
                        offset += 10
    size = os.path.getsize(path)
    end = size
    if size - offset >= 128:
        with open(path, "rb") as f:
            f.seek(size - 128)
            if f.read(3) == b"TAG":
                end = size - 128
    with open(path, "rb") as f:
        f.seek(offset)
        h = hashlib.sha256()
        remaining = end - offset
        while remaining > 0:
            chunk = f.read(min(1 << 20, remaining))
            if not chunk:
                break
            h.update(chunk)
            remaining -= len(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)
    junk_lyrics = [f for f in tag.getall("USLT") if len(f.text or "") > LYRICS_SANE]
    if total <= PROCESS_OVER and not junk_lyrics:
        return None
    largest = max(pics, key=lambda p: len(p.data)) if pics else None
    if not apply:
        parts = []
        if total > PROCESS_OVER: parts.append(f"{total // 1024} KB of art")
        if junk_lyrics: parts.append(f"{sum(len(f.text) for f in junk_lyrics) // 1024} KB of junk lyrics")
        return "would slim " + " + ".join(parts)

    before = audio_region_hash(path)
    tmp = path + ".slimtmp"
    shutil.copy2(path, tmp)
    try:
        note = "art untouched"
        t = ID3(tmp)
        if largest is not None and total > PROCESS_OVER:
            rescued = rescue_cover(os.path.dirname(path), largest.data)
            jpeg, note = compress_cover(largest.data)
            t.delall("APIC")
            t.delall("PIC")
            t.add(APIC(encoding=0, mime="image/jpeg", type=3, desc="", data=jpeg))
        else:
            rescued = False
            jpeg = b""
        for f in junk_lyrics:
            t.delall(f.HashKey)
        v2 = 4 if tag.version >= (2, 4, 0) else 3
        # v1=0 REMOVES any trailing ID3v1 tag: it is a 128-byte legacy appendix fully redundant
        # with the v2 frames, and keeping it in sync is what tripped the audio guard before the
        # guard learned to exclude it. Decided 2026-08-29: strip it wherever we rewrite.
        t.save(v1=0, 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})"
        if junk_lyrics:
            note_out += f" -{sum(len(f.text) for f in junk_lyrics) // 1024} KB junk lyrics"
        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()