import os
import subprocess
import re
import json


def sanitize_filename(name):
    """
    Removes invalid characters from the filename.
    """
    return re.sub(r'[<>:"/\\|?*@]', '_', name)


def extract_subtitle_tracks(video_path):
    """
    Extracts all subtitle tracks from a video file and saves them with their respective titles.

    Parameters:
    video_path (str): Path to the video file.
    """
    try:
        base_name = os.path.splitext(os.path.basename(video_path))[0]
        directory = os.path.dirname(video_path)

        # Get stream info using ffprobe
        probe_cmd = ["ffprobe", "-v", "error", "-show_streams", "-of", "json", video_path]
        probe_result = subprocess.run(probe_cmd, capture_output=True, text=True, check=True)
        streams = json.loads(probe_result.stdout).get('streams', [])

        # Filter subtitle streams
        subtitle_streams = [s for s in streams if s.get('codec_type') == 'subtitle']
        if not subtitle_streams:
            print(f"No subtitle tracks found in {video_path}.")
            return

        for i, stream in enumerate(subtitle_streams):
            title = stream.get('tags', {}).get('title', f"subtitle_{i + 1}")
            language = stream.get('tags', {}).get('language', 'unknown')
            sanitized_title = sanitize_filename(title)
            sanitized_language = sanitize_filename(language)
            output_name = f"{base_name}__{sanitized_title}__{sanitized_language}.srt"
            output_path = os.path.join(directory, output_name)

            # Extract subtitle stream
            print(f"Extracting subtitle: {sanitized_title} ({sanitized_language}) to {output_path}")
            result = subprocess.run([
                "ffmpeg", "-i", video_path,
                "-map", f"0:s:{i}",
                "-c:s", "srt",
                output_path,
                "-y"
            ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

            if result.returncode != 0:
                print(f"ffmpeg error for subtitle stream {i}: {result.stderr}")

        print(f"Subtitle extraction completed for {video_path}.")
    except Exception as e:
        print(f"Error processing {video_path}: {e}")


def process_path(input_path):
    """
    Processes a file or directory, extracting all subtitle tracks from video files.

    Parameters:
    input_path (str): Path to the video file or directory.
    """
    if os.path.isfile(input_path):
        extract_subtitle_tracks(input_path)
    elif os.path.isdir(input_path):
        for file_name in os.listdir(input_path):
            file_path = os.path.join(input_path, file_name)
            if os.path.isfile(file_path) and file_name.lower().endswith(('.mp4', '.mkv', '.avi', '.mov')):
                extract_subtitle_tracks(file_path)
    else:
        print(f"Invalid path: {input_path}")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Extract all subtitle tracks from video files.")
    parser.add_argument("input_path", help="Path to the video file or directory.")
    args = parser.parse_args()

    process_path(args.input_path)