import os
import subprocess
import tempfile
import json
import time
from moviepy import AudioFileClip, ImageClip, concatenate_videoclips

class ProgressCallback:
    """Video generation ko bela progress track garne callback class"""
    def __init__(self, progress_file=None, total_steps=100):
        self.progress_file = progress_file
        self.total_steps = total_steps
        self.current_step = 0
        self.last_update_time = 0
        self.update_interval = 0.5
        
    def update(self, step=1, message="", force=False):
        """Progress lai specified step le update garne"""
        self.current_step += step
        progress = min(100, int((self.current_step / self.total_steps) * 100))
        
        print(f"Progress: {progress}% - {message}")
        
        current_time = time.time()
        if self.progress_file and (force or (current_time - self.last_update_time) > self.update_interval):
            try:
                with open(self.progress_file, 'w') as f:
                    json.dump({
                        "progress": progress,
                        "message": message,
                        "timestamp": current_time
                    }, f)
                self.last_update_time = current_time
            except Exception as e:
                print(f"Error writing progress file: {e}")
    
    def complete(self):
        """Progress lai complete mark garne"""
        if self.progress_file:
            try:
                with open(self.progress_file, 'w') as f:
                    json.dump({
                        "progress": 100,
                        "message": "Complete",
                        "timestamp": time.time()
                    }, f)
            except Exception as e:
                print(f"Error writing progress file: {e}")
        print("Progress: 100% - Complete")

def get_audio_duration(audio_path):
    """FFprobe use garera audio duration nikalne"""
    try:
        cmd = [
            'ffprobe', '-v', 'quiet', '-print_format', 'json', 
            '-show_format', audio_path
        ]
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode == 0:
            data = json.loads(result.stdout)
            return float(data['format']['duration'])
    except Exception as e:
        print(f"Error getting audio duration: {e}")
    return None

