#!/bin/bash

# 初始化清理标记
CLEANUP=1

# 检查输入参数
if [ $# -lt 1 ] || [ $# -gt 2 ]; then
    echo "用法: $0 <m3u8文件路径> [--debug]"
    exit 1
fi

# 解析参数
INPUT_M3U8="$1"
if [ $# -eq 2 ]; then
    if [ "$2" == "--debug" ]; then
        CLEANUP=0
        echo "调试模式已启用，临时文件将保留"
    else
        echo "错误: 第二个参数必须是--debug"
        exit 1
    fi
fi

# 验证输入文件
if [ ! -f "$INPUT_M3U8" ]; then
    echo "错误: 指定的m3u8文件不存在"
    exit 2
fi

# 生成路径信息
INPUT_DIR=$(dirname "$INPUT_M3U8")
BASE_NAME=$(basename "${INPUT_M3U8%.*}")
OUTPUT_MP4="${INPUT_DIR}/${BASE_NAME}.mp4"

# 创建临时目录
declare -g TEMP_DIR DECRYPTED_DIR FIXED_DIR  # 用于trap函数访问
TEMP_DIR="${INPUT_DIR}/temp_${BASE_NAME}_downloads"
DECRYPTED_DIR="${INPUT_DIR}/temp_${BASE_NAME}_decrypted"
FIXED_DIR="${INPUT_DIR}/temp_${BASE_NAME}_fixed"
mkdir -p "$TEMP_DIR" "$DECRYPTED_DIR" "$FIXED_DIR"

# 定义清理函数
cleanup() {
    if [ $CLEANUP -eq 1 ]; then
        echo "清理临时文件..."
        rm -rf "$TEMP_DIR" "$DECRYPTED_DIR" "$FIXED_DIR"
    else
        echo "临时文件保留在:"
        [ -d "$TEMP_DIR" ] && echo "  - 加密分片: $(realpath "$TEMP_DIR")"
        [ -d "$DECRYPTED_DIR" ] && echo "  - 解密文件: $(realpath "$DECRYPTED_DIR")"
        [ -d "$FIXED_DIR" ] && echo "  - 修复文件: $(realpath "$FIXED_DIR")"
    fi
}
trap cleanup EXIT

# 1. 提取AES KEY和IV
KEY_URI=$(grep -o 'URI="[^"]*"' "$INPUT_M3U8" | cut -d'"' -f2)
IV_HEX=$(grep -o 'IV=0x[0-9A-Fa-f]*' "$INPUT_M3U8" | cut -d'x' -f2)

# 转换IV为二进制格式
echo -n "$IV_HEX" | xxd -r -p > "${TEMP_DIR}/iv.bin"

# 2. 下载AES KEY
echo "下载加密密钥..."
if ! curl -s -o "${TEMP_DIR}/key.bin" "$KEY_URI"; then
    echo "错误: 密钥下载失败"
    exit 3
fi

# 3. 安全提取分片URL
mapfile -t SEGMENT_URLS < <(grep -E '^https?://' "$INPUT_M3U8")

# 4. 处理分片
for i in "${!SEGMENT_URLS[@]}"; do
    SEG_NUM=$((i+1))
    ENC_FILE="${TEMP_DIR}/${BASE_NAME}_seg${SEG_NUM}.ts"
    DEC_FILE="${DECRYPTED_DIR}/${BASE_NAME}_decrypted_seg${SEG_NUM}.ts"
    FIXED_FILE="${FIXED_DIR}/${BASE_NAME}_fixed_seg${SEG_NUM}.ts"

    # 下载加密分片
    echo "($SEG_NUM/${#SEGMENT_URLS[@]}) 下载分片..."
    if ! curl -fsS -o "$ENC_FILE" "${SEGMENT_URLS[$i]}"; then
        echo "错误: 分片 $SEG_NUM 下载失败"
        exit 4
    fi

    # 解密分片
    echo "($SEG_NUM/${#SEGMENT_URLS[@]}) 解密分片..."
    if ! openssl aes-128-cbc -d \
        -in "$ENC_FILE" \
        -out "$DEC_FILE" \
        -nosalt \
        -iv "$(xxd -p -c16 "${TEMP_DIR}/iv.bin")" \
        -K "$(xxd -p -c16 "${TEMP_DIR}/key.bin")"; then
        echo "错误: 分片 $SEG_NUM 解密失败"
        exit 5
    fi

    # 修复视频分片
    echo "($SEG_NUM/${#SEGMENT_URLS[@]}) 修复分片..."
    if ! ffmpeg -hide_banner -y \
        -i "$DEC_FILE" \
        -c copy \
        -fflags +genpts \
        -avoid_negative_ts make_zero \
        -strict experimental \
        "$FIXED_FILE" 2>/dev/null; then
        echo "错误: 分片 $SEG_NUM 修复失败"
        exit 6
    fi
done

# 5. 合并修复后的文件
echo "合并修复片段..."
cat "${FIXED_DIR}/${BASE_NAME}_fixed_seg"*.ts > "${TEMP_DIR}/${BASE_NAME}_merged.ts"

# 6. 生成最终文件（更新部分）
echo "生成Chrome兼容视频..."
if ! ffmpeg -hide_banner -y \
    -i "${TEMP_DIR}/${BASE_NAME}_merged.ts" \
    -c:v libx264 \
    -preset medium \
    -profile:v high \
    -level 4.0 \
    -pix_fmt yuv420p \
    -crf 23 \
    -c:a aac \
    -b:a 128k \
    -ar 44100 \
    -movflags +faststart \
    -f mp4 \
    "$OUTPUT_MP4" 2>/dev/null; then
    echo "错误: 最终视频生成失败"
    exit 7
fi

echo -e "\n处理成功！输出文件: $(realpath "$OUTPUT_MP4")"