219 lines
8.3 KiB
Python
219 lines
8.3 KiB
Python
"""Duplicate group resolution: choose which file to keep when consuming analysis."""
|
||
|
||
from datetime import datetime
|
||
import re
|
||
from pathlib import Path
|
||
from typing import Optional, Union
|
||
|
||
from vlm.models import MovieIdentity, SeriesIdentity
|
||
from vlm.utils import is_sample_path
|
||
|
||
|
||
class DuplicateResolutionError(ValueError):
|
||
"""Raised when a duplicate group cannot be resolved deterministically."""
|
||
|
||
|
||
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 quality, then first_seen (input order).
|
||
- by_reputation_quality_time: Prefer reputation first, then quality, then newer modified time.
|
||
- 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", "by_reputation_quality_time", "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:
|
||
raise DuplicateResolutionError(
|
||
"by_quality strategy requires quality comparison data"
|
||
)
|
||
if len(quality_comparison) != len(items):
|
||
raise DuplicateResolutionError(
|
||
"by_quality strategy requires quality data aligned with duplicate items"
|
||
)
|
||
return _by_quality_index(items, quality_comparison)
|
||
if strategy == "by_reputation":
|
||
return _by_reputation_index(items, quality_comparison=quality_comparison)
|
||
if strategy == "by_reputation_quality_time":
|
||
return _by_reputation_quality_time_index(items, quality_comparison=quality_comparison)
|
||
raise DuplicateResolutionError(f"Unsupported duplicate strategy: {strategy}")
|
||
|
||
|
||
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 quality_key(path: Path, qc: dict, idx: int) -> tuple:
|
||
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)
|
||
|
||
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 {}
|
||
is_sample = is_sample_path(path)
|
||
return (is_sample, *quality_key(path, qc, idx))
|
||
|
||
indexed = list(enumerate(items))
|
||
indexed.sort(key=key)
|
||
return indexed[0][0]
|
||
|
||
|
||
def _by_reputation_index(
|
||
items: list[tuple[Path, Union[MovieIdentity, SeriesIdentity]]],
|
||
quality_comparison: Optional[list[dict]] = None,
|
||
) -> int:
|
||
"""Sort by: sample, reputation, quality, then input order."""
|
||
quality_list = quality_comparison or []
|
||
|
||
def quality_key(path: Path, qc: dict, idx: int) -> tuple:
|
||
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)
|
||
# Reputation fallback: prioritize source first (e.g. BluRay over WEB-DL).
|
||
return (-src_tier, -res_tier, -codec_tier, -size, idx)
|
||
|
||
def key(idx_reason: tuple[int, tuple[Path, Union[MovieIdentity, SeriesIdentity]]]) -> tuple:
|
||
idx, (path, identity) = idx_reason
|
||
qc = quality_list[idx] if idx < len(quality_list) else {}
|
||
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
|
||
is_sample = is_sample_path(path)
|
||
# Prefer non-sample, then reputation, then better quality, then lower index.
|
||
return (is_sample, not has_rep, -score, -votes, *quality_key(path, qc, idx))
|
||
|
||
indexed = list(enumerate(items))
|
||
indexed.sort(key=key)
|
||
return indexed[0][0]
|
||
|
||
|
||
def _parse_modified_timestamp(value: Optional[str]) -> float:
|
||
"""Parse ISO modified timestamp to unix epoch seconds; unknown returns 0."""
|
||
if not value:
|
||
return 0.0
|
||
try:
|
||
normalized = value[:-1] + "+00:00" if value.endswith("Z") else value
|
||
return datetime.fromisoformat(normalized).timestamp()
|
||
except (ValueError, TypeError):
|
||
return 0.0
|
||
|
||
|
||
def _by_reputation_quality_time_index(
|
||
items: list[tuple[Path, Union[MovieIdentity, SeriesIdentity]]],
|
||
quality_comparison: Optional[list[dict]] = None,
|
||
) -> int:
|
||
"""Sort by: sample, reputation, quality, modified time (newer first), then order."""
|
||
quality_list = quality_comparison or []
|
||
|
||
def quality_key(path: Path, qc: dict) -> tuple:
|
||
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)
|
||
|
||
def key(idx_reason: tuple[int, tuple[Path, Union[MovieIdentity, SeriesIdentity]]]) -> tuple:
|
||
idx, (path, identity) = idx_reason
|
||
qc = quality_list[idx] if idx < len(quality_list) else {}
|
||
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
|
||
modified_ts = _parse_modified_timestamp(qc.get("modified_timestamp"))
|
||
is_sample = is_sample_path(path)
|
||
return (
|
||
is_sample,
|
||
not has_rep,
|
||
-score,
|
||
-votes,
|
||
*quality_key(path, qc),
|
||
-modified_ts,
|
||
idx,
|
||
)
|
||
|
||
indexed = list(enumerate(items))
|
||
indexed.sort(key=key)
|
||
return indexed[0][0]
|