EpisodeMatcher/episode_matcher.py

#!/usr/bin/env python3
"""
Episode Matcher - Main script for classifying and renaming TV show episode files.

This tool takes a folder of MKV files and:
1. Classifies them as episodes or extras based on file size and duration
2. Moves extras to an "extras" subfolder  
3. Renames episodes using Jellyfin naming convention
4. Uses TVDB data to validate episode matching
5. Detects and handles duplicate episodes

Usage:
    python episode_matcher.py <folder_path> <show_name> <season_number> [options]

Examples:
    python episode_matcher.py "/path/to/episodes" "Game of Thrones" 1
    python episode_matcher.py "/path/to/episodes" "Game of Thrones" 2 --disc-mapping "1:1-3,2:4-6,3:7-9,4:10"
    python episode_matcher.py "/path/to/episodes" "Breaking Bad" 2 --auto-delete-duplicates --dry-run
"""

import argparse
import sys
import os
from pathlib import Path

try:
    from src.TVDBProvider import TVDBClient
    from src.TVDBProvider.tvdb_client import MockTVDBClient
    from src.Matcher import FileClassifier, EpisodeRenamer
    from src.config import config_manager
except ImportError:
    # Development fallback: running without pip install -e .
    import importlib
    _src = os.path.join(os.path.dirname(__file__), 'src')
    _spec_tvdb = importlib.util.spec_from_file_location(
        "TVDBProvider", os.path.join(_src, "TVDBProvider", "__init__.py"))
    _tvdb_mod = importlib.util.module_from_spec(_spec_tvdb)
    _spec_tvdb.loader.exec_module(_tvdb_mod)
    TVDBClient = _tvdb_mod.TVDBClient

    _spec_client = importlib.util.spec_from_file_location(
        "TVDBProvider.tvdb_client", os.path.join(_src, "TVDBProvider", "tvdb_client.py"))
    _client_mod = importlib.util.module_from_spec(_spec_client)
    _spec_client.loader.exec_module(_client_mod)
    MockTVDBClient = _client_mod.MockTVDBClient

    _spec_matcher = importlib.util.spec_from_file_location(
        "Matcher", os.path.join(_src, "Matcher", "__init__.py"))
    _matcher_mod = importlib.util.module_from_spec(_spec_matcher)
    _spec_matcher.loader.exec_module(_matcher_mod)
    FileClassifier = _matcher_mod.FileClassifier
    EpisodeRenamer = _matcher_mod.EpisodeRenamer

    _spec_config = importlib.util.spec_from_file_location(
        "config", os.path.join(_src, "config.py"))
    _config_mod = importlib.util.module_from_spec(_spec_config)
    _spec_config.loader.exec_module(_config_mod)
    config_manager = _config_mod.config_manager


def parse_disc_mapping(mapping_str):
    """Parse disc mapping string like '1:1-3,2:4-6,3:7-9,4:10' into a dict."""
    mapping = {}
    
    if not mapping_str:
        return mapping
    
    try:
        # Split by comma to get each disc mapping
        disc_parts = mapping_str.split(',')
        
        for part in disc_parts:
            # Split by colon to get disc:episodes
            disc_str, episodes_str = part.strip().split(':')
            disc_num = int(disc_str)
            
            # Parse episode range
            if '-' in episodes_str:
                # Range like "1-3"
                start, end = episodes_str.split('-')
                episodes = list(range(int(start), int(end) + 1))
            else:
                # Single episode like "10"
                episodes = [int(episodes_str)]
            
            mapping[disc_num] = episodes
        
        return mapping
    
    except Exception as e:
        print(f"Error parsing disc mapping '{mapping_str}': {e}")
        print("Expected format: '1:1-3,2:4-6,3:7-9,4:10'")
        sys.exit(1)


