#!/usr/bin/env bash
# onwx-Image for AMD (Debian/Ubuntu): Triton-XDNA and native Z-Image weights.
# Requires amdxdna and the XRT development package with pyxrt.
set -e
cd "$(dirname "$0")"
export PYTHONUTF8=1

if ! python3 -c "import sys; sys.exit(0 if sys.version_info >= (3, 10) else 1)" 2>/dev/null; then
    echo "Python 3.10 or newer is needed:  sudo apt install python3"
    exit 1
fi
if [ ! -e /dev/accel/accel0 ]; then
    echo "[setup] note: no NPU device (/dev/accel/accel0). Install the AMD amdxdna driver and XRT with pyxrt."
fi

need=""
webkit=gir1.2-webkit2-4.1
apt-cache show $webkit >/dev/null 2>&1 || webkit=gir1.2-webkit2-4.0
for p in python3-venv python3-gi $webkit wl-clipboard xclip; do
    dpkg -s "$p" >/dev/null 2>&1 || need="$need $p"
done
if [ -n "$need" ]; then
    echo "[setup] installing system packages:$need"
    sudo apt-get install -y $need
fi

install_deps() {
    .venv/bin/python -m pip install --disable-pip-version-check -r requirements.txt "triton-xdna==3.6.0.2026102404+a414baf" \
        --find-links https://github.com/amd/Triton-XDNA/releases/expanded_assets/latest-wheels \
        --find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels-no-rtti-2 \
        --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly \
        --find-links https://github.com/Xilinx/mlir-air/releases/expanded_assets/latest-air-wheels-no-rtti
}
if [ ! -f .venv/bin/python ]; then
    echo "[setup] creating Python environment in .venv ..."
    python3 -m venv --system-site-packages .venv
fi
if ! cmp -s requirements.txt .venv/requirements.done 2>/dev/null; then
    if ! install_deps; then
        echo "[setup] installation failed; rebuilding the Python environment once ..."
        rm -rf .venv
        python3 -m venv --system-site-packages .venv
        install_deps || exit 1
    fi
    cp requirements.txt .venv/requirements.done
fi

.venv/bin/python -c "import pyxrt" 2>/dev/null || \
    echo "[setup] note: pyxrt is not importable. Install XRT's Python bindings before NPU inference."

echo "[setup] Z-Image-Turbo model ..."
.venv/bin/python -m onwx_image.fetch_model
echo "[setup] NPU kernels for every picture size (about 30 minutes the first time; then a minute) ..."
.venv/bin/python -m onwx_image.prebuild || \
    echo "[setup] building the NPU kernels failed: the app builds them at each size's first picture instead."

mkdir -p ~/.local/share/applications
cat > ~/.local/share/applications/onwx-image.desktop <<DESK
[Desktop Entry]
Type=Application
Name=onwx-Image
Comment=Text-to-image on the AMD Ryzen AI NPU (Z-Image-Turbo)
Exec="$PWD/.venv/bin/python" -m onwx_image
Path=$PWD
Icon=$PWD/onwx_image/web/icon.png
Terminal=false
Categories=Graphics;
StartupWMClass=onwx-image
DESK
echo "[setup] application list: onwx-Image"
echo
echo "[start] onwx-Image is starting on the AMD Ryzen AI NPU."
nohup .venv/bin/python -m onwx_image >/dev/null 2>&1 &