#!/usr/bin/env python3

import os
import subprocess
import re
import sys
import cv2
import numpy as np
from PIL import Image, ExifTags
import tempfile
import argparse
import shutil

def natural_sort_key(s):
    """Key for natural sorting of strings with numbers, matching macOS Finder order."""
    s = os.path.basename(s)
    return [int(text) if text.isdigit() else text.lower()
            for text in re.split(r'(\d+)', s)]

def get_video_info(video_file):
    """Get the duration of a video file in seconds."""
    try:
        result = subprocess.run([
            'ffprobe', '-v', 'error',
            '-select_streams', 'v:0',
            '-show_entries', 'stream=duration',
            '-of', 'default=noprint_wrappers=1:nokey=1',
            video_file
        ], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        duration_str = result.stdout.strip()
        if duration_str:
            return float(duration_str)
        else:
            return 0.0
    except Exception as e:
        print(f"Error getting video info for {video_file}: {e}")
        return 0.0

def correct_image_orientation(image):
    """Correct image orientation using EXIF data if available."""
    try:
        exif = image._getexif()
        if exif is not None:
            orientation_key = next((key for key, value in ExifTags.TAGS.items() if value == 'Orientation'), None)
            if orientation_key is not None and orientation_key in exif:
                orientation = exif[orientation_key]
                # print(f"Image has EXIF orientation: {orientation}")
                if orientation == 3:
                    image = image.rotate(180, expand=True)
                elif orientation == 6:
                    image = image.rotate(270, expand=True)
                elif orientation == 8:
                    image = image.rotate(90, expand=True)
    except Exception as e:
        print(f"Error correcting image orientation: {str(e)}")
    return image

def pad_image_to_resolution(image, target_resolution=(1280, 720)):
    """Pad the image with black borders to achieve target resolution while preserving aspect ratio."""
    width, height = image.size
    target_width, target_height = target_resolution

    # Calculate scaling factor
    scale = min(target_width / width, target_height / height)

    # Calculate new dimensions
    new_width = int(width * scale)
    new_height = int(height * scale)

    # Resize image
    resized_image = image.resize((new_width, new_height), Image.LANCZOS)

    # Create new image with black background
    new_image = Image.new('RGB', target_resolution, (0, 0, 0))

    # Calculate position to paste resized image
    left = (target_width - new_width) // 2
    top = (target_height - new_height) // 2

    # Paste resized image onto black background
    new_image.paste(resized_image, (left, top))

    return new_image

def blend_images(image1, image2, alpha):
    """Blend two images (crossfade transition)."""
    return cv2.addWeighted(image1, 1 - alpha, image2, alpha, 0)

def create_video_from_images(images, output_path, fps=30, duration=5):
    """Create a video from images with optional transitions."""
    output_resolution = (1280, 720)
    total_frames = int(fps * duration)
    transition_duration = 1  # seconds
    transition_frames = int(fps * transition_duration)

    # Correct orientation and pad images
    pil_images = []
    for img_path in images:
        try:
            img = Image.open(img_path)
            img = correct_image_orientation(img)
            img = pad_image_to_resolution(img, output_resolution)
            pil_images.append(img)
        except Exception as e:
            print(f"Error processing image {img_path}: {str(e)}")
            return False

    if not pil_images:
        print(f"No images to create video: {output_path}")
        return False

    # Convert to OpenCV images
    opencv_images = [cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) for img in pil_images]

    # Set up video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    video_writer = cv2.VideoWriter(output_path, fourcc, fps, output_resolution)

    if len(opencv_images) == 1:
        # Single image, display for total duration
        for _ in range(total_frames):
            video_writer.write(opencv_images[0])
    else:
        # Multiple images with transitions
        num_images = len(opencv_images)
        image_duration_frames = int((total_frames - (num_images - 1) * transition_frames) / num_images)
        frame_idx = 0
        for i in range(num_images):
            # Show current image
            for _ in range(image_duration_frames):
                video_writer.write(opencv_images[i])
                frame_idx += 1

            # Transition to next image
            if i < num_images - 1:
                for t in range(transition_frames):
                    alpha = t / transition_frames
                    frame = blend_images(opencv_images[i], opencv_images[i + 1], alpha)
                    video_writer.write(frame)
                    frame_idx += 1

    video_writer.release()
    print(f"Video created from images: {output_path}")
    return True