def main():
    """Main function to run the episode matcher."""
    parser = argparse.ArgumentParser(
        description="Classify and rename TV show episode files",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__
    )
    
    parser.add_argument(
        "folder_path", 
        help="Path to folder containing MKV files"
    )
    
    parser.add_argument(
        "show_name",
        help="Name of the TV show (e.g., 'Game of Thrones')"
    )
    
    parser.add_argument(
        "season_number",
        type=int,
        help="Season number (e.g., 1, 2, 3)"
    )
    
    parser.add_argument(
        "--api-key",
        help="TVDB API key (optional, will check config.json first, then use mock data if not found)"
    )
    
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show what would be done without actually renaming/moving files"
    )
    
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Show detailed analysis information"
    )
    
    parser.add_argument(
        "--disc-mapping",
        help="Disc to episode mapping (e.g., '1:1-3,2:4-6,3:7-9,4:10')"
    )
    
    parser.add_argument(
        "--auto-delete-duplicates",
        action="store_true",
        help="Automatically delete duplicate files instead of moving them to 'delete me' folder"
    )
    
    parser.add_argument(
        "--force",
        action="store_true",
        help="Confirm destructive operations (required with --auto-delete-duplicates to actually delete files)"
    )

    parser.add_argument(
        "--extensions",
        help="Comma-separated video extensions to process (e.g., 'mkv,mp4,avi'). Defaults to all common formats."
    )

    parser.add_argument(
        "--fallback-ratio",
        type=float,
        help="Minutes of video per GB when MediaInfo unavailable (default: 45). Use ~25 for 4K, ~60 for low-bitrate."
    )
    
    args = parser.parse_args()
    
    # Validate inputs
    folder_path = Path(args.folder_path)
    if not folder_path.exists():
        print(f"Error: Folder not found: {folder_path}")
        sys.exit(1)
    
    if not folder_path.is_dir():
        print(f"Error: Path is not a directory: {folder_path}")
        sys.exit(1)
    
    if args.season_number < 1:
        print(f"Error: Season number must be positive, got: {args.season_number}")
        sys.exit(1)
    
    print(f"=== Episode Matcher ===")
    print(f"Folder: {folder_path}")
    print(f"Show: {args.show_name}")
    print(f"Season: {args.season_number}")
    print(f"Mode: {'DRY RUN' if args.dry_run else 'LIVE'}")
    
    # Show config status
    if args.verbose:
        print("\n=== Configuration ===")
        config_manager.print_config_status()
    
    print()
    
    # Get API key from config or command line argument
    api_key = args.api_key or config_manager.get_tvdb_api_key()
    
    # Initialize TVDB client
    if api_key:
        print("Using TVDB API with configured key...")
        tvdb_client = TVDBClient()
        if not tvdb_client.authenticate(api_key):
            print("TVDB authentication failed, falling back to mock data")
            tvdb_client = MockTVDBClient()
            tvdb_client.authenticate()
    else:
        print("No API key found in config.json or command line, using mock TVDB data...")
        print("To use real TVDB data, add your API key to config.json")
        tvdb_client = MockTVDBClient()
        tvdb_client.authenticate()
    
    # Get episode information from TVDB
    print(f"Fetching episode data for '{args.show_name}' season {args.season_number}...")
    tvdb_episodes = tvdb_client.get_episode_durations(args.show_name, args.season_number)
    
    if tvdb_episodes:
        print(f"Found {len(tvdb_episodes)} episodes in TVDB data")
        if args.verbose:
            print("\nTVDB Episodes:")
            for ep in tvdb_episodes:
                print(f"  Episode {ep['episode_number']}: {ep['name']} ({ep['runtime']} min)")
    else:
        print("Warning: Could not fetch episode data from TVDB")
        tvdb_episodes = []
    
    # Initialize file classifier
    video_extensions = args.extensions.split(',') if args.extensions else None
    try:
        classifier = FileClassifier(
            str(folder_path),
            video_extensions=video_extensions,
            fallback_duration_minutes_per_gb=args.fallback_ratio,
        )
    except FileNotFoundError as e:
        print(f"Error: {e}")
        sys.exit(1)
    
    if args.verbose:
        classifier.print_analysis()
    
    # Classify files
    expected_episode_count = len(tvdb_episodes) if tvdb_episodes else None
    episodes, extras = classifier.classify_files(expected_episode_count, tvdb_episodes)
    
    print(f"\n=== Classification Results ===")
    print(f"Episodes: {len(episodes)}")
    print(f"Extras: {len(extras)}")
    
    if not episodes:
        print("Error: No episodes found. Check your files and try again.")
        sys.exit(1)
    
    if args.verbose:
        print("\nEpisodes:")
        for ep in episodes:
            print(f"  - {ep['name']} ({ep['size_gb']:.2f} GB, {ep['duration_minutes']:.1f} min)")
        
        if extras:
            print("\nExtras:")
            for ex in extras:
                print(f"  - {ex['name']} ({ex['size_gb']:.2f} GB, {ex['duration_minutes']:.1f} min)")
    
    # Initialize renamer
    disc_mapping = None
    if args.disc_mapping:
        print(f"Using disc mapping: {args.disc_mapping}")
        disc_mapping = parse_disc_mapping(args.disc_mapping)
        if args.verbose:
            print("Parsed disc mapping:")
            for disc, episodes in disc_mapping.items():
                print(f"  Disc {disc}: Episodes {episodes}")
    
    if args.auto_delete_duplicates and not args.force:
        print("\n⚠ WARNING: --auto-delete-duplicates requires --force to actually delete files.")
        print("   Without --force, duplicates will be moved to 'delete me' folder.\n")
    
    renamer = EpisodeRenamer(str(folder_path), args.show_name, args.season_number, 
                            disc_mapping=disc_mapping, auto_delete_duplicates=args.auto_delete_duplicates and args.force)
    
    if args.dry_run:
        print(f"\n=== DRY RUN - No files will be modified ===")
        
        if extras:
            print(f"\nWould move {len(extras)} files to extras folder:")
            for extra in extras:
                print(f"  - {extra['name']}")
        
        print(f"\nWould rename {len(episodes)} episode files:")
        if args.verbose:
            print("Episodes to be renamed:")
            for i, ep in enumerate(episodes):
                print(f"  {i+1}: {ep} (type: {type(ep)})")
        
        # Handle the case where episodes is incorrectly formatted
        if episodes and not isinstance(episodes[0], dict):
            print("Warning: Episodes data is in wrong format, attempting to reconstruct...")
            # This is a temporary workaround - we need to get the actual episode data
            # from the classifier again
            all_files = classifier.file_info
            episode_files = [f for f in all_files if f not in extras]
            episodes = episode_files[:len(tvdb_episodes)] if tvdb_episodes else episode_files
            
        # Simulate the matching process
        if tvdb_episodes:
            matched_episodes = renamer._match_episodes_by_duration(episodes, tvdb_episodes)
        else:
            sorted_episodes = sorted(episodes, key=lambda x: x['size_bytes'], reverse=True)
            matched_episodes = []
            for i, episode in enumerate(sorted_episodes, 1):
                matched_episodes.append({
                    'file_info': episode,
                    'episode_number': i,
                    'tvdb_info': None
                })
        
        for match in matched_episodes:
            episode_number = match['episode_number']
            file_info = match['file_info']
            original_name = file_info['name']  # This should work now
            new_name = renamer._generate_episode_filename(episode_number, Path(original_name).suffix)
            print(f"  Episode {episode_number:2d}: {original_name} → {new_name}")
    
    else:
        # Move extras to folder
        if extras:
            print(f"\n=== Moving {len(extras)} extras to subfolder ===")
            moved_extras = renamer.move_extras_to_folder(extras)
            print(f"Successfully moved {len(moved_extras)} files to extras folder")
        
        # Rename episodes
        print(f"\n=== Renaming {len(episodes)} episodes ===")
        renamed_files = renamer.rename_episodes(episodes, tvdb_episodes)
        
        if renamed_files:
            renamer.print_rename_summary(renamed_files)
            
            # Validate that all episodes were processed
            expected_count = len(tvdb_episodes) if tvdb_episodes else len(episodes)
            success = renamer.validate_season_complete(renamed_files, expected_count)
            
            if success:
                print(f"\n✓ Episode matching completed successfully!")
            else:
                print(f"\n⚠ Episode matching completed with warnings")
                sys.exit(1)
        else:
            print("Error: No episodes were renamed")
            sys.exit(1)


if __name__ == "__main__":
    main()