dl-organizer/src/vlm/duplicate_resolve.py

"""Duplicate group resolution: choose which file to keep when consuming analysis."""

import re
from pathlib import Path
from typing import Optional, Union

from vlm.models import MovieIdentity, SeriesIdentity


def choose_keep_index(
    items: list[tuple[Path, Union[MovieIdentity, SeriesIdentity]]],
    strategy: str,
    quality_comparison: Optional[list[dict]] = None,
) -> int | None:
    """Choose the index of the item to keep in a duplicate group.

    Strategies:
    - by_quality: Prefer higher quality (resolution > source > codec > size). Requires quality_comparison.
    - by_reputation: Prefer items with reputation_score; then sort by score desc,
      then reputation_votes desc, then first_seen (input order). Keep index 0 after sort.
    - first_seen: Keep the first item (index 0).
    - manual: Return None; caller should not generate quarantine ops, only record in metadata.

    Args:
        items: List of (path, identity) for the duplicate group.
        strategy: One of "by_quality", "by_reputation", "first_seen", "manual".
        quality_comparison: List of quality dicts (filename, path, size_bytes, resolution?, codec?)
            aligned with items. Required when strategy is "by_quality".

    Returns:
        Index in [0, len(items)) to keep, or None for manual.
    """
    if strategy == "manual":
        return None
    if strategy == "first_seen":
        return 0
    if strategy == "by_quality":
        if quality_comparison is None or len(quality_comparison) != len(items):
            return 0  # Fallback to first if quality data missing/mismatched
        return _by_quality_index(items, quality_comparison)
    if strategy == "by_reputation":
        return _by_reputation_index(items)
    return 0


def _parse_resolution_tier(resolution: Optional[str], path: Path) -> int:
    """Parse resolution to tier. Higher = better. 2160=4, 1080=3, 720=2, 0=unknown."""
    if resolution:
        m = re.search(r"(\d{3,4})[x×](\d{3,4})", resolution, re.I)
        if m:
            h = int(m.group(2))
            if h >= 2160:
                return 4
            if h >= 1080:
                return 3
            if h >= 720:
                return 2
            return 1
    s = path.name.upper()
    if "2160" in s or "4K" in s or "UHD" in s:
        return 4
    if "1080" in s:
        return 3
    if "720" in s:
        return 2
    if "480" in s or "576" in s:
        return 1
    return 0


def _parse_source_tier(path: Path) -> int:
    """Parse source type from path. BluRay=3, WEB-DL=2, HDTV=1, 0=other."""
    s = path.name.upper()
    if "BLURAY" in s or "BLU-RAY" in s or "BD" in s:
        return 3
    if "WEB-DL" in s or "WEBRIP" in s or "WEB" in s:
        return 2
    if "HDTV" in s or "HDTVRIP" in s:
        return 1
    return 0


def _parse_codec_tier(codec: Optional[str], path: Path) -> int:
    """Parse codec to tier. hevc/x265=2, h264/x264=1, 0=other."""
    if codec:
        c = codec.lower()
        if "hevc" in c or "h265" in c or "x265" in c:
            return 2
        if "h264" in c or "x264" in c or "avc" in c:
            return 1
    s = path.name.upper()
    if "X265" in s or "H265" in s or "HEVC" in s:
        return 2
    if "X264" in s or "H264" in s or "AVC" in s:
        return 1
    return 0


def _by_quality_index(
    items: list[tuple[Path, Union[MovieIdentity, SeriesIdentity]]],
    quality_comparison: list[dict],
) -> int:
    """Sort by: resolution desc, source desc, codec desc, size desc, index asc. Return best index."""
    def key(idx_item: tuple[int, tuple[Path, Union[MovieIdentity, SeriesIdentity]]]) -> tuple:
        idx, (path, _) = idx_item
        qc = quality_comparison[idx] if idx < len(quality_comparison) else {}
        resolution = qc.get("resolution")
        codec = qc.get("codec")
        size = qc.get("size_bytes", 0) or 0
        res_tier = _parse_resolution_tier(resolution, path)
        src_tier = _parse_source_tier(path)
        codec_tier = _parse_codec_tier(codec, path)
        return (-res_tier, -src_tier, -codec_tier, -size, idx)

    indexed = list(enumerate(items))
    indexed.sort(key=key)
    return indexed[0][0]


def _by_reputation_index(
    items: list[tuple[Path, Union[MovieIdentity, SeriesIdentity]]],
) -> int:
    """Sort by: has reputation > no reputation; then score desc; then votes desc; then order. Return 0."""
    def key(idx_reason: tuple[int, tuple[Path, Union[MovieIdentity, SeriesIdentity]]]) -> tuple:
        idx, (path, identity) = idx_reason
        has_rep = identity.reputation_score is not None
        score = identity.reputation_score if identity.reputation_score is not None else -1.0
        votes = identity.reputation_votes if identity.reputation_votes is not None else -1
        # Prefer has reputation (True > False), then higher score, then higher votes, then lower index
        return (not has_rep, -score, -votes, idx)

    indexed = list(enumerate(items))
    indexed.sort(key=key)
    return indexed[0][0]