dl-organizer/src/vlm/planner.py

"""Plan generator for creating execution plans from parsed identities.

This module generates structured execution plans that specify how video files
should be organized based on their parsed identities and configuration templates.
"""

import json
import uuid
from datetime import datetime
from pathlib import Path
from typing import Optional, Union

from vlm.config import Config
from vlm.duplicate_resolve import choose_keep_index
from vlm.utils import ensure_utc, utc_now
from vlm.models import (
    ExecutionPlan,
    FileOperation,
    MovieIdentity,
    SeriesIdentity,
    VideoFile,
)

QUARANTINE_REASON_DUPLICATE = "重复项（已按外部评分保留一条）"
QUARANTINE_REASON_DUPLICATE_BY_QUALITY = "重复项（已按画质保留最优）"


def generate_plan(
    identities: list[tuple[VideoFile, Union[MovieIdentity, SeriesIdentity, None]]],
    config: Config,
    analysis_data: Optional[dict] = None,
) -> ExecutionPlan:
    """Generate an execution plan from parsed identities; optionally apply analysis duplicates.

    When analysis_data is provided and duplicate_keep is not "manual", duplicate groups
    are resolved (one kept, rest quarantined) according to config.duplicate_keep.
    """
    operations = []
    for video_file, identity in identities:
        operation = _create_operation(video_file, identity, config)
        operations.append(operation)

    metadata: dict = {}
    if analysis_data is not None:
        metadata["analysis_source"] = analysis_data.get("metadata", {}).get("source_identities", "")
        metadata["duplicate_groups_considered"] = len(analysis_data.get("duplicates", []))
        metadata["completeness_seasons_with_gaps"] = len(analysis_data.get("completeness", []))

        if config.duplicate_keep != "manual":
            path_to_index = {str(vf.path): i for i, (vf, _) in enumerate(identities)}
            for dup in analysis_data.get("duplicates", []):
                paths = dup.get("files", [])
                indices = [path_to_index[p] for p in paths if p in path_to_index]
                path_to_qc = {qc.get("path"): qc for qc in dup.get("quality_comparison", [])}
                items = []
                valid_indices = []
                for i in indices:
                    identity = identities[i][1]
                    if identity is not None and isinstance(
                        identity, (MovieIdentity, SeriesIdentity)
                    ):
                        items.append((identities[i][0].path, identity))
                        valid_indices.append(i)
                if not items:
                    continue
                quality_list = [
                    path_to_qc.get(str(p), {}) for p, _ in items
                ]
                keep_idx = choose_keep_index(
                    items, config.duplicate_keep, quality_comparison=quality_list
                )
                if keep_idx is None:
                    continue
                keep_identity_index = valid_indices[keep_idx]
                quarantine_indices = set(valid_indices) - {keep_identity_index}
                reason = (
                    QUARANTINE_REASON_DUPLICATE_BY_QUALITY
                    if config.duplicate_keep == "by_quality"
                    else QUARANTINE_REASON_DUPLICATE
                )
                for i in quarantine_indices:
                    vf = identities[i][0]
                    operations[i] = FileOperation(
                        operation_type="quarantine",
                        source_path=vf.path,
                        destination_path=None,
                        reason=reason,
                        has_conflict=False,
                        conflict_reason=None,
                    )

    summary = _generate_summary(operations)
    summary_by_reason = _generate_summary_by_reason(operations)
    human_summary = _generate_human_summary(operations, summary, summary_by_reason, metadata)

    return ExecutionPlan(
        plan_id=str(uuid.uuid4()),
        created_at=utc_now(),
        operations=operations,
        summary=summary,
        summary_by_reason=summary_by_reason,
        human_summary=human_summary,
        metadata=metadata,
    )


def _create_operation(
    video_file: VideoFile,
    identity: Union[MovieIdentity, SeriesIdentity, None],
    config: Config
) -> FileOperation:
    """Create a file operation for a single video file.
    
    Args:
        video_file: The video file to create an operation for
        identity: Parsed identity (MovieIdentity, SeriesIdentity, or None)
        config: Configuration with templates
        
    Returns:
        FileOperation specifying what to do with the file
    """
    # Handle anime category - generate no-op (v1 constraint)
    if video_file.category == "anime":
        return FileOperation(
            operation_type="no-op",
            source_path=video_file.path,
            destination_path=None,
            reason="Anime files not organized in v1",
            has_conflict=False,
            conflict_reason=None
        )
    
    # Handle other category - generate no-op (v1 constraint)
    if video_file.category == "other":
        return FileOperation(
            operation_type="no-op",
            source_path=video_file.path,
            destination_path=None,
            reason="Other files not organized in v1",
            has_conflict=False,
            conflict_reason=None
        )
    
    # Handle files without identity - generate no-op
    if identity is None:
        return FileOperation(
            operation_type="no-op",
            source_path=video_file.path,
            destination_path=None,
            reason="No identity parsed",
            has_conflict=False,
            conflict_reason=None
        )
    
    # Handle movie identity
    if isinstance(identity, MovieIdentity):
        return _create_movie_operation(video_file, identity, config)
    
    # Handle series identity
    if isinstance(identity, SeriesIdentity):
        return _create_series_operation(video_file, identity, config)
    
    # Fallback - should not reach here
    return FileOperation(
        operation_type="no-op",
        source_path=video_file.path,
        destination_path=None,
        reason="Unknown identity type",
        has_conflict=False,
        conflict_reason=None
    )


