import os
import subprocess
import shutil

from pdf2image import convert_from_path

from moviepy.editor import ImageSequenceClip, AudioFileClip

def stitch_frames_to_video(frame_dir: str, audio_path: str | None, output_video_path: str, fps: int = 24):
    """
    Stitches generated image frames into a video and attaches an audio track.

    Args:
        frame_dir: Directory containing the image frames.
        audio_path: Path to the master audio file (can be None for silent video).
        output_video_path: Full path for the output video file (e.g., "output.mp4").
        fps: Frames per second for the output video.
    """
    if not os.path.isdir(frame_dir):
        print(f"Error: Frame directory '{frame_dir}' not found. Cannot stitch video.")
        return

    frame_files = sorted([
        os.path.join(frame_dir, f)
        for f in os.listdir(frame_dir)
        if f.lower().endswith(('.png', '.jpg', '.jpeg'))
    ])

    if not frame_files:
        print("No frames found to stitch. Video will not be created.")
        return

    print(f"Found {len(frame_files)} frames. Proceeding to create video...")

    video_clip = None
    audio_clip = None

    try:
        video_clip = ImageSequenceClip(frame_files, fps=fps)

        if audio_path and os.path.exists(audio_path):
            try:
                audio_clip = AudioFileClip(audio_path)
                video_clip = video_clip.set_audio(audio_clip)
                print("Audio attached successfully to the video clip.")
            except Exception as e:
                print(f"Error loading or attaching audio file '{audio_path}': {e}. Video will be silent.")
        else:
            print(f"Audio file '{audio_path}' not found or not provided. Skipping audio attachment.")

        video_clip.write_videofile(
            output_video_path,
            fps=fps,
            codec="libx264",
            audio_codec="aac",
            temp_audiofile="temp-audio.m4a",
            remove_temp=True,
            threads=4,
            logger=None # Suppress moviepy verbose output
        )
        print(f"Video successfully saved to {output_video_path}")
    except Exception as e:
        print(f"AN ERROR OCCURRED during video writing: {e}")
    finally:
        # Ensure clips are closed to release resources
        if video_clip:
            video_clip.close()
        if audio_clip:
            audio_clip.close()

    # Clean up temporary frames directory
    try:
        shutil.rmtree(frame_dir)
        print(f"Removed temporary frame directory '{frame_dir}'.")
    except OSError as e:
        print(f"Error removing temporary frame directory '{frame_dir}': {e}")

def get_audio_duration(audio_path):
    cmd = ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", audio_path]
    result = subprocess.run(cmd, capture_output=True, text=True)
    try:
        return float(result.stdout.strip())
    except ValueError:
        print(f"Warning: Could not get duration for audio file {audio_path}, using default duration")
        return None

