#!/usr/bin/env python3
"""
recover-azuracast-playlists.py — recover playlist assignment for files that
lost their GROUPING tag (e.g. from a pre-tag-copy upgrade-mp3-to-flac.sh run).

For each AzuraCast file currently unassigned to any playlist:
  1. Resolve the source playlist by matching (artist, title) against the
     Spotify playlists named in $ALEMBIC_CONFIG_DIR/pipeline/*.conf.
  2. Write GROUPING=<playlist> onto the FLAC/MP3 file.
  3. `beet update` so beets sees the new tag.
  4. Assign the file to the matching AzuraCast playlist via API.

Default is dry-run. Pass --apply to actually write tags + assign.

Env / args:
  --azuracast-key KEY    AzuraCast API key (required)
  --azuracast-base URL   default from AZURACAST_BASE env (empty if unset)
  --station N            default 1
  --apply                actually act (default is dry-run)
  --only PLAYLIST        limit recovery to one playlist (testing)
"""
import argparse, base64, json, os, re, subprocess, sys, urllib.parse, urllib.request
from difflib import SequenceMatcher
from pathlib import Path

CONF_DIR = Path(f"{os.environ.get('ALEMBIC_CONFIG_DIR', '/config')}/pipeline")
LIB_ROOT_HOST = f"{os.environ.get('MUSIC_DATA_DIR', '/data/music')}/Library"
# Still "/music" through the migration's Stage 0-3 transitional beets mount;
# revisit at Stage 4 once beets' directory: becomes MUSIC_DATA_DIR/Library.
LIB_ROOT_CONTAINER = "/music"


def normalize(s: str) -> str:
    s = (s or "").lower()
    s = re.sub(r"\s*\([^)]*\)", "", s)
    s = re.sub(r"\s*\[[^\]]*\]", "", s)
    s = re.sub(r"\s*-?\s*(extended|original|radio|club|vip|vocal)\s*(mix|edit|version)\s*$", "", s)
    s = re.sub(r"\s*feat\.?\s.*", "", s)
    s = re.sub(r"\s*ft\.?\s.*", "", s)
    s = re.sub(r"[^a-z0-9]+", "", s)
    return s


def read_conf(path: Path) -> dict:
    out = {}
    for line in path.read_text().splitlines():
        m = re.match(r"\s*([a-z0-9_-]+)\s*=\s*(.*?)\s*$", line)
        if m: out[m.group(1)] = m.group(2)
    return out


def spotify_token(cid: str, csec: str) -> str:
    creds = base64.b64encode(f"{cid}:{csec}".encode()).decode()
    req = urllib.request.Request(
        "https://accounts.spotify.com/api/token",
        data=urllib.parse.urlencode({"grant_type": "client_credentials"}).encode(),
        headers={"Authorization": f"Basic {creds}",
                 "Content-Type": "application/x-www-form-urlencoded"})
    with urllib.request.urlopen(req, timeout=15) as r:
        return json.loads(r.read())["access_token"]


def fetch_spotify_playlist(url: str, token: str) -> list[tuple[str, str]]:
    """Return list of (artists_normalized_joined, title_normalized)."""
    m = re.search(r"playlist[/:]([A-Za-z0-9]+)", url)
    if not m: return []
    pid = m.group(1)
    out = []
    next_url = f"https://api.spotify.com/v1/playlists/{pid}/tracks?limit=100"
    while next_url:
        req = urllib.request.Request(next_url, headers={"Authorization": f"Bearer {token}"})
        with urllib.request.urlopen(req, timeout=20) as r:
            data = json.loads(r.read())
        for item in data["items"]:
            t = item.get("track")
            if not t or t.get("is_local"): continue
            artists = [a["name"] for a in t["artists"]]
            out.append((artists, t["name"]))
        next_url = data.get("next")
    return out


def beets_dump() -> dict:
    """Return path → {grouping, artist, albumartist, album, title}"""
    proc = subprocess.run(
        ["beet", "ls", "-f",
         "$path\t$grouping\t$artist\t$albumartist\t$album\t$title"],
        capture_output=True, text=True, check=True)
    out = {}
    for line in proc.stdout.splitlines():
        parts = line.split("\t")
        if len(parts) < 6: continue
        out[parts[0]] = {
            "grouping": parts[1],
            "artist": parts[2],
            "albumartist": parts[3],
            "album": parts[4],
            "title": parts[5],
        }
    return out


