#!/usr/bin/env python3
"""
MP4 to JPEG Converter with ultra-simplified C array generation
生成极简C语言数组，只包含原始数据和偏移地址
"""

import cv2
import os
import argparse
import sys

def align_to(value, alignment):
    """将值对齐到指定的字节边界"""
    return ((value + alignment - 1) // alignment) * alignment

def save_as_c_array(images_data, output_c_file, array_name="image_data"):
    """
    将图片数据保存为最简化的C语言数组
    
    Args:
        images_data: 图片数据字节串列表
        output_c_file: 输出C文件路径
        array_name: C数组名称前缀
    """
    
    if not images_data:
        print("Warning: No image data to save")
        return False
    
    print(f"Generating simplified C array file: {output_c_file}")
    
    try:
        # 计算总数据大小和偏移地址
        total_size = 0
        offsets = []
        
        for data in images_data:
            data_size = len(data)
            aligned_size = align_to(data_size, 8)
            offsets.append(total_size)
            total_size += aligned_size
        
        print(f"  Total images: {len(images_data)}")
        print(f"  Total data size: {total_size} bytes")
        
        with open(output_c_file, 'w', encoding='utf-8') as f:
            # 文件头
            f.write("/* Auto-generated by MP4 to JPEG converter */\n")
            f.write(f"/* Images: {len(images_data)}, Total size: {total_size} bytes, 8-byte aligned */\n\n")
            
            # 图片数量宏定义
            f.write(f"#define {array_name.upper()}_COUNT {len(images_data)}\n\n")
            
            # 原始数据数组
            f.write(f"/* All image data (8-byte aligned) */\n")
            f.write(f"__attribute((aligned(32))) ATTR_PSRAM_DATA_SECTION uint8_t {array_name}_raw[{total_size}] = {{\n")
            
            # 写入所有图片数据（连续存储）
            all_data = bytearray()
            for i, data in enumerate(images_data):
                data_size = len(data)
                aligned_size = align_to(data_size, 8)
                
                # 添加图片数据
                all_data.extend(data)
                
                # 添加对齐填充
                padding_size = aligned_size - data_size
                all_data.extend(b'\x00' * padding_size)
            
            # 写入数据，每行16个字节
            for i in range(0, len(all_data), 16):
                line = "    "
                for j in range(16):
                    if i + j < len(all_data):
                        byte_val = all_data[i + j]
                        line += f"0x{byte_val:02X}, "
                f.write(line + "\n")
            
            f.write("};\n\n")
            
            # 偏移地址数组
            f.write(f"/* Array of image data offset */\n")
            f.write(f"const uint32_t {array_name}_offset[{array_name.upper()}_COUNT] = {{\n")
            
            # 每行4个偏移值
            for i in range(0, len(offsets), 4):
                line = "    "
                for j in range(4):
                    idx = i + j
                    if idx < len(offsets):
                        offset = offsets[idx]
                        f.write(f"{offset:8d},  /* Image {idx:3d} */\n")
            
            f.write("};\n")
        
        print(f"C array file generated successfully!")
        return True
        
    except Exception as e:
        print(f"Error generating C array file: {e}")
        import traceback
        traceback.print_exc()
        return False

def save_as_simple_c_array(images_data, output_c_file, array_name="image_data"):
    """
    保存为更简单的C数组格式，每行一个偏移值
    """
    
    if not images_data:
        print("Warning: No image data to save")
        return False
    
    print(f"Generating simple C array file: {output_c_file}")
    
    try:
        # 计算总数据大小和偏移地址
        total_size = 0
        offsets = []
        
        for data in images_data:
            data_size = len(data)
            aligned_size = align_to(data_size, 8)
            offsets.append(total_size)
            total_size += aligned_size
        
        print(f"  Total images: {len(images_data)}")
        print(f"  Total data size: {total_size} bytes")
        
        with open(output_c_file, 'w', encoding='utf-8') as f:
            # 文件头
            f.write("/* Auto-generated by MP4 to JPEG converter */\n")
            f.write(f"/* Images: {len(images_data)}, Total size: {total_size} bytes */\n\n")
            
            # 图片数量宏定义
            f.write(f"#define {array_name.upper()}_COUNT {len(images_data)}\n\n")
            
            # 原始数据数组
            f.write(f"__attribute((aligned(32))) ATTR_PSRAM_DATA_SECTION uint8_t {array_name}_raw[{total_size}] = {{\n")
            
            # 写入所有图片数据（连续存储）
            all_data = bytearray()
            for data in images_data:
                data_size = len(data)
                aligned_size = align_to(data_size, 8)
                
                # 添加图片数据
                all_data.extend(data)
                
                # 添加对齐填充
                padding_size = aligned_size - data_size
                all_data.extend(b'\x00' * padding_size)
            
            # 写入数据，每行16个字节
            bytes_written = 0
            while bytes_written < len(all_data):
                line = "    "
                for j in range(16):
                    if bytes_written < len(all_data):
                        byte_val = all_data[bytes_written]
                        line += f"0x{byte_val:02X}, "
                        bytes_written += 1
                f.write(line + "\n")
            
            f.write("};\n\n")
            
            # 偏移地址数组
            f.write(f"const uint32_t {array_name}_offset[{array_name.upper()}_COUNT] = {{\n")
            
            # 每行一个偏移值
            for i, offset in enumerate(offsets):
                f.write(f"    {offset},  /* Image {i} */\n")
            
            f.write("};\n")
        
        print(f"Simple C array file generated successfully!")
        return True
        
    except Exception as e:
        print(f"Error generating simple C array file: {e}")
        return False

def mp4_to_jpg(video_path, output_dir, width=None, height=None, 
               quality=95, start_time=0, duration=None, prefix='frame',
               generate_c_array=False, c_array_name="image_data", c_array_file=None):
    """
    Convert MP4 video to JPEG images with ultra-simplified C array generation
    """
    
    # Create output directory
    os.makedirs(output_dir, exist_ok=True)
    print(f"Output directory: {output_dir}")
    
    # Check input file
    if not os.path.exists(video_path):
        print(f"Error: File '{video_path}' does not exist")
        return False, []
    
    print(f"Input video: {video_path}")
    
    # Open video
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error: Cannot open video file")
        return False, []
    
    # Get video info
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    if fps <= 0 or total_frames <= 0:
        print("Error: Invalid video file")
        cap.release()
        return False, []
    
    print(f"Video FPS: {fps:.2f}")
    print(f"Total frames: {total_frames}")
    
    # Calculate frame range
    start_frame = int(start_time * fps)
    
    if duration is not None:
        end_frame = int((start_time + duration) * fps)
        end_frame = min(end_frame, total_frames)
    else:
        end_frame = total_frames
    
    print(f"Start frame: {start_frame}")
    print(f"End frame: {end_frame}")
    print(f"Frames to extract: {end_frame - start_frame}")
    
    # Set start position
    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
    
    # Get original size
    original_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    original_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    print(f"Original size: {original_width} x {original_height}")
    
    # Determine output size
    if width is None and height is None:
        output_width, output_height = original_width, original_height
        print("Keeping original size")
    elif width is not None and height is not None:
        output_width, output_height = width, height
        print(f"Target size: {output_width} x {output_height}")
    else:
        # Calculate maintaining aspect ratio
        if width is None:
            aspect_ratio = original_width / original_height
            output_height = height
            output_width = int(output_height * aspect_ratio)
        else:
            aspect_ratio = original_height / original_width
            output_width = width
            output_height = int(output_width * aspect_ratio)
        print(f"Resizing to: {output_width} x {output_height}")
    
    # Validate quality
    if quality < 0 or quality > 100:
        print(f"Warning: Quality {quality} out of range. Using 95")
        quality = 95
    
    print(f"JPEG quality: {quality}")
    if generate_c_array:
        print(f"Will generate C array with name: {c_array_name}")
    
    # Extract frames
    frame_count = 0
    saved_count = 0
    images_data = []  # 存储图片数据(bytes)
    
    try:
        while cap.isOpened():
            if start_frame + frame_count >= end_frame:
                break
                
            ret, frame = cap.read()
            if not ret:
                break
            
            # Resize
            if (output_width, output_height) != (original_width, original_height):
                frame = cv2.resize(frame, (output_width, output_height))
            
            # Save image to file
            filename = f"{prefix}_{saved_count:06d}.jpg"
            output_path = os.path.join(output_dir, filename)
            
            # 使用imwrite保存图片
            success = cv2.imwrite(output_path, frame, [cv2.IMWRITE_JPEG_QUALITY, quality])
            
            if success:
                # 读取保存的图片文件数据
                try:
                    with open(output_path, 'rb') as img_file:
                        img_data = img_file.read()
                    
                    images_data.append(img_data)
                    saved_count += 1
                    
                    if saved_count % 50 == 0:
                        print(f"  Saved {saved_count} images...")
                except Exception as e:
                    print(f"  Warning: Failed to read saved image {saved_count}: {e}")
            else:
                print(f"  Warning: Failed to save image {saved_count}")
            
            frame_count += 1
            
    except KeyboardInterrupt:
        print("\nInterrupted by user")
    except Exception as e:
        print(f"Error during extraction: {e}")
    finally:
        cap.release()
    
    print(f"\nImage extraction completed!")
    print(f"Total images saved: {saved_count}")
    print(f"Location: {output_dir}")
    
    # 生成C数组文件
    if generate_c_array and images_data:
        if c_array_file is None:
            c_array_file = os.path.join(output_dir, f"{c_array_name}.c")
        
        success = save_as_simple_c_array(images_data, c_array_file, c_array_name)
        
        if success:
            print(f"C array generation completed!")
        else:
            print(f"C array generation failed!")
    
    return saved_count > 0, images_data

def main():
    parser = argparse.ArgumentParser(
        description='Convert MP4 to JPEG images with ultra-simplified C array generation',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Basic conversion (no C array)
  python mp4_2_jpg.py video.mp4 -o images
  
  # With specific size and quality
  python mp4_2_jpg.py video.mp4 -o images -w 640 -ht 480 -q 90
  
  # With time range
  python mp4_2_jpg.py video.mp4 -o images -s 30 -d 10 -w 320
  
  # Generate C array file
  python mp4_2_jpg.py video.mp4 -o images -w 240 -c -n video_frames
        """
    )
    
    # Required argument
    parser.add_argument('input', help='Input video file path')
    
    # Optional arguments
    parser.add_argument('-o', '--output', default='output_frames', 
                       help='Output directory (default: output_frames)')
    parser.add_argument('-w', '--width', type=int, 
                       help='Output image width (pixels)')
    parser.add_argument('-ht', '--height', type=int, 
                       help='Output image height (pixels)')
    parser.add_argument('-q', '--quality', type=int, default=95,
                       help='JPEG quality (0-100, default: 95)')
    parser.add_argument('-s', '--start', type=float, default=0,
                       help='Start time in seconds (default: 0)')
    parser.add_argument('-d', '--duration', type=float, 
                       help='Duration to extract in seconds')
    parser.add_argument('-p', '--prefix', default='frame',
                       help='Filename prefix (default: frame)')
    
    # C array generation arguments
    parser.add_argument('-c', '--c-array', action='store_true',
                       help='Generate C array file with image data')
    parser.add_argument('-n', '--c-name', default='image_data',
                       help='C array name prefix (default: image_data)')
    parser.add_argument('-cf', '--c-file', 
                       help='Output C file path (default: <output_dir>/<c_name>.c)')
    
    # If no args, show help
    if len(sys.argv) == 1:
        parser.print_help()
        return
    
    args = parser.parse_args()
    
    # Call conversion function
    success, images_data = mp4_to_jpg(
        video_path=args.input,
        output_dir=args.output,
        width=args.width,
        height=args.height,
        quality=args.quality,
        start_time=args.start,
        duration=args.duration,
        prefix=args.prefix,
        generate_c_array=args.c_array,
        c_array_name=args.c_name,
        c_array_file=args.c_file
    )
    
    if not success:
        print("Conversion failed!")
        sys.exit(1)

if __name__ == "__main__":
    main()