def create_video_with_ffmpeg(image_paths, audio_mapping, output_video, output_dir, fps=24):
    temp_dir = os.path.join(output_dir, "temp")
    os.makedirs(temp_dir, exist_ok=True)
    slide_videos_dir = os.path.join(output_dir, "slide_videos")
    os.makedirs(slide_videos_dir, exist_ok=True)
    slide_videos_list_path = os.path.join(temp_dir, "slide_videos_list.txt")
    individual_slide_videos = []
    try:
        with open(slide_videos_list_path, "w") as slide_list_file:
            for i, image_path in enumerate(image_paths, start=1):
                audio_file = audio_mapping.get(i)
                is_last_slide = (i == len(image_paths))
                duration = 1 if is_last_slide else 3
                if audio_file:
                    audio_filename = os.path.basename(audio_file)
                    original_audio_path = os.path.join(output_dir, "audio_check", f"slide_{i}_{audio_filename}")
                    if os.path.exists(original_audio_path):
                        fetched_duration = get_audio_duration(original_audio_path)
                        if fetched_duration is not None and fetched_duration > 0:
                             duration = fetched_duration
                slide_video_path = os.path.join(slide_videos_dir, f"slide_{i}.mp4")
                individual_slide_videos.append(slide_video_path)
                slide_list_file.write(f"file '{os.path.abspath(slide_video_path)}'\n")
                slide_ffmpeg_cmd = [
                    "ffmpeg", "-y",
                    "-loop", "1",
                    "-i", image_path,
                ]
                if audio_file and os.path.exists(original_audio_path):
                    slide_ffmpeg_cmd.extend(["-i", original_audio_path])
                    slide_ffmpeg_cmd.extend([
                        "-map", "0:v", "-map", "1:a",
                        "-c:a", "aac", "-b:a", "192k",
                    ])
                else:
                    slide_ffmpeg_cmd.extend([
                        "-f", "lavfi", "-t", str(duration),
                        "-i", "anullsrc=channel_layout=stereo:sample_rate=44100",
                        "-map", "0:v", "-map", "1:a",
                        "-c:a", "aac", "-b:a", "192k",
                    ])
                slide_ffmpeg_cmd.extend([
                    "-c:v", "libx264", "-preset", "medium",
                    "-profile:v", "high", "-level", "4.1",
                    "-pix_fmt", "yuv420p",
                    "-vf", f"scale=1920:1080:force_original_aspect_ratio=decrease:eval=frame,pad=1920:1080:(ow-iw)/2:(oh-ih)/2:color=black,fps={fps}",
                    "-t", str(duration),
                    "-shortest",
                    slide_video_path
                ])
                print(f"Creating slide {i} video: {' '.join(slide_ffmpeg_cmd)}")
                subprocess.run(slide_ffmpeg_cmd, check=True)
        print("Combining all slide videos...")
        temp_video_path = os.path.join(temp_dir, "temp_video.mp4")
        filter_cmd = [
            'ffmpeg', '-y'
        ]
        for video in individual_slide_videos:
            filter_cmd.extend(['-i', video])
        input_segments = "".join(f"[{i}:v][{i}:a]" for i in range(len(individual_slide_videos)))
        filter_cmd.extend([
            '-filter_complex',
            f'{input_segments}concat=n={len(individual_slide_videos)}:v=1:a=1[outv][outa]',
            '-map', '[outv]',
            '-map', '[outa]',
            '-c:v', 'libx264',
            '-preset', 'medium',
            '-crf', '22',
            '-c:a', 'aac',
            '-b:a', '192k',
            '-ar', '48000',
            temp_video_path
        ])
        print(f"Running filter complex command: {' '.join(filter_cmd)}")
        subprocess.run(filter_cmd, check=True)
        shutil.copy2(temp_video_path, output_video)
        print(f"Video created successfully: {output_video}")
        return output_video
    except Exception as e:
        print(f"Error creating video: {e}")
        raise
    finally:
        try:
            shutil.rmtree(temp_dir)
        except Exception as e:
            print(f"Warning: Error cleaning up temporary files: {e}")

