# -*- coding: utf-8 -*-
"""
视频按大小分割脚本
将视频分割为指定大小（默认180MB）的片段
"""
import os
import subprocess
import json
from pathlib import Path
import imageio_ffmpeg

# 配置
INPUT_DIR = Path(__file__).parent / "原视频"
OUTPUT_DIR = Path(__file__).parent / "分割后"
TARGET_DURATION_MIN = 30  # 目标片段时长(分钟)
MERGE_THRESHOLD_MIN = 10  # 最后一段小于此时长则合并到前一段
FFMPEG_PATH = imageio_ffmpeg.get_ffmpeg_exe()

def get_video_info(video_path):
    """获取视频信息：时长和文件大小（快速，只读取头部）"""
    import re
    # 只读取文件头部信息，不处理整个视频
    cmd = [
        FFMPEG_PATH,
        '-i', str(video_path),
    ]

    result = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8')
    # 从stderr中解析时长 (ffmpeg输出信息在stderr)
    duration_match = re.search(r'Duration: (\d+):(\d+):(\d+\.?\d*)', result.stderr)
    if duration_match:
        h, m, s = duration_match.groups()
        duration = int(h) * 3600 + int(m) * 60 + float(s)
    else:
        raise ValueError("无法获取视频时长")

    file_size = os.path.getsize(video_path)
    return duration, file_size

def split_video(video_path):
    """将视频按时长分割"""
    video_path = Path(video_path)
    video_name = video_path.stem
    video_ext = video_path.suffix

    # 获取视频信息
    duration, file_size = get_video_info(video_path)
    file_size_mb = file_size / (1024 * 1024)
    duration_min = duration / 60

    print(f"处理: {video_path.name}", flush=True)
    print(f"  时长: {duration_min:.1f}分钟, 大小: {file_size_mb:.1f}MB", flush=True)

    target_duration = TARGET_DURATION_MIN * 60  # 转为秒
    merge_threshold = MERGE_THRESHOLD_MIN * 60  # 转为秒

    # 如果视频短于目标时长，直接复制
    if duration <= target_duration:
        print(f"  视频短于{TARGET_DURATION_MIN}分钟，无需分割")
        output_path = OUTPUT_DIR / f"{video_name}-1{video_ext}"
        import shutil
        shutil.copy2(video_path, output_path)
        print(f"  已复制到: {output_path.name}")
        return [output_path]

    # 计算片段时间点
    segments = []
    start_time = 0
    while start_time < duration:
        remaining = duration - start_time
        if remaining <= target_duration + merge_threshold:
            # 剩余时间不多，作为最后一段
            segments.append((start_time, remaining))
            break
        else:
            segments.append((start_time, target_duration))
            start_time += target_duration

    print(f"  分割为 {len(segments)} 个片段，每段约 {TARGET_DURATION_MIN} 分钟", flush=True)

    output_files = []
    for segment_index, (start_time, actual_duration) in enumerate(segments, 1):
        output_path = OUTPUT_DIR / f"{video_name}-{segment_index}{video_ext}"

        cmd = [
            FFMPEG_PATH,
            '-y',  # 覆盖输出文件
            '-i', str(video_path),
            '-ss', str(start_time),
            '-t', str(actual_duration),
            '-c', 'copy',  # 不重新编码，速度快
            '-avoid_negative_ts', 'make_zero',
            str(output_path)
        ]

        print(f"  分割片段 {segment_index}: {start_time/60:.1f}min - {(start_time + actual_duration)/60:.1f}min", flush=True)

        subprocess.run(cmd, capture_output=True, encoding='utf-8')

        # 检查输出文件大小
        if output_path.exists():
            out_size = os.path.getsize(output_path) / (1024 * 1024)
            print(f"    -> {output_path.name} ({out_size:.1f}MB, {actual_duration/60:.1f}min)")
            output_files.append(output_path)

    return output_files

def main():
    print("=" * 50)
    print("视频分割工具")
    print(f"输入目录: {INPUT_DIR}")
    print(f"输出目录: {OUTPUT_DIR}")
    print(f"目标时长: {TARGET_DURATION_MIN}分钟/段")
    print("=" * 50)

    # 创建输出目录
    OUTPUT_DIR.mkdir(exist_ok=True)

    # 支持的视频格式
    video_extensions = {'.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm'}

    # 获取所有视频文件
    video_files = [f for f in INPUT_DIR.iterdir()
                   if f.is_file() and f.suffix.lower() in video_extensions]

    if not video_files:
        print("未找到视频文件!")
        return

    print(f"找到 {len(video_files)} 个视频文件\n")

    # 处理每个视频
    for video_file in video_files:
        try:
            split_video(video_file)
            print()
        except Exception as e:
            print(f"  错误: {e}\n")

    print("=" * 50)
    print("处理完成!")
    print("=" * 50)

if __name__ == "__main__":
    main()