dl-organizer/src/vlm/plan_review.py

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

from __future__ import annotations

import csv
import re
from pathlib import Path

from vlm.models import ExecutionPlan, FileOperation
from vlm.utils import is_sample_path


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 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
        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 save_review_csv(rows: list[dict[str, str]], output: Path) -> None:
    """Write review rows to CSV."""
    output.parent.mkdir(parents=True, exist_ok=True)
    fields = ["index", "operation_type", "risk_flags", "source_path", "destination_path", "reason"]
    with open(output, "w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)