def create_individual_segments_ffmpeg_synchronized(image_dir, audio_dir, temp_dir, gap_duration=0.5, progress_callback=None, job_id=None):
    """
    FFmpeg use garera synchronized individual segments create garne
    Audio duration anusar image duration set garne
    
    Args:
        image_dir: Image files directory
        audio_dir: Audio files directory  
        temp_dir: Temporary files store garne directory
        gap_duration: Segments bich ko gap
        progress_callback: Progress tracking
        job_id: Job specific files filter garne
    
    Returns:
        List of created segment video files with timing info
    """
    segment_files = []
    timing_info = []
    
    # Image files ko sorted list
    image_files = sorted([f for f in os.listdir(image_dir) 
                         if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp'))])
    
    if job_id:
        image_files = [f for f in image_files if job_id in f]
    
    print(f"Found {len(image_files)} image files")
    
    if progress_callback is None:
        progress_callback = ProgressCallback(total_steps=len(image_files) * 2)  # segments + gaps
    
    # Ensure temp directory exists
    os.makedirs(temp_dir, exist_ok=True)
    
    current_start_time = 0.0
    
    for i, image_file in enumerate(image_files):
        progress_callback.update(0, f"Processing segment {i+1}/{len(image_files)}")
        
        # Find matching audio file
        if job_id and "segment_" in image_file:
            segment_pattern = image_file.split("_")[0] + "_" + image_file.split("_")[1]
            audio_candidates = [f for f in os.listdir(audio_dir) 
                              if segment_pattern in f and job_id in f]
        else:
            base_name = os.path.splitext(image_file)[0]
            audio_candidates = [f for f in os.listdir(audio_dir) 
                              if f.startswith(base_name)]
        
        if not audio_candidates:
            print(f"No matching audio found for {image_file}")
            continue
            
        image_path = os.path.join(image_dir, image_file)
        audio_path = os.path.join(audio_dir, audio_candidates[0])
        
        # Get audio duration for exact synchronization
        audio_duration = get_audio_duration(audio_path)
        if audio_duration is None:
            print(f"Could not get duration for {audio_path}")
            continue
        
        print(f"Audio duration for {audio_candidates[0]}: {audio_duration:.3f} seconds")
        
        # Create segment using FFmpeg with exact duration synchronization
        segment_output = os.path.join(temp_dir, f"segment_{i:03d}.mp4")
        
        try:
            # Use FFmpeg to create video segment with exact audio duration
            ffmpeg_cmd = [
                'ffmpeg', '-y',
                '-loop', '1',
                '-i', image_path,          # Image input (looped)
                '-i', audio_path,          # Audio input
                '-c:v', 'libx264',         # Video codec
                '-c:a', 'aac',             # Audio codec
                '-t', str(audio_duration), # Exact duration from audio
                '-pix_fmt', 'yuv420p',     # Pixel format for compatibility
                '-r', '24',                # Frame rate
                '-shortest',               # Stop when shortest input ends
                segment_output
            ]
            
            result = subprocess.run(ffmpeg_cmd, capture_output=True, text=True)
            
            if result.returncode == 0:
                segment_files.append(segment_output)
                
                # Record timing info
                segment_info = {
                    'file': segment_output,
                    'start_time': current_start_time,
                    'duration': audio_duration,
                    'end_time': current_start_time + audio_duration,
                    'type': 'segment'
                }
                timing_info.append(segment_info)
                current_start_time += audio_duration
                
                print(f"Created segment: {segment_output} ({audio_duration:.3f}s)")
                progress_callback.update(1, f"Segment {i+1} created")
                
                # Add gap segment if not the last one
                if i < len(image_files) - 1 and gap_duration > 0:
                    gap_output = os.path.join(temp_dir, f"gap_{i:03d}.mp4")
                    
                    # Create gap video with exact duration
                    gap_cmd = [
                        'ffmpeg', '-y',
                        '-loop', '1',
                        '-i', image_path,          # Use same image for gap
                        '-t', str(gap_duration),   # Exact gap duration
                        '-c:v', 'libx264',
                        '-pix_fmt', 'yuv420p',
                        '-r', '24',
                        '-an',                     # No audio for gap
                        gap_output
                    ]
                    
                    gap_result = subprocess.run(gap_cmd, capture_output=True, text=True)
                    if gap_result.returncode == 0:
                        segment_files.append(gap_output)
                        
                        # Record gap timing info
                        gap_info = {
                            'file': gap_output,
                            'start_time': current_start_time,
                            'duration': gap_duration,
                            'end_time': current_start_time + gap_duration,
                            'type': 'gap'
                        }
                        timing_info.append(gap_info)
                        current_start_time += gap_duration
                        
                        print(f"Created gap segment: {gap_output} ({gap_duration}s)")
                        progress_callback.update(1, f"Gap {i+1} created")
                    else:
                        print(f"Failed to create gap: {gap_result.stderr}")
            else:
                print(f"FFmpeg error for {image_file}: {result.stderr}")
                
        except Exception as e:
            print(f"Error creating segment {i}: {e}")
    
    # Print timing summary
    print("\n=== TIMING SUMMARY ===")
    total_duration = 0
    for info in timing_info:
        print(f"{info['type'].upper()}: {info['start_time']:.3f}s - {info['end_time']:.3f}s ({info['duration']:.3f}s)")
        total_duration = max(total_duration, info['end_time'])
    print(f"Total video duration: {total_duration:.3f}s")
    print("====================\n")
    
    return segment_files, timing_info

def concatenate_segments_ffmpeg_verified(segment_files, output_path, temp_dir, timing_info=None, progress_callback=None):
    """
    FFmpeg use garera segments lai concatenate garne with verification
    
    Args:
        segment_files: Video segment files ko list
        output_path: Final output video path
        temp_dir: Temporary directory
        timing_info: Timing information for verification
        progress_callback: Progress tracking
    """
    if not segment_files:
        print("No segments to concatenate")
        return False
    
    if progress_callback is None:
        progress_callback = ProgressCallback(total_steps=3)
    
    # Create concat file for FFmpeg
    concat_file = os.path.join(temp_dir, 'concat_list.txt')
    
    try:
        # Verify each segment before concatenation
        progress_callback.update(0, "Verifying segments...")
        verified_segments = []
        
        for i, segment_file in enumerate(segment_files):
            if os.path.exists(segment_file) and os.path.getsize(segment_file) > 0:
                # Get actual duration of segment
                duration = get_audio_duration(segment_file)
                if duration is not None:
                    print(f"Verified segment {i}: {duration:.3f}s - {os.path.basename(segment_file)}")
                    verified_segments.append(segment_file)
                else:
                    print(f"Warning: Could not verify duration for {segment_file}")
                    verified_segments.append(segment_file)  # Include anyway
            else:
                print(f"Warning: Segment file missing or empty: {segment_file}")
        
        if not verified_segments:
            print("No valid segments found for concatenation")
            return False
        
        progress_callback.update(1, f"Verified {len(verified_segments)} segments")
        
        # Write concat file
        with open(concat_file, 'w') as f:
            for segment_file in verified_segments:
                # Escape path for FFmpeg
                escaped_path = segment_file.replace("'", "'\"'\"'")
                f.write(f"file '{escaped_path}'\n")
        
        progress_callback.update(0, "Creating final video...")
        
        # Use FFmpeg concat demuxer - most reliable for sequential playback
        concat_cmd = [
            'ffmpeg', '-y',
            '-f', 'concat',
            '-safe', '0',
            '-i', concat_file,
            '-c', 'copy',      # Copy streams without re-encoding for speed and quality
            '-avoid_negative_ts', 'make_zero',  # Ensure proper timestamp alignment
            output_path
        ]
        
        print("Running FFmpeg concatenation...")
        result = subprocess.run(concat_cmd, capture_output=True, text=True)
        
        if result.returncode == 0:
            progress_callback.update(1, "Video concatenation complete")
            
            # Verify final output
            final_duration = get_audio_duration(output_path)
            if final_duration:
                print(f"Final video duration: {final_duration:.3f}s")
                
                # Compare with expected duration
                if timing_info:
                    expected_duration = max(info['end_time'] for info in timing_info)
                    duration_diff = abs(final_duration - expected_duration)
                    if duration_diff < 0.1:  # Within 100ms tolerance
                        print("✓ Duration verification passed")
                    else:
                        print(f"⚠ Duration mismatch: expected {expected_duration:.3f}s, got {final_duration:.3f}s")
            
            progress_callback.complete()
            print(f"Final video created: {output_path}")
            return True
        else:
            print(f"FFmpeg concatenation error: {result.stderr}")
            return False
            
    except Exception as e:
        print(f"Error in concatenation: {e}")
        return False

def create_video_memory_efficient_synchronized(image_dir, audio_dir, output_path, gap_duration=0.5, 
                                             progress_callback=None, job_id=None, cleanup_temp=True):
    """
    Memory efficient video creation with perfect audio-image synchronization
    
    Args:
        image_dir: Image files directory
        audio_dir: Audio files directory
        output_path: Final output video path
        gap_duration: Gap between segments (seconds)
        progress_callback: Progress tracking callback
        job_id: Job specific file filtering
        cleanup_temp: Whether to cleanup temporary files after creation
    
    Returns:
        Success boolean and timing information
    """
    # Create temporary directory
    with tempfile.TemporaryDirectory() as temp_dir:
        print(f"Using temporary directory: {temp_dir}")
        
        if progress_callback is None:
            progress_callback = ProgressCallback(total_steps=100)
        
        # Step 1: Create synchronized individual segments
        progress_callback.update(0, "Creating synchronized segments...")
        segment_files, timing_info = create_individual_segments_ffmpeg_synchronized(
            image_dir, audio_dir, temp_dir, gap_duration, progress_callback, job_id
        )
        
        if not segment_files:
            print("No segments were created")
            return False, None
        
        # Step 2: Concatenate segments with verification
        progress_callback.update(0, "Concatenating segments...")
        success = concatenate_segments_ffmpeg_verified(
            segment_files, output_path, temp_dir, timing_info, progress_callback
        )
        
        if cleanup_temp:
            print("Temporary files will be automatically cleaned up")
        else:
            print(f"Temporary files kept in: {temp_dir}")
            
        return success, timing_info

def create_video_from_directories(image_dir, audio_dir, output_path, gap_duration=0.5, 
                                             progress_callback=None, job_id=None):
    """
    Main function with perfect synchronization
    Ensures each segment has exactly matching audio-image duration
    Maintains sequential order without overlaps
    """
    print("Creating synchronized video with exact audio-image matching...")
    print(f"Gap duration between segments: {gap_duration}s")
    
    # Validate input directories
    if not os.path.exists(image_dir):
        print(f"Error: Image directory not found: {image_dir}")
        return False
    
    if not os.path.exists(audio_dir):
        print(f"Error: Audio directory not found: {audio_dir}")
        return False
    
    # Check for FFmpeg
    try:
        subprocess.run(['ffmpeg', '-version'], capture_output=True, check=True)
    except (subprocess.CalledProcessError, FileNotFoundError):
        print("Error: FFmpeg not found. Please install FFmpeg to use this function.")
        return False
    
    success, timing_info = create_video_memory_efficient_synchronized(
        image_dir, audio_dir, output_path, gap_duration, progress_callback, job_id
    )
    
    if success and timing_info:
        print("\n✓ Video creation successful with perfect synchronization!")
        print("✓ No segment overlaps or timing issues")
        print("✓ Audio and image durations exactly matched")
        return True
    else:
        print("✗ Video creation failed")
        return False

# if __name__ == "__main__":
#     # Test with synchronized approach
#     print("Testing synchronized video creation...")
    
#     success = create_video_from_directories_synchronized(
#         './images/', 
#         './audio/', 
#         'slideshow_synchronized.mp4',
#         gap_duration=0.5  # 500ms gap between segments
#     )
    
#     if success:
#         print("✓ Synchronized video created successfully!")
#     else:
#         print("✗ Video creation failed!")