def combine_slides_and_visualizations(slides_dir, vis_dir, output_path):
    try:
        slide_videos = sorted([f for f in os.listdir(slides_dir) if f.startswith('slide_') and f.endswith('.mp4')],
                             key=lambda x: int(x.split('_')[1].split('.')[0]))
        vis_videos = sorted([f for f in os.listdir(vis_dir) if f.startswith('vis_') and f.endswith('.mp4')],
                           key=lambda x: int(x.split('_')[1].split('.')[0]))
        if not slide_videos:
            print("No slide videos found.")
            return False
        temp_dir = "temp_processed_videos"
        os.makedirs(temp_dir, exist_ok=True)
        all_videos = []
        if len(slide_videos) >= 1:
            all_videos.append(os.path.join(slides_dir, slide_videos[0]))
        if len(slide_videos) >= 2:
            all_videos.append(os.path.join(slides_dir, slide_videos[1]))
        vis_index = 0
        slide_index = 2
        while vis_index < len(vis_videos) or slide_index < len(slide_videos):
            if vis_index < len(vis_videos):
                all_videos.append(os.path.join(vis_dir, vis_videos[vis_index]))
                vis_index += 1
            if slide_index < len(slide_videos):
                all_videos.append(os.path.join(slides_dir, slide_videos[slide_index]))
                slide_index += 1
        processed_videos = []
        for i, video_path in enumerate(all_videos):
            duration_cmd = f'ffprobe -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 "{video_path}"'
            duration_output = subprocess.check_output(duration_cmd, shell=True).decode('utf-8').strip()
            duration = float(duration_output)
            output_temp = os.path.join(temp_dir, f"processed_{i}.mp4")
            cmd = (f'ffmpeg -i "{video_path}" -c:v libx264 -preset medium -crf 22 '
                   f'-c:a aac -b:a 192k -ar 48000 -to {duration} '
                   f'-vsync 1 -async 1 -copyts -avoid_negative_ts make_zero '
                   f'"{output_temp}"')
            print(f"Processing video {i+1}/{len(all_videos)}: {cmd}")
            subprocess.call(cmd, shell=True)
            processed_videos.append(output_temp)
        output_dir = os.path.dirname(output_path)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)
        normalized_videos = []
        for i, video_path in enumerate(processed_videos):
            normalized_path = os.path.join(temp_dir, f"normalized_{i}.mp4")
            normalize_cmd = (f'ffmpeg -y -i "{video_path}" -vf "setsar=1:1" -c:v libx264 -preset medium '
                            f'-crf 22 -c:a copy "{normalized_path}"')
            print(f"Normalizing video {i} SAR: {normalize_cmd}")
            subprocess.call(normalize_cmd, shell=True)
            normalized_videos.append(normalized_path)
        input_segments = "".join(f"[{i}:v][{i}:a]" for i in range(len(normalized_videos)))
        filter_cmd = (
            'ffmpeg -y ' + ' '.join([f'-i "{v}"' for v in normalized_videos]) + ' '
            f'-filter_complex "{input_segments}concat=n={len(normalized_videos)}:v=1:a=1[outv][outa]" '
            '-map "[outv]" -map "[outa]" -c:v libx264 -preset medium -crf 22 '
            f'-c:a aac -b:a 192k -ar 48000 "{output_path}"'
        )
        print(f"Creating video using filter complex: {filter_cmd}")
        subprocess.call(filter_cmd, shell=True)
        shutil.rmtree(temp_dir, ignore_errors=True)
        return os.path.exists(output_path)
    except Exception as e:
        print(f"Error combining videos: {e}")
        import traceback
        traceback.print_exc()
        return False

def fix_video_compatibility(input_video, output_video=None):
    if output_video is None:
        base, ext = os.path.splitext(input_video)
        output_video = f"{base}_fixed.mp4"
    cmd = [
        'ffmpeg',
        '-i', input_video,
        '-c:v', 'libx264',
        '-preset', 'medium',
        '-profile:v', 'high',
        '-level', '4.1',
        '-pix_fmt', 'yuv420p',
        '-vf', 'scale=1920:1080:force_original_aspect_ratio=decrease,pad=1920:1080:(ow-iw)/2:(oh-ih)/2',
        '-r', '24',
        '-c:a', 'aac',
        '-b:a', '192k',
        '-movflags', '+faststart',
        '-y',
        output_video
    ]
    print("Running FFmpeg command to fix video compatibility:")
    print(" ".join(cmd))
    try:
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=True
        )
        while True:
            output = process.stderr.readline()
            if output == '' and process.poll() is not None:
                break
            if output:
                print(output.strip())
        if process.returncode == 0:
            print(f"\nVideo successfully fixed and saved to: {output_video}")
            return True
        else:
            print("\nError: FFmpeg process failed")
            return False
    except Exception as e:
        print(f"\nError: {str(e)}")
        return False