#!/bin/bash

# FFMpeg Video Upscale using upscayl-bin for AI and FFmpeg Lanczos for Manual Upscaling.
# Added: Deep IDET Interlace Analysis to prevent AI artifacting.
# Concept and Testing by Glen MacArthur, coded by Google Gemini 3.

trap 'pkill -P $$; rm -rf "$TMP_DIR"; exit' SIGINT SIGTERM EXIT

BINARY="/opt/upscayl/upscayl-bin"
MODELS_DIR="/opt/upscayl/models"
INPUT_FILE="$1"
ICON="/usr/local/share/icons/custom/ffmpeg.png"
IMAGE="/usr/local/share/icons/custom/ffmpeg-app.png"

[ -z "$INPUT_FILE" ] && { echo "Usage: ./ff_upscayl_video <video_file>"; exit 1; }

DIRNAME=$(dirname -- "$INPUT_FILE")
BASENAME=$(basename -- "$INPUT_FILE" | cut -f1 -d'.')
TMP_DIR=$(mktemp -d)

handle_error() {
    yad --error --title="Pipeline Failed" --width=520 --center --text="<b>Conversion Failure</b>\n\nCheck input integrity or filter settings."
    rm -rf "$TMP_DIR"
    exit 1
}

# 1. YAD UI 
USER_OPTS=$(yad --title="FFmpeg Video Upscale" --center --form --separator="|" --width=720 \
    --window-icon="$ICON" --image="$IMAGE"  \
    --text="<b>Select Upscale Preset:</b>\n*AI models optimized for Vulkan GPU Drivers, 16Gb of RAM Recommended." \
    --field="Presets:CB" "AI: HQ 4X (realesrgan-x4plus) - [Very Slow]!AI: General/Animation 4X (realesr-animevideov3-x4) - [Slow]!AI: General/Animation 2X (realesr-animevideov3-x2) - [Fast]!non-AI: 2X (Lanczos/x264) - [Fastest]!non-AI: 2X (Lanczos/ProRes/WAV) - [Fastest]" \
    --field="--- Restoration ---:RO" "" \
    --field="Deinterlace (bwdif):CHK" "FALSE" \
    --field="Run Auto-Interlace Detection (idet):CHK" "FALSE" \
    --field="Deblock (spp):CHK" "FALSE" \
    --field="Enable Fast Denoise (hqdn3d):CHK" "FALSE" \
    --field="--- Enhancement ---:RO" "" \
    --field="Subtle Video Enhancement (EQ):CHK" "FALSE" \
    --field="Enable Smart Sharpening (CAS):CHK" "FALSE" \
    --field="Add Film Grain:CHK" "FALSE" \
    --field="Apply 1s Fade In/Out:CHK" "FALSE" \
    --button="Start Processing:0" --button="Cancel:1")

[ $? -ne 0 ] && { rm -rf "$TMP_DIR"; exit 1; }

PRESET=$(echo "$USER_OPTS" | cut -d'|' -f1)
DEINTERLACE=$(echo "$USER_OPTS" | cut -d'|' -f3)
RUN_IDET=$(echo "$USER_OPTS" | cut -d'|' -f4)
DEBLOCK=$(echo "$USER_OPTS" | cut -d'|' -f5)
DENOISE=$(echo "$USER_OPTS" | cut -d'|' -f6)
EQ_BUMP=$(echo "$USER_OPTS" | cut -d'|' -f8)
SHARPEN_ENABLED=$(echo "$USER_OPTS" | cut -d'|' -f9)
GRAIN=$(echo "$USER_OPTS" | cut -d'|' -f10)
FADES=$(echo "$USER_OPTS" | cut -d'|' -f11)