def create_custom_video_from_images(images, durations, output_path, fps=30):
    """Create a video from images with specified durations and transitions."""
    output_resolution = (1280, 720)

    # Correct orientation and pad images
    pil_images = []
    for img_path in images:
        try:
            img = Image.open(img_path)
            img = correct_image_orientation(img)
            img = pad_image_to_resolution(img, output_resolution)
            pil_images.append(img)
        except Exception as e:
            print(f"Error processing image {img_path}: {str(e)}")
            return False

    if not pil_images:
        print(f"No images to create video: {output_path}")
        return False

    # Total frames based on specified durations
    total_frames = sum([int(fps * d) for d in durations])
    transition_duration = 1  # seconds
    transition_frames = int(fps * transition_duration)

    # Convert to OpenCV images
    opencv_images = [cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) for img in pil_images]

    # Set up video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    video_writer = cv2.VideoWriter(output_path, fourcc, fps, output_resolution)

    frame_idx = 0
    num_images = len(opencv_images)

    for i in range(num_images):
        image_duration_frames = int(fps * durations[i])

        # Show current image
        for _ in range(image_duration_frames):
            video_writer.write(opencv_images[i])
            frame_idx += 1

        # Transition to next image
        if i < num_images - 1:
            for t in range(transition_frames):
                alpha = t / transition_frames
                frame = blend_images(opencv_images[i], opencv_images[i + 1], alpha)
                video_writer.write(frame)
                frame_idx += 1

    video_writer.release()
    print(f"Custom video created from images: {output_path}")
    return True

def process_images_and_create_videos(image_files, output_videos_folder):
    """Process images and create videos, handling grouped images."""
    # Sort images naturally
    image_files.sort(key=natural_sort_key)

    # Group images
    groups = {}
    single_images = []
    for img in image_files:
        base_name = re.sub(r'_(original|restored|colored)\.\w+$', '', img)
        suffix = re.search(r'_(original|restored|colored)\.\w+$', img)
        if suffix:
            key = base_name
            if key not in groups:
                groups[key] = {}
            groups[key][suffix.group(1)] = img
        else:
            single_images.append(img)

    # Create videos from grouped images
    video_paths = []
    os.makedirs(output_videos_folder, exist_ok=True)

    # Process grouped images
    for idx, (key, group_imgs) in enumerate(sorted(groups.items()), start=1):
        images = []
        durations = []
        # Ordered as per specified durations
        if 'original' in group_imgs:
            images.append(group_imgs['original'])
            durations.append(1)
        if 'restored' in group_imgs:
            images.append(group_imgs['restored'])
            durations.append(1)
        if 'colored' in group_imgs:
            images.append(group_imgs['colored'])
            durations.append(5)
        output_path = os.path.join(output_videos_folder, f"group_video_{idx}.mp4")
        success = create_custom_video_from_images(images, durations, output_path)
        if success:
            video_paths.append(output_path)

    # Process single images
    for idx, img_path in enumerate(sorted(single_images, key=natural_sort_key), start=len(groups)+1):
        output_path = os.path.join(output_videos_folder, f"single_video_{idx}.mp4")
        success = create_video_from_images([img_path], output_path)
        if success:
            video_paths.append(output_path)

    return video_paths

def concatenate_videos_in_batches(video_files, output, batch_size=25, transition_duration=1):
    """Concatenate videos in batches using FFmpeg with VideoToolbox hardware acceleration."""
    if len(video_files) == 0:
        print("No videos to concatenate.")
        return False

    if len(video_files) == 1:
        # Only one video, rename it to the output
        shutil.copy(video_files[0], output)
        print(f"Single video renamed to final output: {output}")
        return True

    temp_videos = []
    num_batches = (len(video_files) + batch_size - 1) // batch_size

    print(f"Total videos: {len(video_files)}, Batch size: {batch_size}, Number of batches: {num_batches}")

    for batch_num in range(num_batches):
        batch_videos = video_files[batch_num * batch_size:(batch_num + 1) * batch_size]
        temp_output = os.path.join(tempfile.gettempdir(), f"batch_{batch_num}.mp4")
        success = concatenate_videos_batch(batch_videos, temp_output, transition_duration)
        if not success:
            return False
        temp_videos.append(temp_output)

    # Now concatenate all batch videos
    success = concatenate_videos_batch(temp_videos, output, transition_duration)
    # Clean up temporary batch files
    for temp_video in temp_videos:
        if os.path.exists(temp_video):
            os.remove(temp_video)

    if success:
        print(f"Final video created: {output}")
    else:
        print("Error occurred during video concatenation.")
    return success

