#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
BIN_DIR="${ROOT_DIR}/bin"
MODEL_DIR="${ROOT_DIR}/models"
TMP_DIR="${ROOT_DIR}/.tmp"
FORCE=0
RELEASE_VERSION="${REALESRGAN_RELEASE_VERSION:-v0.2.5.0}"
ASSET_NAME="${REALESRGAN_RELEASE_ASSET:-realesrgan-ncnn-vulkan-20220424-ubuntu.zip}"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --force)
      FORCE=1
      shift
      ;;
    --version)
      RELEASE_VERSION="${2:?--version requires a value}"
      shift 2
      ;;
    *)
      echo "Unknown argument: $1" >&2
      exit 1
      ;;
  esac
done

if [[ "${RELEASE_VERSION}" == "latest" ]]; then
  RELEASE_URL="${REALESRGAN_RELEASE_URL:-https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/${ASSET_NAME}}"
else
  RELEASE_URL="${REALESRGAN_RELEASE_URL:-https://github.com/xinntao/Real-ESRGAN/releases/download/${RELEASE_VERSION}/${ASSET_NAME}}"
fi

mkdir -p "${BIN_DIR}" "${MODEL_DIR}" "${TMP_DIR}"

if [[ -f "${BIN_DIR}/realesrgan-ncnn-vulkan" && "${FORCE}" -ne 1 ]]; then
  echo "Binary already exists at ${BIN_DIR}/realesrgan-ncnn-vulkan. Use --force to replace it."
  exit 0
fi

ARCHIVE_PATH="${TMP_DIR}/${ASSET_NAME}"
EXTRACT_DIR="${TMP_DIR}/realesrgan-extract"
rm -rf "${EXTRACT_DIR}"
mkdir -p "${EXTRACT_DIR}"

echo "Downloading ${RELEASE_URL}"
curl -L --fail --output "${ARCHIVE_PATH}" "${RELEASE_URL}"

if [[ ! -f "${ARCHIVE_PATH}" ]]; then
  echo "Download failed: archive not found at ${ARCHIVE_PATH}" >&2
  exit 1
fi

unzip -o "${ARCHIVE_PATH}" -d "${EXTRACT_DIR}" >/dev/null

BINARY_PATH="$(find "${EXTRACT_DIR}" -type f -name 'realesrgan-ncnn-vulkan' | head -n 1)"

if [[ -z "${BINARY_PATH}" ]]; then
  echo "Could not find realesrgan-ncnn-vulkan in the downloaded archive." >&2
  exit 1
fi

install -m 0755 "${BINARY_PATH}" "${BIN_DIR}/realesrgan-ncnn-vulkan"

find "${EXTRACT_DIR}" -type f \( -name '*.param' -o -name '*.bin' \) -print0 | while IFS= read -r -d '' model_file; do
  install -m 0644 "${model_file}" "${MODEL_DIR}/$(basename "${model_file}")"
done

echo "Installed binary: ${BIN_DIR}/realesrgan-ncnn-vulkan"
echo "Installed models in: ${MODEL_DIR}"