#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
H264文件检查工具
检查生成的test.h264文件是否符合播放要求
"""

import os
import sys

def check_h264_file(file_path):
    """检查H264文件是否符合播放要求"""
    print(f"检查文件: {file_path}\n")
    
    # 1. 检查文件是否存在
    if not os.path.exists(file_path):
        print("❌ 文件不存在！")
        print(f"   预期路径: {file_path}")
        print("\n提示：")
        print("   - 如果是Android设备，需要先通过adb pull下载文件")
        print(f"   - 或者将文件复制到当前目录")
        return False
    
    # 2. 检查文件大小
    file_size = os.path.getsize(file_path)
    print(f"✓ 文件大小: {file_size} 字节 ({file_size / 1024:.2f} KB)")
    
    if file_size == 0:
        print("❌ 文件为空！")
        return False
    
    if file_size < 100:
        print("⚠️  警告: 文件太小，可能只有SPS/PPS或1帧数据")
    
    # 3. 读取文件内容并分析
    try:
        with open(file_path, 'rb') as f:
            data = f.read()
    except Exception as e:
        print(f"❌ 读取文件失败: {e}")
        return False
    
    # 检查NALU单元
    nal_start_code_3 = bytes([0x00, 0x00, 0x01])
    nal_start_code_4 = bytes([0x00, 0x00, 0x00, 0x01])
    
    nalu_list = []
    i = 0
    
    # 查找所有NALU单元
    while i < len(data):
        # 查找起始码
        found = False
        
        # 查找4字节起始码
        if i + 4 <= len(data) and data[i:i+4] == nal_start_code_4:
            start_pos = i + 4
            found = True
        # 查找3字节起始码
        elif i + 3 <= len(data) and data[i:i+3] == nal_start_code_3:
            # 确保前面不是0x00 (避免误判)
            if i == 0 or data[i-1] != 0x00:
                start_pos = i + 3
                found = True
        
        if found:
            # 查找下一个起始码
            next_start = -1
            
            # 先找4字节起始码
            for j in range(start_pos, len(data) - 3):
                if j + 4 <= len(data) and data[j:j+4] == nal_start_code_4:
                    next_start = j
                    break
            
            # 如果没找到4字节的，找3字节的
            if next_start == -1:
                for j in range(start_pos, len(data) - 2):
                    if j + 3 <= len(data) and data[j:j+3] == nal_start_code_3:
                        # 确保前面不是0x00
                        if j == 0 or data[j-1] != 0x00:
                            next_start = j
                            break
            
            # 提取NALU数据
            if next_start == -1:
                nalu_data = data[start_pos:]
            else:
                nalu_data = data[start_pos:next_start]
            
            if len(nalu_data) > 0:
                # 获取NALU类型 (第一个字节的低5位)
                nal_type = nalu_data[0] & 0x1F
                nalu_list.append({
                    'type': nal_type,
                    'size': len(nalu_data),
                    'name': get_nalu_type_name(nal_type)
                })
            
            i = next_start if next_start != -1 else len(data)
        else:
            i += 1
    
    # 4. 分析NALU单元
    print(f"\n✓ 找到 {len(nalu_list)} 个NALU单元\n")
    
    if len(nalu_list) == 0:
        print("❌ 未找到任何NALU单元！文件格式可能不正确")
        return False
    
    # 统计不同类型的NALU
    nal_type_count = {}
    for nalu in nalu_list:
        nal_type = nalu['type']
        if nal_type not in nal_type_count:
            nal_type_count[nal_type] = 0
        nal_type_count[nal_type] += 1
    
    print("NALU类型统计:")
    for nal_type in sorted(nal_type_count.keys()):
        count = nal_type_count[nal_type]
        name = get_nalu_type_name(nal_type)
        print(f"  {name} (类型{nal_type}): {count} 个")
    
    # 5. 检查关键要求
    print("\n检查项:")
    
    has_sps = 7 in nal_type_count
    has_pps = 8 in nal_type_count
    has_idr = 5 in nal_type_count
    has_non_idr = 1 in nal_type_count
    
    # SPS/PPS检查
    if has_sps:
        print(f"  ✓ 包含SPS (序列参数集) - {nal_type_count[7]} 个")
    else:
        print("  ❌ 缺少SPS (序列参数集) - 必需！")
    
    if has_pps:
        print(f"  ✓ 包含PPS (图像参数集) - {nal_type_count[8]} 个")
    else:
        print("  ❌ 缺少PPS (图像参数集) - 必需！")
    
    # 关键帧检查
    if has_idr:
        print(f"  ✓ 包含IDR关键帧 - {nal_type_count[5]} 个")
    else:
        print("  ❌ 缺少IDR关键帧")
    
    # 非关键帧检查
    if has_non_idr:
        print(f"  ✓ 包含非IDR帧 - {nal_type_count[1]} 个")
    else:
        print("  ⚠️  没有非IDR帧（只有关键帧）")
    
    # 6. 总体评估
    print("\n" + "="*50)
    can_play = has_sps and has_pps and has_idr
    
    if can_play:
        if len(nalu_list) >= 3:
            print("✅ 文件应该可以播放！")
            print(f"   包含完整的SPS/PPS和 {len(nalu_list)} 个NALU单元")
        else:
            print("⚠️  文件结构完整，但帧数较少")
            print("   建议录制更长时间以获得更多帧")
    else:
        print("❌ 文件可能无法播放")
        if not has_sps or not has_pps:
            print("   原因: 缺少SPS/PPS参数集")
        if not has_idr:
            print("   原因: 缺少IDR关键帧")
    
    print("="*50)
    
    return can_play

def get_nalu_type_name(nal_type):
    """获取NALU类型名称"""
    nal_names = {
        1: "非IDR编码片",
        2: "编码片数据分区A",
        3: "编码片数据分区B",
        4: "编码片数据分区C",
        5: "IDR图像编码片",
        6: "SEI (补充增强信息)",
        7: "SPS (序列参数集)",
        8: "PPS (图像参数集)",
        9: "访问单元分隔符",
        10: "序列结束",
        11: "流结束",
        12: "填充数据"
    }
    return nal_names.get(nal_type, f"未知类型{nal_type}")

if __name__ == "__main__":
    # 默认文件路径
    default_path = "test.h264"
    
    # 如果提供了命令行参数，使用参数作为文件路径
    file_path = sys.argv[1] if len(sys.argv) > 1 else default_path
    
    print("H264文件检查工具")
    print("="*50)
    
    success = check_h264_file(file_path)
    
    sys.exit(0 if success else 1)