#!/usr/bin/env python3
"""Embed album art that iTunes had but never wrote into the audio files.

lintunes shows art read live from each file's embedded tags, so any art that
lived only in iTunes' Album Artwork cache is invisible until it's embedded. This
backfills it, in order of reliability, and writes nothing it isn't sure about:

  tier2 (album)  for a track with no embedded art, copy a same-album track's
                 embedded art (iTunes shows one cover per album)
  tier3 (cache)  decode iTunes' Album Artwork/Cache/<lib>-<PID>.itc (keyed by
                 the track's persistent id) and embed it

Download/Store-only art (an opaque id with no track link from here) is *not*
guessed at -- it's reported so you can decide separately (see --audit).

Defaults to --dry-run: nothing is modified until you pass --write. Run the audit
first (scripts/audit_artwork.py) to see how many tracks each tier covers.

Usage:
    python3 scripts/recover_artwork.py --data-dir ./data           # preview
    python3 scripts/recover_artwork.py --data-dir ./data --write   # apply
"""
import argparse
import glob
import json
import os
import sys
from collections import defaultdict
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from lintunes import tagging  # noqa: E402
from lintunes.itc import decode_itc_file  # noqa: E402


def album_key(t: dict):
    album = (t.get("album") or "").strip().lower()
    if not album:
        return None
    who = (t.get("album_artist") or t.get("artist") or "").strip().lower()
    return (who, album)


def build_cache_index(artwork_dir: Path) -> dict[str, str]:
    """persistent id (upper) -> path of its Cache/*.itc."""
    index = {}
    for f in glob.glob(str(artwork_dir / "Cache" / "**" / "*.itc"),
                       recursive=True):
        token = os.path.basename(f).rsplit(".", 1)[0].split("-")[-1].upper()
        index.setdefault(token, f)
    return index


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--data-dir", required=True, help="lintunes data dir")
    parser.add_argument("--artwork-dir",
                        help="iTunes 'Album Artwork' dir (default: alongside the "
                             "music drive root inferred from track locations)")
    parser.add_argument("--write", action="store_true",
                        help="actually embed art (default is a dry run)")
    parser.add_argument("--limit", type=int, default=0,
                        help="process only the first N tracks")
    args = parser.parse_args()
    dry = not args.write

    library_path = Path(args.data_dir) / "library.json"
    if not library_path.exists():
        sys.exit(f"Library not found: {library_path}")
    tracks = json.load(open(library_path, encoding="utf-8"))
    items = list(tracks.values())
    if args.limit:
        items = items[:args.limit]

    # Locate the Album Artwork dir: a sibling of "iTunes Media" on the drive.
    artwork_dir = Path(args.artwork_dir) if args.artwork_dir else None
    if artwork_dir is None:
        for t in items:
            loc = t.get("location") or ""
            idx = loc.find("iTunes Media")
            if idx != -1:
                artwork_dir = Path(loc[:idx]) / "Album Artwork"
                break
    if artwork_dir is None or not artwork_dir.exists():
        sys.exit("Could not find the 'Album Artwork' dir; pass --artwork-dir.")
    print(f"{'DRY RUN' if dry else 'WRITING'} | Album Artwork: {artwork_dir}\n")

    cache_index = build_cache_index(artwork_dir)

    # Pass 1: embedded-art map + album grouping (one mutagen read per file).
    embedded = {}
    album_art_source = {}    # album key -> a location that has embedded art
    for t in items:
        loc = t.get("location") or ""
        has = bool(loc) and os.path.exists(loc) and \
            tagging.read_embedded_artwork(loc) is not None
        embedded[id(t)] = has
        if has:
            key = album_key(t)
            if key is not None:
                album_art_source.setdefault(key, loc)

    # Pass 2: recover.
    stats = defaultdict(int)
    errors = []
    for t in items:
        if embedded[id(t)]:
            continue
        loc = t.get("location") or ""
        if not loc or not os.path.exists(loc):
            stats["file_missing"] += 1
            continue
        pid = (t.get("persistent_id") or "").upper()
        label = f"{t.get('artist','?')} — {t.get('name','?')}"

        # tier2: an album-mate has embedded art.
        key = album_key(t)
        src = album_art_source.get(key) if key is not None else None
        if src and src != loc:
            data = tagging.read_embedded_artwork(src)
            if data:
                mime = "image/png" if data[:8] == b"\x89PNG\r\n\x1a\n" \
                    else "image/jpeg"
                stats["tier2_album"] += 1
                if not dry:
                    try:
                        tagging.write_artwork(loc, data, mime)
                    except Exception as e:  # noqa: BLE001
                        errors.append((label, repr(e)))
                        stats["tier2_album"] -= 1
                        stats["error"] += 1
                continue

        # tier3: this track's own Cache/*.itc.
        itc = cache_index.get(pid)
        if itc:
            decoded = decode_itc_file(itc)
            if decoded is None:
                stats["tier3_raw_skip"] += 1  # ARGb/PNGf raw, can't recover
                continue
            data, mime = decoded
            stats["tier3_cache"] += 1
            if not dry:
                try:
                    tagging.write_artwork(loc, data, mime)
                except Exception as e:  # noqa: BLE001
                    errors.append((label, repr(e)))
                    stats["tier3_cache"] -= 1
                    stats["error"] += 1
            continue

        stats["unrecovered"] += 1

    print("Result:")
    for k in ("tier2_album", "tier3_cache", "tier3_raw_skip", "unrecovered",
              "file_missing", "error"):
        print(f"  {k:<16}{stats[k]:>8}")
    recovered = stats["tier2_album"] + stats["tier3_cache"]
    verb = "would embed" if dry else "embedded"
    print(f"\n{verb} art on {recovered} tracks.")
    if dry:
        print("Re-run with --write to apply.")
    if errors:
        print(f"\n{len(errors)} write error(s):")
        for label, err in errors[:20]:
            print(f"  - {label}: {err}")


if __name__ == "__main__":
    main()