#!/usr/bin/env python3
"""
find-fuzzy-dupes.py — find acoustic duplicates in the beets library using
the existing $ALEMBIC_CONFIG_DIR/beets/fingerprints.db (Chromaprint).

Pairs whose duration differs by <=5s and Chromaprint similarity >= 0.92 are
flagged as duplicates. Groups are transitively merged. Within each group the
ranking is: FLAC > MP3/other, then largest size wins.

DRY RUN by default. Pass --apply to delete losers via `beet remove -d -f`.

Usage:
    find-fuzzy-dupes.py
    find-fuzzy-dupes.py --apply
    find-fuzzy-dupes.py --threshold 0.95 --apply
"""

import argparse
import json
import os
import re
import sqlite3
import subprocess
import sys
import time
from collections import defaultdict

import numpy as np

FINGERPRINT_DB = f"{os.environ.get('ALEMBIC_CONFIG_DIR', '/config')}/beets/fingerprints.db"
SCAN_CACHE = "/tmp/find-fuzzy-dupes-scan-cache.json"
DEFAULT_THRESHOLD = 0.92
DURATION_TOLERANCE = 5  # seconds
# Library dedup compares same-track recordings — offsets are tiny. Use a
# smaller window than DJ-import (±80) so a 4k-file library scans in minutes.
MAX_OFFSET = 15


# Precomputed bitmasks for popcount (reused across compares).
_M1 = np.uint32(0x55555555)
_M2 = np.uint32(0x33333333)
_M4 = np.uint32(0x0F0F0F0F)
_H01 = np.uint32(0x01010101)


def _sim_at_offset(fp1: np.ndarray, fp2: np.ndarray, offset: int) -> float:
    if offset >= 0:
        n = min(fp1.size - offset, fp2.size)
        if n < 100:
            return 0.0
        xor = np.bitwise_xor(fp1[offset : offset + n], fp2[:n])
    else:
        n = min(fp1.size, fp2.size + offset)
        if n < 100:
            return 0.0
        xor = np.bitwise_xor(fp1[:n], fp2[-offset : -offset + n])
    v = xor - ((xor >> 1) & _M1)
    v = (v & _M2) + ((v >> 2) & _M2)
    v = (v + (v >> 4)) & _M4
    diff_bits = int(((v * _H01) >> 24).sum())
    return 1.0 - diff_bits / (n * 32.0)


def _compare_fingerprints(fp1: np.ndarray, fp2: np.ndarray, max_offset: int = MAX_OFFSET) -> float:
    if fp1.size < 50 or fp2.size < 50:
        return 0.0
    # Fast pre-filter: try offset 0 first. If unrelated, skip the offset
    # search entirely. Same recordings cluster within a few offsets of 0,
    # so a sub-0.50 base score is a definitive "not a duplicate."
    base = _sim_at_offset(fp1, fp2, 0)
    if base < 0.50:
        return base
    best = base
    if best >= 0.99:
        return best
    for offset in range(1, max_offset + 1):
        for o in (offset, -offset):
            sim = _sim_at_offset(fp1, fp2, o)
            if sim > best:
                best = sim
                if best >= 0.99:
                    return best
    return best


_DESKTOP_RE = re.compile(r"-DESKTOP-[A-Z0-9]+", re.IGNORECASE)
_DUP_SUFFIX_RE = re.compile(r"\(\d+\)\.[a-z]+$|\.\d+\.[a-z]+$", re.IGNORECASE)
# DJ/club-friendly versions to prefer when same audio matches at different
# labellings. Bonus is small enough not to override format/size when those
# signal a real quality difference.
_EXTENDED_RE = re.compile(
    r"\b(extended|club|dj edit|dj mix|long version|original mix)\b",
    re.IGNORECASE,
)


