#!/bin/bash

# =================配置区域=================
# 基础路径配置
BASE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
INPUT_DIR="$BASE_DIR/videos_to_process"
OUTPUT_DIR="$BASE_DIR/enhanced_videos"
APP_DIR="$BASE_DIR/app"
EXEC_PATH="$APP_DIR/realesrgan-ncnn-vulkan"

# Real-ESRGAN 配置
# 模型可选: realesrgan-x4plus (适合真实影像), realesrgan-x4plus-anime (适合动漫)
MODEL_NAME="realesrgan-x4plus"
SCALE=4
# 显卡ID，自动探测一般填0，如果多卡可修改
GPU_ID=0 

# 下载链接
DOWNLOAD_URL="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-ubuntu.zip"
ZIP_FILE="realesrgan-ncnn-vulkan-20220424-ubuntu.zip"

# =========================================

# 颜色定义
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m' # No Color

# 0. 检查依赖 (ffmpeg)
if ! command -v ffmpeg &> /dev/null; then
    echo -e "${RED}Error: ffmpeg 未安装。${NC}"
    echo "请运行: sudo apt update && sudo apt install ffmpeg"
    exit 1
fi

# 1. 准备目录和环境
mkdir -p "$OUTPUT_DIR"
mkdir -p "$APP_DIR"

# 检查并下载 Real-ESRGAN
if [ ! -f "$EXEC_PATH" ]; then
    echo -e "${YELLOW}Real-ESRGAN 可执行文件未找到，准备下载...${NC}"
    
    # 进入 app 目录下载
    cd "$APP_DIR" || exit
    
    if [ -f "$ZIP_FILE" ]; then
        rm "$ZIP_FILE"
    fi
    
    echo "正在下载: $DOWNLOAD_URL"
    wget -q --show-progress "$DOWNLOAD_URL"
    
    if [ $? -ne 0 ]; then
        echo -e "${RED}下载失败，请检查网络连接。${NC}"
        exit 1
    fi

    echo "正在解压..."
    unzip -q "$ZIP_FILE"
    
    # 赋予执行权限
    chmod +x realesrgan-ncnn-vulkan
    
    # 清理压缩包
    rm "$ZIP_FILE"
    
    echo -e "${GREEN}Real-ESRGAN 安装完成！${NC}"
    
    # 回到基础目录
    cd "$BASE_DIR" || exit
else
    echo -e "${GREEN}Real-ESRGAN 已安装，跳过下载。${NC}"
fi

