dl-organizer/src/vlm/plan_review.py

"""Plan review helpers for flagging high-risk operations before execution."""

from __future__ import annotations

import csv
import random
import re
from pathlib import Path
from typing import Any

from vlm.models import ExecutionPlan, FileOperation, MovieIdentity, SeriesIdentity
from vlm.review_display import display_path
from vlm.utils import format_size, is_sample_path

REVIEW_CSV_BASE_FIELDS = [
    "index",
    "operation_type",
    "risk_flags",
    "source_path",
    "destination_path",
    "reason",
]

REVIEW_CSV_ENRICHED_FIELDS = [
    "title",
    "category",
    "season",
    "episode",
    "source_name",
    "dest_name",
    "rel_source",
    "rel_dest",
    "quality_hint",
    "duplicate_group_id",
]

REVIEW_CSV_ALL_FIELDS = REVIEW_CSV_BASE_FIELDS + REVIEW_CSV_ENRICHED_FIELDS

GROUP_BY_CHOICES = ("none", "reason", "title", "duplicate")


def _extract_season_episode(operation: FileOperation) -> tuple[int | None, int | None]:
    season: int | None = None
    episode: int | None = None
    for text in [
        str(operation.destination_path) if operation.destination_path else "",
        operation.reason,
        operation.source_path.name,
    ]:
        if season is None:
            m = re.search(r"S(\d{1,3})E(\d{1,3})", text, re.IGNORECASE)
            if m:
                season = int(m.group(1))
                episode = int(m.group(2))
            else:
                m = re.search(r"Season\s+(\d{1,3})", text, re.IGNORECASE)
                if m:
                    season = int(m.group(1))
        if episode is None:
            m = re.search(r"E(\d{1,3})", text, re.IGNORECASE)
            if m:
                episode = int(m.group(1))
    return season, episode


def _quality_hint_from_record(record: dict[str, Any]) -> str:
    vm = record.get("video_metadata") or {}
    if not isinstance(vm, dict):
        vm = {}
    parts: list[str] = []
    if vm.get("resolution"):
        parts.append(str(vm["resolution"]))
    if vm.get("codec"):
        parts.append(str(vm["codec"]))
    size = vm.get("size_bytes")
    if size:
        parts.append(format_size(int(size)))
    return " / ".join(parts)


def build_identity_lookup(identities_data: dict[str, Any]) -> dict[str, dict[str, str]]:
    """Map source file path strings to review enrichment fields."""
    lookup: dict[str, dict[str, str]] = {}

    def _add_record(record: dict[str, Any], category: str) -> None:
        path = record.get("path")
        if not path:
            return
        path_key = str(path)
        title = str(record.get("display_title") or record.get("title") or "")
        ctx: dict[str, str] = {
            "title": title,
            "category": category,
            "quality_hint": _quality_hint_from_record(record),
        }
        if category == "movie":
            year = record.get("year")
            if year is not None:
                ctx["episode"] = ""
                ctx["season"] = ""
        elif category == "series":
            season = record.get("season")
            episodes = record.get("episodes") or []
            if season is not None:
                ctx["season"] = str(season)
            if episodes:
                ctx["episode"] = ",".join(str(e) for e in episodes)
        lookup[path_key] = ctx
        try:
            lookup[str(Path(path_key).resolve())] = ctx
        except OSError:
            pass

    for rec in identities_data.get("movies", []) or []:
        if isinstance(rec, dict):
            _add_record(rec, "movie")
    for rec in identities_data.get("series", []) or []:
        if isinstance(rec, dict):
            _add_record(rec, "series")
    for section in ("anime", "other"):
        for rec in identities_data.get(section, []) or []:
            if isinstance(rec, dict):
                _add_record(rec, section)

    return lookup


