import os
import re
import subprocess
import sys

def get_episode_code(filename):
    """Extract episode code like S01E01 from filename."""
    match = re.search(r"S\d+E\d+", filename)
    return match.group(0) if match else None

def append_videos_in_chunks(folder_path, chunk_size=4):
    video_files = []

    # Collect all video files
    for root, dirs, files in os.walk(folder_path):
        for file in sorted(files):
            if file.lower().endswith(('.mp4', '.mkv', '.mov', '.ts')):
                full_path = os.path.join(root, file)
                video_files.append(full_path)

    # Process files in chunks of 4
    for i in range(0, len(video_files), chunk_size):
        chunk = video_files[i:i + chunk_size]
        if not chunk:
            continue

        # Use the episode code of the first file in the chunk for output name
        base_filename = os.path.basename(chunk[0])
        episode_code = get_episode_code(base_filename) or f"group_{i//chunk_size + 1}"
        output_file = os.path.join(folder_path, f"{episode_code}.mkv")

        # Create the temporary list file
        temp_list_file = os.path.join(folder_path, f"{episode_code}_list.txt")
        with open(temp_list_file, "w", encoding="utf-8") as f:
            for video in chunk:
                f.write(f"file '{video}'\n")

        # Run ffmpeg to concatenate the files
        try:
            print(f"Processing chunk starting with {episode_code}...")
            subprocess.run([
                "ffmpeg", "-f", "concat", "-safe", "0", "-i", temp_list_file,
                "-map", "0", "-c", "copy", output_file
            ], check=True)
            print(f"Created {output_file}")
        except subprocess.CalledProcessError as e:
            print(f"Error processing {episode_code}: {e}")
        # finally:
            # os.remove(temp_list_file)

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(f"Usage: {sys.argv[0]} <folder_path>")
        sys.exit(1)

    folder_path = sys.argv[1]
    append_videos_in_chunks(folder_path)