def rank_file(host_path: str) -> int:
    """Lower is better. FLAC > non-FLAC; clean filename > sync-conflict
    artifact; extended/club mixes beat radio edits at same format; then
    larger wins. Penalty weights:
      DESKTOP-suffix (OneDrive sync conflict)  > (N)/.N. (Windows/beets dup)
    so an exclusively-(1)-marked file beats one that's both (1) and DESKTOP."""
    try:
        size = os.path.getsize(host_path)
    except OSError:
        size = 0
    ext_score = 1 if host_path.lower().endswith(".flac") else 0
    desktop = 1 if _DESKTOP_RE.search(host_path) else 0
    dup = 1 if _DUP_SUFFIX_RE.search(host_path) else 0
    extended = 1 if _EXTENDED_RE.search(host_path) else 0
    return (
        10_000_000_000
        - ext_score * 1_000_000_000_000
        + desktop * 500_000_000_000
        + dup * 100_000_000_000
        - extended * 50_000_000_000
        - size
    )


class UnionFind:
    def __init__(self, ids):
        self.parent = {i: i for i in ids}

    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra != rb:
            self.parent[ra] = rb


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--apply", action="store_true", help="actually delete losers")
    ap.add_argument("--threshold", type=float, default=DEFAULT_THRESHOLD)
    ap.add_argument("--duration-tol", type=int, default=DURATION_TOLERANCE)
    ap.add_argument("--no-cache", action="store_true",
                    help="ignore the scan cache and re-fingerprint everything")
    args = ap.parse_args()

    if not os.path.exists(FINGERPRINT_DB):
        print(f"ERROR: fingerprint index missing at {FINGERPRINT_DB}", file=sys.stderr)
        print("Run build-fingerprint-index.py first.", file=sys.stderr)
        return 2

    conn = sqlite3.connect(f"file:{FINGERPRINT_DB}?mode=ro", uri=True)
    rows = list(conn.execute("SELECT beets_id, path, duration, fingerprint FROM fingerprints"))
    conn.close()

    entries = []
    for beets_id, path, duration, fp_csv in rows:
        fp = np.fromstring(fp_csv, dtype=np.uint32, sep=",")
        if fp.size == 0:
            continue
        if not os.path.exists(path):
            continue
        entries.append((beets_id, path, duration, fp))
    print(f"[fuzzy-dupes] loaded {len(entries)} fingerprints (skipped missing/empty)")

    # Bucket by duration to make pairwise scan O(N * window) instead of O(N^2).
    by_dur = defaultdict(list)
    for idx, (_, _, dur, _) in enumerate(entries):
        by_dur[dur].append(idx)

    # Try cache first. Cache key = (fingerprint mtime, threshold, duration_tol).
    fp_mtime = os.path.getmtime(FINGERPRINT_DB)
    cache_key = {"fp_mtime": fp_mtime, "threshold": args.threshold,
                 "duration_tol": args.duration_tol, "n_entries": len(entries)}
    pair_scores: dict[tuple[int, int], float] = {}
    if not args.no_cache and os.path.exists(SCAN_CACHE):
        try:
            with open(SCAN_CACHE) as f:
                cache = json.load(f)
            if cache.get("key") == cache_key:
                pair_scores = {tuple(map(int, k.split(","))): v
                               for k, v in cache["pairs"].items()}
                print(f"[fuzzy-dupes] loaded {len(pair_scores)} pair scores from cache "
                      f"({SCAN_CACHE})")
        except (OSError, json.JSONDecodeError, KeyError):
            pair_scores = {}

    if not pair_scores:
        start = time.time()
        compared = 0
        matched = 0
        for i, (_, path_i, dur_i, fp_i) in enumerate(entries):
            for d in range(dur_i - args.duration_tol, dur_i + args.duration_tol + 1):
                for j in by_dur.get(d, []):
                    if j <= i:
                        continue
                    compared += 1
                    fp_j = entries[j][3]
                    score = _compare_fingerprints(fp_i, fp_j)
                    if score >= args.threshold:
                        pair_scores[(i, j)] = score
                        matched += 1
            if (i + 1) % 250 == 0:
                elapsed = time.time() - start
                rate = (i + 1) / elapsed if elapsed > 0 else 0
                print(f"[fuzzy-dupes] scanned {i + 1}/{len(entries)} "
                      f"({rate:.1f}/s, {compared} compared, {matched} matches)", flush=True)
        try:
            with open(SCAN_CACHE, "w") as f:
                json.dump({
                    "key": cache_key,
                    "pairs": {f"{i},{j}": s for (i, j), s in pair_scores.items()},
                }, f)
            print(f"[fuzzy-dupes] cached {len(pair_scores)} pair scores → {SCAN_CACHE}")
        except OSError as e:
            print(f"[fuzzy-dupes] WARN: could not write cache: {e}", file=sys.stderr)

    uf = UnionFind(range(len(entries)))
    for (i, j) in pair_scores:
        uf.union(i, j)

    # Group by union-find root.
    groups: dict[int, list[int]] = defaultdict(list)
    for idx in range(len(entries)):
        groups[uf.find(idx)].append(idx)

    dup_groups = [g for g in groups.values() if len(g) >= 2]
    print(f"\n[fuzzy-dupes] found {len(dup_groups)} duplicate groups "
          f"covering {sum(len(g) for g in dup_groups)} files\n")

    if not dup_groups:
        return 0

    delete_targets: list[tuple[int, str]] = []
    skipped_transitive = 0
    for group in sorted(dup_groups, key=lambda g: -len(g)):
        ranked = sorted(group, key=lambda idx: rank_file(entries[idx][1]))
        keeper_idx = ranked[0]
        keeper_path = entries[keeper_idx][1]
        keeper_fp = entries[keeper_idx][3]
        keeper_size = os.path.getsize(keeper_path)
        keeper_ext = keeper_path.rsplit(".", 1)[-1].upper()
        print(f"--- group ({len(group)} files) ---")
        print(f"  KEEP   [{keeper_ext} {keeper_size / 1024 / 1024:.1f}M] {keeper_path}")
        for loser_idx in ranked[1:]:
            beets_id, loser_path, _, loser_fp = entries[loser_idx]
            loser_size = os.path.getsize(loser_path) if os.path.exists(loser_path) else 0
            loser_ext = loser_path.rsplit(".", 1)[-1].upper()
            # Verify direct similarity to the keeper. Transitive union-find can
            # pull in unrelated tracks via a chain of partial matches (e.g. a
            # game-OST jingle group where A~B and B~C but A and C are unrelated).
            pair_key = (min(keeper_idx, loser_idx), max(keeper_idx, loser_idx))
            score = pair_scores.get(pair_key)
            if score is None:
                score = _compare_fingerprints(keeper_fp, loser_fp)
            if score < args.threshold:
                print(f"  SKIP   [{loser_ext} {loser_size / 1024 / 1024:.1f}M sim={score:.3f} to keeper] {loser_path}")
                skipped_transitive += 1
                continue
            print(f"  DELETE [{loser_ext} {loser_size / 1024 / 1024:.1f}M sim={score:.3f}] {loser_path}")
            delete_targets.append((beets_id, loser_path))
        print()
    if skipped_transitive:
        print(f"[fuzzy-dupes] skipped {skipped_transitive} transitive false-positives "
              f"(chain-grouped but sim<{args.threshold} to keeper)\n")

    if not args.apply:
        print(f"\n[fuzzy-dupes] DRY RUN — pass --apply to delete {len(delete_targets)} files")
        return 0

    print(f"\n[fuzzy-dupes] APPLY: deleting {len(delete_targets)} files via beet remove -d -f")
    failed = 0
    for beets_id, path in delete_targets:
        # Escape regex metacharacters for path:: regex query
        escaped = re.escape(path)
        result = subprocess.run(
            ["beet", "remove", "-d", "-f", f"path::{escaped}"],
            capture_output=True, text=True,
        )
        if result.returncode != 0:
            print(f"  FAILED: {path} — {result.stderr.strip()}", file=sys.stderr)
            failed += 1
    print(f"[fuzzy-dupes] done. {len(delete_targets) - failed} deleted, {failed} failed.")
    return 1 if failed else 0


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