# 2. 开始扫描视频并处理
shopt -s nullglob
video_files=("$INPUT_DIR"/*.mp4 "$INPUT_DIR"/*.mov "$INPUT_DIR"/*.avi "$INPUT_DIR"/*.mkv)

if [ ${#video_files[@]} -eq 0 ]; then
    echo -e "${YELLOW}在 $INPUT_DIR 中没有找到支持的视频文件。${NC}"
    exit 0
fi

echo -e "${GREEN}找到 ${#video_files[@]} 个视频文件，开始处理...${NC}"

# 进度条绘制用的全长字符串
FULL_BAR="##################################################"
BAR_WIDTH=50

for INPUT_VIDEO in "${video_files[@]}"; do
    BASE_NAME=$(basename "$INPUT_VIDEO")
    FILENAME="${BASE_NAME%.*}"
    OUTPUT_VIDEO="$OUTPUT_DIR/${FILENAME}_x${SCALE}.mp4"
    
    echo "----------------------------------------------------------------"
    echo -e "正在处理: ${YELLOW}$BASE_NAME${NC}"
    echo "----------------------------------------------------------------"

    # 创建临时目录
    TMP_FRAMES=$(mktemp -d -p "$BASE_DIR" "tmp_frames_XXXXXX")
    OUT_FRAMES=$(mktemp -d -p "$BASE_DIR" "out_frames_XXXXXX")
    
    # Step 1: 拆帧
    echo -e "${YELLOW}[1/3] 正在提取帧 (使用 ffmpeg)...${NC}"
    # 使用 png 格式以保证无损画质传入 AI
    ffmpeg -v error -stats -i "$INPUT_VIDEO" -qscale:v 1 -qmin 1 -fps_mode passthrough "$TMP_FRAMES/frame%08d.png"
    
    if [ $? -ne 0 ]; then
        echo -e "${RED}拆帧失败，跳过此视频。${NC}"
        rm -rf "$TMP_FRAMES" "$OUT_FRAMES"
        continue
    fi

    # 统计总帧数 (ls -1U 速度较快，不排序)
    TOTAL_FRAMES=$(ls -1U "$TMP_FRAMES" | wc -l)
    echo "总帧数: $TOTAL_FRAMES"

    # Step 2: AI 超分
    echo -e "${YELLOW}[2/3] 正在进行 AI 超分 (Real-ESRGAN GPU)...${NC}"
    echo "模型: $MODEL_NAME | 放大倍数: x$SCALE"
    
    # 1. 以后台方式启动 (&)，并将标准输出和错误输出都扔进黑洞
    "$EXEC_PATH" \
        -i "$TMP_FRAMES" \
        -o "$OUT_FRAMES" \
        -n "$MODEL_NAME" \
        -s "$SCALE" \
        -g "$GPU_ID" \
        -f png > /dev/null 2>&1 &
    
    PID=$! # 获取进程ID

    # 2. 循环监控进度
    while kill -0 $PID 2> /dev/null; do
        # 统计已生成的文件数
        CURRENT_FRAMES=$(ls -1U "$OUT_FRAMES" | wc -l)
        
        # 计算百分比
        if [ "$TOTAL_FRAMES" -gt 0 ]; then
            PERCENT=$(( CURRENT_FRAMES * 100 / TOTAL_FRAMES ))
        else
            PERCENT=0
        fi
        
        # 防止文件系统统计延迟导致超过 100%
        if [ "$PERCENT" -gt 100 ]; then PERCENT=100; fi
        
        # 计算进度条填充长度
        FILLED_LEN=$(( PERCENT * BAR_WIDTH / 100 ))
        
        # 打印进度条 (\r 回车不换行)
        printf "\r[${GREEN}%-${BAR_WIDTH}s${NC}] %3d%% (%d/%d)" "${FULL_BAR:0:FILLED_LEN}" "$PERCENT" "$CURRENT_FRAMES" "$TOTAL_FRAMES"
        
        sleep 1
    done
    
    # 等待进程彻底退出并获取退出码
    wait $PID
    EXIT_CODE=$?
    
    echo "" # 进度条跑完换个行

    if [ $EXIT_CODE -ne 0 ]; then
        echo -e "${RED}AI 超分失败，请检查显卡驱动或显存是否足够。${NC}"
        rm -rf "$TMP_FRAMES" "$OUT_FRAMES"
        continue
    fi

    # Step 3: 合成视频
    echo -e "${YELLOW}[3/3] 正在合成视频 (High Quality H.265)...${NC}"
    
    # 获取原视频帧率
    FPS=$(ffprobe -v error -select_streams v -of default=noprint_wrappers=1:nokey=1 -show_entries stream=r_frame_rate "$INPUT_VIDEO")
    
    # 修复点：给 map 参数加上双引号 "1:a:0?"，防止 nullglob 导致参数丢失
    ffmpeg -v error -stats \
        -framerate "$FPS" -i "$OUT_FRAMES/frame%08d.png" \
        -i "$INPUT_VIDEO" -map 0:v:0 -map "1:a:0?" -c:a copy \
        -c:v libx265 -preset veryslow -crf 18 -pix_fmt yuv420p -tag:v hvc1 \
        "$OUTPUT_VIDEO"
        
    if [ $? -ne 0 ]; then
        echo -e "${YELLOW}H.265 编码失败，尝试回退到 H.264...${NC}"
        # 修复点：fallback 命令同样加上引号
        ffmpeg -v error -stats \
            -framerate "$FPS" -i "$OUT_FRAMES/frame%08d.png" \
            -i "$INPUT_VIDEO" -map 0:v:0 -map "1:a:0?" -c:a copy \
            -c:v libx264 -preset veryslow -crf 18 -pix_fmt yuv420p \
            "$OUTPUT_VIDEO"
    fi

    # 清理临时文件
    echo -e "${GREEN}清理临时文件...${NC}"
    rm -rf "$TMP_FRAMES" "$OUT_FRAMES"
    
    echo -e "${GREEN}完成! 输出文件: $OUTPUT_VIDEO${NC}"
done

echo -e "${GREEN}所有任务已完成。${NC}"