def _create_movie_operation(
    video_file: VideoFile,
    identity: MovieIdentity,
    config: Config
) -> FileOperation:
    """Create operation for a movie file.
    
    Args:
        video_file: The movie file
        identity: Parsed movie identity
        config: Configuration with templates
        
    Returns:
        FileOperation for organizing the movie
    """
    # Explicitly blocked by manual review workflow.
    if identity.review_status == "rejected":
        return FileOperation(
            operation_type="no-op",
            source_path=video_file.path,
            destination_path=None,
            reason="Movie rejected during manual review",
            has_conflict=False,
            conflict_reason=None
        )

    # If movie needs review (no year or low-confidence enrichment), generate no-op
    if identity.needs_review or identity.year is None:
        return FileOperation(
            operation_type="no-op",
            source_path=video_file.path,
            destination_path=None,
            reason="Movie needs manual review (no year found)",
            has_conflict=False,
            conflict_reason=None
        )
    
    # Apply movie directory template
    target_dir = config.movie_template.format(
        title=identity.title,
        year=identity.year
    )
    
    # Get file extension
    ext = video_file.path.suffix
    
    # Apply movie filename template
    target_filename = config.movie_filename_template.format(
        title=identity.title,
        year=identity.year,
        ext=ext
    )
    
    # Construct full destination path
    destination = config.library_root / target_dir / target_filename
    
    # Check if source and destination are the same
    if video_file.path.resolve() == destination.resolve():
        return FileOperation(
            operation_type="no-op",
            source_path=video_file.path,
            destination_path=None,
            reason="File already at target location",
            has_conflict=False,
            conflict_reason=None
        )
    
    # Determine operation type (move or rename)
    if video_file.path.parent == destination.parent:
        operation_type = "rename"
    else:
        operation_type = "move"
    
    # Check for conflicts - destination file already exists
    has_conflict = destination.exists()
    conflict_reason = None
    if has_conflict:
        conflict_reason = f"Destination file already exists: {destination}"
    
    return FileOperation(
        operation_type=operation_type,
        source_path=video_file.path,
        destination_path=destination,
        reason=f"Organize movie: {identity.title} ({identity.year})",
        has_conflict=has_conflict,
        conflict_reason=conflict_reason
    )


def _create_series_operation(
    video_file: VideoFile,
    identity: SeriesIdentity,
    config: Config
) -> FileOperation:
    """Create operation for a series file.
    
    Args:
        video_file: The series file
        identity: Parsed series identity
        config: Configuration with templates
        
    Returns:
        FileOperation for organizing the series episode
    """
    # Explicitly blocked by manual review workflow.
    if identity.review_status == "rejected":
        return FileOperation(
            operation_type="no-op",
            source_path=video_file.path,
            destination_path=None,
            reason="Series rejected during manual review",
            has_conflict=False,
            conflict_reason=None
        )

    # If series needs review (no season or no episodes), generate no-op
    if identity.needs_review or identity.season is None or len(identity.episodes) == 0:
        return FileOperation(
            operation_type="no-op",
            source_path=video_file.path,
            destination_path=None,
            reason="Series needs manual review (no season/episode found)",
            has_conflict=False,
            conflict_reason=None
        )
    
    # Apply series directory template
    target_dir = config.series_template.format(
        title=identity.title,
        season=identity.season
    )
    
    # Get file extension
    ext = video_file.path.suffix
    
    # Apply series filename template
    # For multi-episode files, use the first episode number
    target_filename = config.series_filename_template.format(
        season=identity.season,
        episode=identity.episodes[0],
        ext=ext
    )
    
    # Construct full destination path
    destination = config.library_root / target_dir / target_filename
    
    # Check if source and destination are the same
    if video_file.path.resolve() == destination.resolve():
        return FileOperation(
            operation_type="no-op",
            source_path=video_file.path,
            destination_path=None,
            reason="File already at target location",
            has_conflict=False,
            conflict_reason=None
        )
    
    # Determine operation type (move or rename)
    if video_file.path.parent == destination.parent:
        operation_type = "rename"
    else:
        operation_type = "move"
    
    # Check for conflicts - destination file already exists
    has_conflict = destination.exists()
    conflict_reason = None
    if has_conflict:
        conflict_reason = f"Destination file already exists: {destination}"
    
    return FileOperation(
        operation_type=operation_type,
        source_path=video_file.path,
        destination_path=destination,
        reason=f"Organize series: {identity.title} S{identity.season:02d}E{identity.episodes[0]:02d}",
        has_conflict=has_conflict,
        conflict_reason=conflict_reason
    )