def build_duplicate_path_maps(
    analysis_data: dict[str, Any] | None,
) -> tuple[dict[str, str], dict[str, dict[str, Any]]]:
    """Return path -> duplicate_group_id and path -> quality_comparison entry."""
    path_to_group: dict[str, str] = {}
    path_to_quality: dict[str, dict[str, Any]] = {}
    if not analysis_data:
        return path_to_group, path_to_quality

    for group_idx, dup in enumerate(analysis_data.get("duplicates", []) or []):
        if not isinstance(dup, dict):
            continue
        identity = dup.get("identity") or {}
        if not isinstance(identity, dict):
            identity = {}
        title = identity.get("title", "")
        itype = identity.get("type", "")
        year = identity.get("year")
        season = identity.get("season")
        episodes = identity.get("episodes") or []
        if itype == "movie" and year is not None:
            group_id = f"movie:{title}:{year}"
        elif season is not None and episodes:
            group_id = f"series:{title}:{season}:{episodes[0]}"
        else:
            group_id = f"dup:{group_idx}"

        quality_list = dup.get("quality_comparison") or []
        if not isinstance(quality_list, list):
            quality_list = []
        for file_path in dup.get("files", []) or []:
            path_key = str(file_path)
            path_to_group[path_key] = group_id
            for qc in quality_list:
                if isinstance(qc, dict) and str(qc.get("path")) == path_key:
                    path_to_quality[path_key] = qc
                    break

    return path_to_group, path_to_quality


def build_review_context(
    video_file_path: Path,
    identity: MovieIdentity | SeriesIdentity | None,
    *,
    category: str | None = None,
    duplicate_group_id: str = "",
    keep_candidate: bool | None = None,
) -> dict[str, Any]:
    """Build plan-time review_context for a single operation."""
    ctx: dict[str, Any] = {}
    if category:
        ctx["category"] = category
    elif identity is not None:
        ctx["category"] = "movie" if isinstance(identity, MovieIdentity) else "series"

    if identity is not None:
        ctx["title"] = identity.title
        if isinstance(identity, MovieIdentity):
            if identity.year is not None:
                ctx["year"] = identity.year
        else:
            if identity.season is not None:
                ctx["season"] = identity.season
            if identity.episodes:
                ctx["episode"] = identity.episodes[0]

    if duplicate_group_id:
        ctx["duplicate_group_id"] = duplicate_group_id
    if keep_candidate is not None:
        ctx["keep_candidate"] = keep_candidate

    return ctx


def enrich_review_row(
    row: dict[str, str],
    *,
    library_root: Path | None = None,
    identity_lookup: dict[str, dict[str, str]] | None = None,
    path_to_duplicate_group: dict[str, str] | None = None,
    operation: FileOperation | None = None,
) -> dict[str, str]:
    """Add optional enrichment columns to a review row dict."""
    enriched = dict(row)
    source_path = row.get("source_path", "")
    dest_path = row.get("destination_path", "")

    enriched["source_name"] = Path(source_path).name if source_path else ""
    enriched["dest_name"] = Path(dest_path).name if dest_path else ""

    if library_root is not None and source_path:
        enriched["rel_source"] = display_path(Path(source_path), library_root)
    if library_root is not None and dest_path:
        enriched["rel_dest"] = display_path(Path(dest_path), library_root)

    ctx = (operation.review_context if operation else None) or {}
    if isinstance(ctx, dict) and ctx:
        if ctx.get("title"):
            enriched["title"] = str(ctx["title"])
        if ctx.get("category"):
            enriched["category"] = str(ctx["category"])
        if ctx.get("season") is not None:
            enriched["season"] = str(ctx["season"])
        if ctx.get("episode") is not None:
            enriched["episode"] = str(ctx["episode"])
        if ctx.get("duplicate_group_id"):
            enriched["duplicate_group_id"] = str(ctx["duplicate_group_id"])

    if identity_lookup and source_path:
        id_ctx = identity_lookup.get(source_path) or identity_lookup.get(
            str(Path(source_path).resolve())
        )
        if id_ctx:
            for key in ("title", "category", "season", "episode", "quality_hint"):
                if id_ctx.get(key) and not enriched.get(key):
                    enriched[key] = id_ctx[key]

    if path_to_duplicate_group and source_path:
        gid = path_to_duplicate_group.get(source_path)
        if gid and not enriched.get("duplicate_group_id"):
            enriched["duplicate_group_id"] = gid

    return enriched