# --- IDET ANALYSIS PASS (STRICT & BASE-10 COMPLIANT) ---
CALCULATED_PARITY="auto"
if [[ "$RUN_IDET" == "TRUE" ]]; then
    IDET_LOG=$(ffmpeg -i "$INPUT_FILE" -vf "idet" -frames:v 1000 -an -f null /dev/null 2>&1)
    
    TFF_STR=$(echo "$IDET_LOG" | grep -i "Repeated Fields" | awk '{print $3}' | sed 's/[^0-9]//g' | head -n1)
    BFF_STR=$(echo "$IDET_LOG" | grep -i "Repeated Fields" | awk '{print $5}' | sed 's/[^0-9]//g' | head -n1)
    
    TFF_INT=$((10#${TFF_STR:-0}))
    BFF_INT=$((10#${BFF_STR:-0}))
    
    if [ "$TFF_INT" -ge "$BFF_INT" ]; then
        CALCULATED_PARITY="tff"
    else
        CALCULATED_PARITY="bff"
    fi
    
    if [ "$TFF_INT" -gt 0 ] || [ "$BFF_INT" -gt 0 ]; then
        DEINTERLACE="TRUE"
    fi
fi

# Define baseline components
FILTERS=()
AUDIO_FILTERS=("loudnorm")
RAW_DUR=$(ffprobe -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 "$INPUT_FILE")
TOTAL_SECONDS=$(printf "%.2f" "$RAW_DUR")

# Target Frame Rate Resolution
FPS=$(ffprobe -v error -select_streams v:0 -show_entries stream=r_frame_rate -of default=noprint_wrappers=1:nokey=1 "$INPUT_FILE" | bc -l)

# --- FIXED FRAME-ACCURATE FADE LAYOUT CALCULATIONS ---
if [[ "$FADES" == "TRUE" ]]; then
    CURR_FRAMES=$(echo "$TOTAL_SECONDS * $FPS" | bc | cut -d'.' -f1)
    F_LEN=$(echo "$FPS" | bc | cut -d'.' -f1)
    O_START=$(( CURR_FRAMES - F_LEN ))
    [[ "$O_START" -lt 0 ]] && O_START=0
    
    AUDIO_OUT_START=$(echo "$TOTAL_SECONDS - 1.0" | bc -l)
fi

case "$PRESET" in
    *"HQ 4X"*)            MODEL="realesrgan-x4plus"; SCALE="4"; MODE="AI" ;;
    *"Animation 4X"*)    MODEL="realesr-animevideov3-x4"; SCALE="4"; MODE="AI" ;;
    *"Animation 2X"*)    MODEL="realesr-animevideov3-x2"; SCALE="2"; MODE="AI" ;;
    *)                    MODE="MANUAL" ;;
esac

# Append Pre-processing filters
[[ "$DEINTERLACE" == "TRUE" ]] && FILTERS+=("bwdif=mode=1:parity=${CALCULATED_PARITY}:deint=1")
[[ "$DEBLOCK" == "TRUE" ]] && FILTERS+=("spp=6")
[[ "$DENOISE" == "TRUE" ]] && FILTERS+=("hqdn3d=1.5:1.0:3.0:2.0")
[[ "$EQ_BUMP" == "TRUE" ]] && FILTERS+=("eq=contrast=1.1:saturation=1.1")
[[ "$SHARPEN_ENABLED" == "TRUE" ]] && FILTERS+=("cas=strength=0.5")
[[ "$GRAIN" == "TRUE" ]] && FILTERS+=("noise=alls=10:allf=t")