def az_request(base: str, key: str, path: str, method="GET", body=None) -> dict:
    url = f"{base}{path}"
    headers = {"X-API-Key": key, "Accept": "application/json"}
    data = None
    if body is not None:
        headers["Content-Type"] = "application/json"
        data = json.dumps(body).encode()
    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    with urllib.request.urlopen(req, timeout=30) as r:
        raw = r.read()
        return json.loads(raw) if raw else {}


def az_all_files(base: str, key: str, station: int) -> list[dict]:
    out = []
    page = 1
    while True:
        d = az_request(base, key, f"/api/station/{station}/files?per_page=500&page={page}")
        out.extend(d["rows"])
        if page >= d["total_pages"]: break
        page += 1
    return out


def az_playlists_by_name(base: str, key: str, station: int) -> dict:
    return {p["name"]: p["id"]
            for p in az_request(base, key, f"/api/station/{station}/playlists")}


def write_grouping(host_path: str, grouping: str) -> bool:
    """Write GROUPING tag in-place. Return True on success."""
    ext = host_path.rsplit(".", 1)[-1].lower()
    if ext == "flac":
        subprocess.run(["metaflac", "--remove-tag=GROUPING", host_path],
                       capture_output=True, check=False)
        return subprocess.run(["metaflac", f"--set-tag=GROUPING={grouping}", host_path],
                              capture_output=True).returncode == 0
    if ext == "mp3":
        return subprocess.run(["id3v2", "--TIT1", grouping, host_path],
                              capture_output=True).returncode == 0
    return False