def enrich_review_rows(
    rows: list[dict[str, str]],
    plan: ExecutionPlan,
    *,
    library_root: Path | None = None,
    identity_lookup: dict[str, dict[str, str]] | None = None,
    path_to_duplicate_group: dict[str, str] | None = None,
) -> list[dict[str, str]]:
    """Enrich all review rows using plan operations and optional lookups."""
    op_by_index = {i + 1: op for i, op in enumerate(plan.operations)}
    result: list[dict[str, str]] = []
    for row in rows:
        idx = int(row["index"])
        op = op_by_index.get(idx)
        result.append(
            enrich_review_row(
                row,
                library_root=library_root,
                identity_lookup=identity_lookup,
                path_to_duplicate_group=path_to_duplicate_group,
                operation=op,
            )
        )
    return result


def review_plan(
    plan: ExecutionPlan,
    season_threshold: int = 20,
    episode_threshold: int = 40,
) -> tuple[list[dict[str, str]], dict[str, int]]:
    """Build review rows and aggregate risk counters."""
    rows: list[dict[str, str]] = []
    counters = {
        "total_operations": len(plan.operations),
        "high_risk_operations": 0,
        "manual_review": 0,
        "sample_source": 0,
        "high_season": 0,
        "high_episode": 0,
        "conflicts": 0,
    }

    for idx, op in enumerate(plan.operations, start=1):
        flags: list[str] = []
        reason_l = op.reason.casefold()
        if "manual review" in reason_l:
            flags.append("manual_review")
            counters["manual_review"] += 1
        if is_sample_path(op.source_path):
            flags.append("sample_source")
            counters["sample_source"] += 1
        season, episode = _extract_season_episode(op)
        if season is not None and season >= season_threshold:
            flags.append("high_season")
            counters["high_season"] += 1
        if episode is not None and episode >= episode_threshold:
            flags.append("high_episode")
            counters["high_episode"] += 1
        if op.has_conflict:
            flags.append("conflict")
            counters["conflicts"] += 1
        rc = op.review_context if hasattr(op, "review_context") else {}
        if isinstance(rc, dict) and rc.get("duplicate_group_id") and op.operation_type == "quarantine":
            if "duplicate" not in flags:
                flags.append("duplicate")
        elif "duplicate" in reason_l and "duplicate" not in flags:
            flags.append("duplicate")
        if flags:
            counters["high_risk_operations"] += 1
            rows.append(
                {
                    "index": str(idx),
                    "operation_type": op.operation_type,
                    "risk_flags": "|".join(flags),
                    "source_path": str(op.source_path),
                    "destination_path": str(op.destination_path) if op.destination_path else "",
                    "reason": op.reason,
                }
            )

    return rows, counters


def append_sample_safe_rows(
    plan: ExecutionPlan,
    rows: list[dict[str, str]],
    sample_count: int,
    *,
    rng: random.Random | None = None,
) -> list[dict[str, str]]:
    """Append random non-high-risk move operations for spot-checking."""
    if sample_count < 1:
        return rows

    high_risk_indices = {int(r["index"]) for r in rows}
    candidates: list[tuple[int, FileOperation]] = []
    for idx, op in enumerate(plan.operations, start=1):
        if idx in high_risk_indices:
            continue
        if op.operation_type != "move":
            continue
        candidates.append((idx, op))

    if not candidates:
        return rows

    picker = rng or random.Random()
    picked = picker.sample(candidates, min(sample_count, len(candidates)))
    extra: list[dict[str, str]] = []
    for idx, op in sorted(picked, key=lambda x: x[0]):
        extra.append(
            {
                "index": str(idx),
                "operation_type": op.operation_type,
                "risk_flags": "spot_check",
                "source_path": str(op.source_path),
                "destination_path": str(op.destination_path) if op.destination_path else "",
                "reason": op.reason,
            }
        )
    return rows + extra