def _generate_summary(operations: list[FileOperation]) -> dict:
    """Generate summary statistics for operations.
    
    Args:
        operations: List of file operations
        
    Returns:
        Dictionary with operation counts by type
    """
    summary = {
        "total": len(operations),
        "move": 0,
        "rename": 0,
        "quarantine": 0,
        "no-op": 0
    }
    
    for operation in operations:
        op_type = operation.operation_type
        if op_type in summary:
            summary[op_type] += 1

    return summary


def _generate_summary_by_reason(operations: list[FileOperation]) -> dict[str, int]:
    """Aggregate operation counts by reason string."""
    out: dict[str, int] = {}
    for op in operations:
        r = op.reason
        out[r] = out.get(r, 0) + 1
    return out


def _generate_human_summary(
    operations: list[FileOperation],
    summary: dict,
    summary_by_reason: dict[str, int],
    metadata: dict,
) -> str:
    """Build a short Chinese narrative summary of the plan."""
    total = summary.get("total", len(operations))
    move = summary.get("move", 0)
    rename = summary.get("rename", 0)
    quarantine = summary.get("quarantine", 0)
    noop = summary.get("no-op", 0)
    conflicts = sum(1 for op in operations if op.has_conflict)

    parts = [
        f"本计划共 {total} 条操作：move {move}，rename {rename}，quarantine {quarantine}，no-op {noop}。"
    ]
    if summary_by_reason:
        reason_lines = [f"  - {r}: {c} 条" for r, c in sorted(summary_by_reason.items(), key=lambda x: -x[1])[:10]]
        parts.append("原因分布：")
        parts.extend(reason_lines)
    if conflicts > 0:
        parts.append(f"冲突 {conflicts} 条。")
    if metadata:
        dup = metadata.get("duplicate_groups_considered", 0)
        gaps = metadata.get("completeness_seasons_with_gaps", 0)
        if dup or gaps:
            parts.append(f"依据 analysis：重复组 {dup} 个；剧集缺口 {gaps} 季。")
    return "\n".join(parts)


def save_plan(plan: ExecutionPlan, output_path: Path) -> None:
    """Save execution plan to JSON file.
    
    Serializes the execution plan to a human-readable and editable JSON format.
    Includes plan_id, created_at timestamp, operations list, and summary.
    
    Args:
        plan: ExecutionPlan to save
        output_path: Path where the JSON file should be saved
    """
    # Convert ExecutionPlan to dictionary
    plan_dict = {
        "plan_id": plan.plan_id,
        "created_at": plan.created_at.isoformat(),
        "operations": [
            {
                "operation_type": op.operation_type,
                "source_path": str(op.source_path),
                "destination_path": str(op.destination_path) if op.destination_path else None,
                "reason": op.reason,
                "has_conflict": op.has_conflict,
                "conflict_reason": op.conflict_reason
            }
            for op in plan.operations
        ],
        "summary": plan.summary,
        "summary_by_reason": plan.summary_by_reason,
        "human_summary": plan.human_summary,
        "metadata": plan.metadata,
    }
    
    # Write to JSON file with indentation for human readability
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(plan_dict, f, indent=2, ensure_ascii=False)


def load_plan(input_path: Path) -> ExecutionPlan:
    """Load execution plan from JSON file.
    
    Deserializes an execution plan from JSON format, reconstructing all
    data structures including Path and datetime objects.
    
    Args:
        input_path: Path to the JSON file to load
        
    Returns:
        ExecutionPlan reconstructed from JSON
        
    Raises:
        FileNotFoundError: If the input file does not exist
        json.JSONDecodeError: If the file contains invalid JSON
        KeyError: If required fields are missing from the JSON
    """
    with open(input_path, 'r', encoding='utf-8') as f:
        plan_dict = json.load(f)
    
    # Reconstruct FileOperation objects
    operations = [
        FileOperation(
            operation_type=op["operation_type"],
            source_path=Path(op["source_path"]),
            destination_path=Path(op["destination_path"]) if op["destination_path"] else None,
            reason=op["reason"],
            has_conflict=op["has_conflict"],
            conflict_reason=op.get("conflict_reason")
        )
        for op in plan_dict["operations"]
    ]
    
    # Reconstruct ExecutionPlan (normalize naive datetime to UTC for backward compatibility)
    created_at = ensure_utc(datetime.fromisoformat(plan_dict["created_at"]))
    return ExecutionPlan(
        plan_id=plan_dict["plan_id"],
        created_at=created_at,
        operations=operations,
        summary=plan_dict["summary"],
        summary_by_reason=plan_dict.get("summary_by_reason", {}),
        human_summary=plan_dict.get("human_summary", ""),
        metadata=plan_dict.get("metadata", {}),
    )