#!/usr/bin/env python3
import argparse
import json
import os
import subprocess
import sys

def read_metadata_json(json_file):
    """
    Read metadata from a JSON file.

    Args:
        json_file (str): Path to the JSON file

    Returns:
        dict: Metadata information
    """
    try:
        with open(json_file, 'r', encoding='utf-8') as f:
            metadata = json.load(f)
        return metadata
    except Exception as e:
        print(f"❌ Error reading JSON file: {str(e)}")
        return None

def write_metadata_to_video(metadata, input_file, output_file=None):
    """
    Write metadata to a video file using mkvmerge.

    Args:
        metadata (dict): Metadata information
        input_file (str): Path to the input video file
        output_file (str, optional): Path to the output video file

    Returns:
        bool: True if successful, False otherwise
    """
    if not os.path.isfile(input_file):
        print(f"❌ Input file not found: {input_file}")
        return False

    if not output_file:
        # Create a temporary output file
        base_name, ext = os.path.splitext(input_file)
        output_file = f"{base_name}_modified{ext}"

    # Start building the mkvmerge command
    cmd = ["mkvmerge", "-o", output_file]

    # Add global metadata (title)
    if "title" in metadata:
        cmd.extend(["--title", metadata["title"]])

    # Process audio tracks
    for i, track in enumerate(metadata.get("audio_tracks", [])):
        track_id = i + 1  # mkvmerge track IDs start from 0 for video, 1 for first audio

        # Set language
        if "language" in track:
            cmd.extend([f"--language", f"{track_id}:{track['language']}"])

        # Set title/name
        if "name" in track and track["name"]:
            cmd.extend([f"--track-name", f"{track_id}:{track['name']}"])

        # Set disposition flags
        flags = track.get("flags", {})

        if flags.get("default", False):
            cmd.extend([f"--default-track", f"{track_id}:yes"])
        else:
            cmd.extend([f"--default-track", f"{track_id}:no"])

        if flags.get("forced", False):
            cmd.extend([f"--forced-track", f"{track_id}:yes"])
        else:
            cmd.extend([f"--forced-track", f"{track_id}:no"])

        if flags.get("original", False):
            cmd.extend([f"--original-flag", f"{track_id}:yes"])
        else:
            cmd.extend([f"--original-flag", f"{track_id}:no"])

    # Process subtitle tracks
    for i, track in enumerate(metadata.get("subtitle_tracks", [])):
        # Calculate track ID (assuming all audio tracks come before subtitle tracks)
        track_id = len(metadata.get("audio_tracks", [])) + i + 1

        # Set language
        if "language" in track:
            cmd.extend([f"--language", f"{track_id}:{track['language']}"])

        # Set title/name
        if "name" in track and track["name"]:
            cmd.extend([f"--track-name", f"{track_id}:{track['name']}"])

        # Set disposition flags
        flags = track.get("flags", {})

        if flags.get("default", False):
            cmd.extend([f"--default-track", f"{track_id}:yes"])
        else:
            cmd.extend([f"--default-track", f"{track_id}:no"])

        if flags.get("forced", False):
            cmd.extend([f"--forced-track", f"{track_id}:yes"])
        else:
            cmd.extend([f"--forced-track", f"{track_id}:no"])

    # Add input file
    cmd.append(input_file)

    # Execute the mkvmerge command
    print(f"🔄 Writing metadata to {os.path.basename(output_file)}")
    print(f"Command: {' '.join(cmd)}")

    try:
        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        if result.returncode != 0:
            print(f"❌ Error writing metadata: {result.stderr}")
            return False

        print(f"✅ Metadata written to {output_file}")
        return True
    except Exception as e:
        print(f"❌ Error executing mkvmerge: {str(e)}")
        return False

def main():
    parser = argparse.ArgumentParser(description="Write metadata from JSON to video file.")
    parser.add_argument("json_file", help="Path to input JSON metadata file")
    parser.add_argument("video_file", help="Path to input video file")
    parser.add_argument("-o", "--output", help="Path to output video file")
    args = parser.parse_args()

    json_file = args.json_file
    video_file = args.video_file
    output_file = args.output

    if not os.path.isfile(json_file):
        print(f"❌ JSON file not found: {json_file}")
        sys.exit(1)

    if not os.path.isfile(video_file):
        print(f"❌ Video file not found: {video_file}")
        sys.exit(1)

    # Read metadata from JSON
    metadata = read_metadata_json(json_file)
    if not metadata:
        sys.exit(1)

    # Write metadata to video
    success = write_metadata_to_video(metadata, video_file, output_file)
    if not success:
        sys.exit(1)

if __name__ == "__main__":
    main()