def group_review_rows(
    rows: list[dict[str, str]],
    group_by: str,
) -> list[dict[str, str]]:
    """Reorder review rows for display/export grouping (does not change plan)."""
    if group_by == "none" or not rows:
        return rows

    def sort_key(row: dict[str, str]) -> tuple:
        if group_by == "reason":
            return (row.get("reason", ""), int(row.get("index", 0)))
        if group_by == "title":
            return (row.get("title", ""), row.get("duplicate_group_id", ""), int(row.get("index", 0)))
        if group_by == "duplicate":
            gid = row.get("duplicate_group_id", "")
            if not gid:
                gid = f"_ungrouped_{row.get('index', '')}"
            return (gid, int(row.get("index", 0)))
        return (int(row.get("index", 0)),)

    return sorted(rows, key=sort_key)


def check_review_requirements(
    plan: ExecutionPlan,
    plan_path: Path,
    review_csv: Path,
    *,
    season_threshold: int = 20,
    episode_threshold: int = 40,
) -> list[str]:
    """Return human-readable errors when review is required but missing or stale."""
    _, counters = review_plan(
        plan,
        season_threshold=season_threshold,
        episode_threshold=episode_threshold,
    )
    high_risk = counters.get("high_risk_operations", 0)
    if high_risk == 0:
        return []

    errors: list[str] = []
    if not review_csv.is_file():
        errors.append(
            f"Review required: {high_risk} high-risk operation(s) but review CSV not found: {review_csv}. "
            f"Run: vlm review-plan --input {plan_path}"
        )
        return errors

    try:
        if review_csv.stat().st_mtime < plan_path.stat().st_mtime:
            errors.append(
                f"Review CSV is older than plan ({review_csv}). "
                f"Re-run review-plan and apply-review before execute --confirm."
            )
    except OSError as exc:
        errors.append(f"Could not compare review CSV timestamps: {exc}")

    return errors


def prepare_review_rows(
    plan: ExecutionPlan,
    *,
    season_threshold: int = 20,
    episode_threshold: int = 40,
    library_root: Path | None = None,
    identities_data: dict[str, Any] | None = None,
    analysis_data: dict[str, Any] | None = None,
    sample_safe: int = 0,
    group_by: str = "none",
) -> tuple[list[dict[str, str]], dict[str, int]]:
    """Build, enrich, optionally sample and group review rows."""
    rows, counters = review_plan(
        plan,
        season_threshold=season_threshold,
        episode_threshold=episode_threshold,
    )
    if sample_safe > 0:
        rows = append_sample_safe_rows(plan, rows, sample_safe)

    identity_lookup = (
        build_identity_lookup(identities_data) if identities_data else None
    )
    path_to_dup_group, _ = build_duplicate_path_maps(analysis_data)
    rows = enrich_review_rows(
        rows,
        plan,
        library_root=library_root,
        identity_lookup=identity_lookup,
        path_to_duplicate_group=path_to_dup_group,
    )
    if group_by and group_by != "none":
        rows = group_review_rows(rows, group_by)
    return rows, counters


def save_review_csv(rows: list[dict[str, str]], output: Path) -> None:
    """Write review rows to CSV."""
    output.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = list(REVIEW_CSV_BASE_FIELDS)
    for row in rows:
        for key in row:
            if key in REVIEW_CSV_ENRICHED_FIELDS and key not in fieldnames:
                fieldnames.append(key)
    for f in REVIEW_CSV_ENRICHED_FIELDS:
        if f in fieldnames:
            continue
        if any(f in row for row in rows):
            fieldnames.append(f)

    with open(output, "w", encoding="utf-8", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=fieldnames, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(rows)