#!/bin/bash

# Image Upscale using Upscayl with FFmpeg for Filtering
# Concept and Testing by Glen MacArthur, coded by Google Gemini 3
BINARY="/opt/upscayl/upscayl-bin"
MODELS_DIR="/opt/upscayl/models"
ICON="/usr/local/share/icons/custom/resize.png"
IMAGE="/usr/local/share/icons/custom/resize-app.png"

INPUT_FILE="$1"
[ -z "$INPUT_FILE" ] && exit 1

# File Prep
DIRNAME=$(dirname -- "$INPUT_FILE")
FILENAME=$(basename -- "$INPUT_FILE")
BASENAME="${FILENAME%.*}"
TMP_FILE=$(mktemp --suffix=.png)

# 3. YAD UI with Presets & Filters
USER_OPTS=$(yad --title="Image Upscale" --window-icon="$ICON" --image="$IMAGE" \
    --form --separator="|" --width=720 --center \
    --text="<b>Select Upscale Preset:</b>\n*AI models optimized for Vulkan GPU Drivers" \
    --field="Select Preset:CB" "Highest Quality 4X (clean sharp sources)!General 4X (good quality sources)!General/Digital 4X (anime or digital sources)!General/Digital 2X (general use 2X)" \
    --field="Source File:RO" "$FILENAME" \
    --field="Output Format:CB" "png!jpg!webp" \
    --field="Subtle Color Enhancement (EQ):CHK" "FALSE" \
    --field="Enable Smart Sharpening (CAS):CHK" "FALSE" \
    --field="Enable Film Grain:CHK" "FALSE" \
    --button="Upscale:0" --button="Cancel:1")

[ $? -ne 0 ] && exit 1

# Parse YAD output
PRESET=$(echo "$USER_OPTS" | cut -d'|' -f1)
FORMAT=$(echo "$USER_OPTS" | cut -d'|' -f3)
EQ_BUMP=$(echo "$USER_OPTS" | cut -d'|' -f4)
SHARPEN=$(echo "$USER_OPTS" | cut -d'|' -f5)
GRAIN=$(echo "$USER_OPTS" | cut -d'|' -f6)

case "$PRESET" in
    "Highest Quality 4X (clean sharp sources)") MODEL="4x-UltraSharp-fp16"; SCALE="4" ;;
    "General 4X (good quality sources)") MODEL="realesrgan-x4plus"; SCALE="4" ;;
    "General/Digital 4X (anime or digital sources)") MODEL="realesrgan-x4plus-anime"; SCALE="4" ;;
    "General/Digital 2X (general use 2X)") MODEL="realesr-animevideov3-x2"; SCALE="2" ;;
esac

# 4. FILTER CONSTRUCTION
FILTERS=()
[[ "$EQ_BUMP" == "TRUE" ]] && FILTERS+=("eq=contrast=1.1:saturation=1.1")
[[ "$SHARPEN" == "TRUE" ]] && FILTERS+=("cas=strength=0.5")
[[ "$GRAIN" == "TRUE" ]] && FILTERS+=("noise=alls=10:allf=t")

FINAL_OUTPUT="${DIRNAME}/${BASENAME}_${SCALE}x.${FORMAT}"

# 5. AI Execution (to Temp file)
stdbuf -oL "$BINARY" -i "$INPUT_FILE" -o "$TMP_FILE" -s "$SCALE" -m "$MODELS_DIR" -n "$MODEL" -f "png" 2>&1 | \
    sed -u 's/%//g' | \
    yad --progress --title="Upscaling Image..." \
    --center --width=500 --window-icon="$ICON" \
    --text="Processing AI Stage: $FILENAME" \
    --auto-close --no-buttons

# 6. Post-Processing (Filters)
if [ ${#FILTERS[@]} -gt 0 ]; then
    notify-send "Finalizing" "Applying color/sharpening filters..." --icon="$ICON"
    
    VF_ARG="-vf $(IFS=,; echo "${FILTERS[*]}")"
    # Using -hide_banner and -loglevel error keeps the background noise down
    ffmpeg -hide_banner -loglevel error -i "$TMP_FILE" $VF_ARG -q:v 1 "$FINAL_OUTPUT" -y
else
    mv "$TMP_FILE" "$FINAL_OUTPUT"
fi

# 7. Cleanup & Notification
rm -f "$TMP_FILE"
if [ -s "$FINAL_OUTPUT" ]; then
    notify-send "Upscale Complete" "Saved to: $(basename "$FINAL_OUTPUT")" --icon="$ICON"
fi