import os
import ffmpeg
import subprocess
import re
import json


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


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

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

        # Use ffmpeg to get information about the video file
        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', [])

        # Process each audio stream
        audio_streams = [s for s in streams if s.get('codec_type') == 'audio']
        if not audio_streams:
            print(f"No audio tracks found in {video_path}.")
            return

        for i, stream in enumerate(audio_streams):
            # Extract title or fallback to index
            title = stream.get('tags', {}).get('title', f"track_{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}.mp3"
            output_path = os.path.join(directory, output_name)

            # Extract the specific audio stream
            print(f"Extracting audio: {sanitized_title} ({sanitized_language}) to {output_path}")
            result = subprocess.run([
                "ffmpeg", "-i", video_path,
                "-map", f"0:a:{i}",  # Select the specific audio stream
                "-c:a", "mp3",  # Set the codec to mp3
                output_path,
                "-y"  # Overwrite if exists
            ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

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

        print(f"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 audio tracks from video files.

    Parameters:
    input_path (str): Path to the video file or directory.
    """
    if os.path.isfile(input_path):
        # Process a single file
        extract_audio_tracks(input_path)
    elif os.path.isdir(input_path):
        # Process all video files in the directory
        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_audio_tracks(file_path)
    else:
        print(f"Invalid path: {input_path}")

if __name__ == "__main__":
    import argparse

    # Set up argument parser
    parser = argparse.ArgumentParser(description="Extract all audio tracks from video files.")
    parser.add_argument("input_path", help="Path to the video file or directory.")
    
    # Parse arguments
    args = parser.parse_args()

    # Process the input path
    process_path(args.input_path)