def beet_update(container_paths: list[str], log) -> None:
    """Re-read tags from disk into beets DB for the given container paths.
    Uses `path:<exact>` query (not `path::<regex>`) and pipes "y" since
    beet update prompts interactively for each change."""
    if not container_paths: return
    for cp in container_paths:
        subprocess.run(
            ["beet", "update", "--nomove", f"path:{cp}"],
            input="y\n", capture_output=True, text=True, check=False)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--azuracast-key", required=True)
    ap.add_argument("--azuracast-base", default=os.environ.get("AZURACAST_BASE", ""))
    ap.add_argument("--station", type=int, default=1)
    ap.add_argument("--apply", action="store_true")
    ap.add_argument("--only", default=None, help="limit recovery to one playlist name")
    args = ap.parse_args()

    print("[1/5] Loading beets metadata...", flush=True)
    beets = beets_dump()
    print(f"      {len(beets)} tracks in beets")

    print("[2/5] Fetching Spotify playlist contents...", flush=True)
    # Build (title_norm) → list of (playlist_name, artists_list)
    title_idx: dict[str, list[tuple[str, list[str]]]] = {}
    confs = sorted(CONF_DIR.glob("*.conf"))
    confs = [c for c in confs if not c.name.startswith("_")]
    if not confs:
        print(f"ERROR: no *.conf files in {CONF_DIR}/", file=sys.stderr); return 1
    first = read_conf(confs[0])
    token = spotify_token(first["spotify-id"], first["spotify-secret"])
    pl_count = 0
    for c in confs:
        cfg = read_conf(c)
        url = cfg.get("input"); pname = c.stem
        if not url or not url.startswith("https://open.spotify.com/"): continue
        try:
            tracks = fetch_spotify_playlist(url, token)
        except Exception as e:
            print(f"      WARN: {pname}: {e}"); continue
        for artists, title in tracks:
            title_idx.setdefault(normalize(title), []).append((pname, artists))
        pl_count += 1
        print(f"      {pname}: {len(tracks)} tracks")
    print(f"      indexed {pl_count} playlists, {sum(len(v) for v in title_idx.values())} title→playlist entries")

    print("[3/5] Fetching AzuraCast playlists + files...", flush=True)
    playlists_by_name = az_playlists_by_name(args.azuracast_base, args.azuracast_key, args.station)
    print(f"      {len(playlists_by_name)} AzuraCast playlists")
    all_files = az_all_files(args.azuracast_base, args.azuracast_key, args.station)
    unassigned = [r for r in all_files if not r.get("playlists")]
    print(f"      {len(all_files)} total files, {len(unassigned)} unassigned")

    print("[4/5] Matching unassigned files to playlists...", flush=True)
    matches: list[dict] = []  # rows with .target_playlist
    unmatched: list[dict] = []
    multi_match: list[dict] = []  # ambiguous
    no_az_playlist: list[dict] = []

    for r in unassigned:
        cp = f"{LIB_ROOT_CONTAINER}/{r['path']}"
        bmeta = beets.get(cp, {})
        # 1) trust beets GROUPING if present
        target = bmeta.get("grouping", "").strip() if bmeta else ""
        source = "beets_grouping" if target else None

        if not target:
            # 2) resolve via Spotify (artist, title) match
            title_n = normalize(bmeta.get("title") or r.get("title") or "")
            file_artists = bmeta.get("artist") or r.get("artist") or ""
            candidates = title_idx.get(title_n, [])
            if not candidates:
                unmatched.append({"id": r["id"], "path": r["path"], "reason": "no spotify title hit"})
                continue
            # Score candidates by artist overlap
            file_artist_n = normalize(file_artists)
            scored = []
            for pname, artists in candidates:
                joined_n = normalize("; ".join(artists))
                score = SequenceMatcher(None, file_artist_n, joined_n).ratio()
                scored.append((score, pname))
            scored.sort(reverse=True)
            top_score, top_pname = scored[0]
            if top_score < 0.4:
                unmatched.append({"id": r["id"], "path": r["path"], "reason": f"low artist score {top_score:.2f}"})
                continue
            # If second-best is within 0.05, it's ambiguous between playlists
            distinct_playlists = {p for _, p in scored}
            if len(distinct_playlists) > 1 and len(scored) > 1 and scored[1][0] > top_score - 0.05 and scored[1][1] != top_pname:
                multi_match.append({"id": r["id"], "path": r["path"], "candidates": [p for _, p in scored[:3]]})
                continue
            target = top_pname
            source = "spotify_match"

        if args.only and target != args.only:
            continue
        if target not in playlists_by_name:
            no_az_playlist.append({"id": r["id"], "path": r["path"], "target": target})
            continue

        matches.append({
            "id": r["id"], "path": r["path"], "target": target, "source": source,
            "host_path": f"{LIB_ROOT_HOST}/{r['path']}",
            "container_path": cp,
            "playlist_id": playlists_by_name[target],
        })

    # Report
    by_target: dict[str, int] = {}
    for m in matches:
        by_target[m["target"]] = by_target.get(m["target"], 0) + 1
    print()
    print(f"=== Match summary ===")
    print(f"  matched: {len(matches)}")
    for pname, n in sorted(by_target.items(), key=lambda x: -x[1]):
        print(f"    {pname}: {n}")
    print(f"  unmatched (no spotify hit): {len(unmatched)}")
    print(f"  ambiguous (multiple playlists tied): {len(multi_match)}")
    print(f"  no matching AzuraCast playlist: {len(no_az_playlist)}")
    if unmatched[:5]:
        print("  unmatched sample:")
        for u in unmatched[:5]:
            print(f"    {u['path']}  ({u['reason']})")
    if multi_match[:5]:
        print("  ambiguous sample:")
        for u in multi_match[:5]:
            print(f"    {u['path']}  candidates={u['candidates']}")

    if not args.apply:
        print()
        print("DRY RUN — pass --apply to (a) write GROUPING tags, (b) beet update, (c) assign in AzuraCast")
        return 0

    print()
    print("[5/5] Applying...", flush=True)
    tag_ok = 0; tag_fail = 0
    az_ok = 0; az_fail = 0
    for m in matches:
        # Write tag
        if write_grouping(m["host_path"], m["target"]):
            tag_ok += 1
        else:
            tag_fail += 1
            print(f"  TAG-FAIL {m['host_path']}")
        # Assign in AzuraCast: PUT /api/station/N/file/{id} with playlists field.
        # Send ONLY the playlists key — sending the full record back causes
        # AzuraCast to reject with HTTP 500 (validation fails on derived fields).
        try:
            az_request(args.azuracast_base, args.azuracast_key,
                       f"/api/station/{args.station}/file/{m['id']}",
                       method="PUT", body={"playlists": [{"id": m["playlist_id"]}]})
            az_ok += 1
        except Exception as e:
            az_fail += 1
            print(f"  AZ-FAIL {m['path']}: {e}")

    print(f"  tags written: {tag_ok}  (failed: {tag_fail})")
    print(f"  azuracast assigned: {az_ok}  (failed: {az_fail})")

    # beet update so beets DB sees the new GROUPING
    print("  running beet update on changed paths...")
    beet_update([m["container_path"] for m in matches], None)
    print("  done")
    return 0


if __name__ == "__main__":
    sys.exit(main())