# Processing Logic
if [[ "$MODE" == "AI" ]]; then
    notify-send "FFmpeg" "Extracting source frames..." --icon="$ICON"
    mkdir -p "$TMP_DIR/upscaled"
    
    # FIX: Isolated extraction pipeline to prevent empty video track failures
    ffmpeg -hide_banner -i "$INPUT_FILE" -q:v 1 "$TMP_DIR/frame_%08d.png" -an -y > /dev/null 2>&1
    ffmpeg -hide_banner -i "$INPUT_FILE" -vn -c:a flac "$TMP_DIR/audio.mkv" -y > /dev/null 2>&1
    
    TOTAL_FRAMES=$(find "$TMP_DIR" -name "frame_*.png" | wc -l)
    [[ "$TOTAL_FRAMES" -eq 0 ]] && handle_error
    
    {
        count=0
        for f in "$TMP_DIR"/frame_*.png; do
            "$BINARY" -i "$f" -o "$TMP_DIR/upscaled/$(basename "$f")" -m "$MODELS_DIR" -n "$MODEL" -s "$SCALE" -f "png" > /dev/null 2>&1
            ((count++))
            echo "$((count * 100 / TOTAL_FRAMES))"
            echo "# Upscaling [$BASENAME] - Frame $count / $TOTAL_FRAMES"
        done
    } | yad --progress --title="FFmpeg Video Upscale" --window-icon="$ICON" --image="$IMAGE" --center --width=600 --auto-close

    [ -z "$(ls -A "$TMP_DIR/upscaled")" ] && handle_error
    
    notify-send "Upscaler" "Assembling target video..." --icon="$ICON"
    
    # Append post-upscale video fades to the filter array
    [[ "$FADES" == "TRUE" ]] && FILTERS+=("fade=t=in:st=0:n=$F_LEN" "fade=t=out:st=$O_START:n=$F_LEN")
    [[ "$FADES" == "TRUE" ]] && AUDIO_FILTERS+=("afade=t=in:st=0:d=1" "afade=t=out:st=${AUDIO_OUT_START}:d=1")
    
    VF_ARG=$(IFS=,; echo "${FILTERS[*]}")
    [[ -n "$VF_ARG" ]] && VF_ARG="-vf $VF_ARG"
    
    AF_ARG=$(IFS=,; echo "${AUDIO_FILTERS[*]}")
    [[ -n "$AF_ARG" ]] && AF_ARG="-af $AF_ARG"
    
    ffmpeg -hide_banner -loglevel error -framerate "$FPS" -i "$TMP_DIR/upscaled/frame_%08d.png" -i "$TMP_DIR/audio.mkv" \
        $VF_ARG -c:v libx264 -crf 18 -c:a aac -b:a 320k $AF_ARG -vsync 1 -async 1 \
        "${DIRNAME}/${BASENAME}_upscaled.mp4" -y
    [ $? -ne 0 ] && handle_error
else
    # Manual Scaling Layout Pipeline
    FILTERS=("scale=iw*2:ih*2:flags=lanczos+accurate_rnd" "${FILTERS[@]}")
    [[ "$FADES" == "TRUE" ]] && FILTERS+=("fade=t=in:st=0:n=$F_LEN" "fade=t=out:st=$O_START:n=$F_LEN")
    [[ "$FADES" == "TRUE" ]] && AUDIO_FILTERS+=("afade=t=in:st=0:d=1" "afade=t=out:st=${AUDIO_OUT_START}:d=1")
    
    VF_ARG="-vf $(IFS=,; echo "${FILTERS[*]}")"
    AF_ARG="-af $(IFS=,; echo "${AUDIO_FILTERS[*]}")"
    
    if [[ "$PRESET" == *"ProRes/WAV"* ]]; then
        CMD="ffmpeg -i '$INPUT_FILE' $VF_ARG -c:v prores_ks -profile:v 3 -c:a pcm_s16le $AF_ARG '${DIRNAME}/${BASENAME}_prores.mov' -y"
    else
        CMD="ffmpeg -i '$INPUT_FILE' $VF_ARG -c:v libx264 -crf 18 -preset slow -c:a aac -b:a 320k $AF_ARG '${DIRNAME}/${BASENAME}_manual.mp4' -y"
    fi
    MANUAL_RUNNER="$TMP_DIR/manual_run.sh"
    echo -e "#!/bin/bash\n$CMD\necho\nread -p 'Press Enter to close...'" > "$MANUAL_RUNNER"
    chmod +x "$MANUAL_RUNNER"
    terminology -T "FFmpeg Manual Upscale" -e "$MANUAL_RUNNER"
fi

rm -rf "$TMP_DIR"
notify-send "Upscale Complete" "File processed successfully." --icon="$ICON"
exit 0