def concatenate_videos_batch(video_files, output, transition_duration):
    """Concatenate a batch of videos using FFmpeg with VideoToolbox hardware acceleration."""
    if len(video_files) == 0:
        return False
    if len(video_files) == 1:
        # Only one video, copy it to the output
        shutil.copy(video_files[0], output)
        return True

    # Prepare input arguments
    inputs = []
    for video in video_files:
        inputs.extend(['-i', video])

    # Ensure videos are sorted correctly
    video_files.sort(key=natural_sort_key)

    # Get durations and build filter_complex
    filter_complex = ''
    total_offset = 0

    for i, video in enumerate(video_files):
        duration = get_video_info(video)
        if i == 0:
            # Use nv12 pixel format compatible with VideoToolbox
            filter_complex += f'[{i}:v]format=nv12[vid{i}]; '
        else:
            offset = max(total_offset - transition_duration, 0)
            prev = i - 1
            filter_complex += f'[vid{prev}][{i}:v]xfade=transition=fade:duration={transition_duration}:offset={offset}[vid{i}]; '
        total_offset += duration

    # Trim trailing whitespace and semi-colon
    filter_complex = filter_complex.rstrip('; ')

    # Final output stream
    final_index = len(video_files) - 1
    final_map = f'[vid{final_index}]'

    # Encoding settings
    video_codec = "h264_videotoolbox"  # Use VideoToolbox hardware acceleration
    bitrate = '5000k'                  # Set a reasonable bitrate

    # Build the full ffmpeg command
    ffmpeg_command = [
        'ffmpeg',
        '-y',  # Overwrite without asking
        *inputs,
        '-filter_complex', filter_complex,
        '-map', final_map,
        '-c:v', video_codec,
        '-b:v', bitrate,         # Set video bitrate
        '-pix_fmt', 'yuv420p',   # Ensure compatibility
        output
    ]

    # Print the command (for debugging)
    print("Running FFmpeg command for batch concatenation:")
    print(' '.join(ffmpeg_command))

    # Execute the command
    try:
        subprocess.run(ffmpeg_command, check=True)
        print(f"Batch video created: {output}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"Error occurred during batch video concatenation: {e}")
        return False

def main():
    # Parse command-line arguments
    parser = argparse.ArgumentParser(description="Process images and concatenate videos with transitions using VideoToolbox (macOS).")
    parser.add_argument('--batch_size', type=int, default=25, help='Number of videos to concatenate per batch (default: 25)')
    parser.add_argument('--transition_duration', type=int, default=1, help='Duration of transitions in seconds (default: 1)')
    args = parser.parse_args()

    batch_size = args.batch_size
    transition_duration = args.transition_duration

    # Output file name
    output = "final_output.mp4"

    # Ensure the output_videos folder is created
    output_videos_folder = 'output_videos'
    os.makedirs(output_videos_folder, exist_ok=True)

    # Process images and create videos
    image_extensions = ('.jpg', '.jpeg', '.png')
    image_files = [f for f in os.listdir('.') if f.lower().endswith(image_extensions) and os.path.isfile(f)]
    if image_files:
        image_files.sort(key=natural_sort_key)
        video_files = process_images_and_create_videos(image_files, output_videos_folder)
    else:
        video_files = []

    # Get all .mp4 files in the output_videos folder and current directory
    if os.path.exists(output_videos_folder):
        video_files += [os.path.join(output_videos_folder, f) for f in os.listdir(output_videos_folder) if f.lower().endswith('.mp4')]
    video_files += [f for f in os.listdir('.') if f.lower().endswith('.mp4') and os.path.isfile(f) and f != output]

    # Remove duplicates while preserving order
    seen = set()
    video_files = [x for x in video_files if x not in seen and not seen.add(x)]

    # Sort videos naturally
    video_files.sort(key=natural_sort_key)

    # Check if there are at least two videos
    if len(video_files) < 2:
        print("Need at least two .mp4 files to concatenate with transitions.")
        sys.exit(1)

    # Concatenate videos in batches using VideoToolbox hardware acceleration
    success = concatenate_videos_in_batches(video_files, output, batch_size=batch_size, transition_duration=transition_duration)
    if success:
        print(f"Video concatenation completed successfully. Output file: {output}")
    else:
        print("Error occurred during video concatenation.")
        sys.exit(1)

if __name__ == '__main__':
    main() keep all the fucntionality of the code and all the functions, check at the end, modify the code because the batching is broken, now only short video of about five pictures rendered is saved in final video, modify the code in a way that each batch video of batch number is saved in temp_batch_video the same way as videos are saved in output_videos, both temp_batch_video and output_videos shall be cleared or created at the star of execution in case there are some old videos from previous execution, this means that for example for batch 25 all batch videos shall be stored in temp_batch_video and at the end they shall be taken in sorted order gstat --printf="%N\n" . | sort -V | paste -sd "," - and concatenated together using fade transition in the same way as other videos, double check the batch merging stores the temp batch videos and that it takes all of them as input for final video, double check that the code has all the functionality