#!/usr/bin/env bash
# onw AMD - Ore no NPU ga konna ni ugoku wake nai: install the resident top-bar app (Ubuntu, AMD Ryzen AI NPU).
# The NPU engine is FastFlowLM; the amdxdna driver must be loaded.
# Installs the system pieces the indicator and the window need (AppIndicator, WebKit; asks for sudo once), creates .venv,
# installs onw + the tray, adds "onw" to the application list and starts it in the top bar.
# GNOME shows indicators through the "Ubuntu AppIndicators" extension, which Ubuntu enables by default.
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 (Ubuntu 22.04 and later ship it):  sudo apt install python3"
    exit 1
fi

need=""
# the window is a WebKit web view (pywebview): 4.1 on Ubuntu 24.04 and later, 4.0 before
webkit=gir1.2-webkit2-4.1
apt-cache show $webkit >/dev/null 2>&1 || webkit=gir1.2-webkit2-4.0
# wl-clipboard / xclip: "Copy the OpenAI API URL" in the top-bar menu (Wayland / X11)
for p in python3-venv python3-gi gir1.2-ayatanaappindicator3-0.1 $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

# FastFlowLM needs a high memlock limit on Linux (its own guide: unlimited, then reboot); Ubuntu's default fails
# `flm validate`. limits.d covers PAM logins, the systemd drop-ins the desktop session and its apps.
relogin=""
if [ "$(ulimit -l)" != "unlimited" ] && [ ! -f /etc/security/limits.d/99-onw-amd-memlock.conf ]; then
    echo "[setup] raising the memlock limit for the NPU (FastFlowLM needs it; asks for sudo)"
    printf '*    soft    memlock    unlimited\n*    hard    memlock    unlimited\n' |
        sudo tee /etc/security/limits.d/99-onw-amd-memlock.conf >/dev/null
    for d in system user; do
        sudo mkdir -p /etc/systemd/$d.conf.d
        printf '[Manager]\nDefaultLimitMEMLOCK=infinity\n' | sudo tee /etc/systemd/$d.conf.d/99-onw-amd-memlock.conf >/dev/null
    done
    relogin=1
fi

install_onw() {
    # the bundled pip is fine (a newer pip once failed to install packages that have scripts, so it is not upgraded)
    .venv/bin/python -m pip install -e ".[tray]"
}
if [ ! -f .venv/installed-tray.txt ] || ! .venv/bin/python -c "import webview" >/dev/null 2>&1; then
    if [ ! -f .venv/bin/python ]; then
        echo "[setup] creating Python environment in .venv ..."
        # system site packages: the GObject bindings for the indicator and WebKit come from apt (python3-gi)
        python3 -m venv --system-site-packages .venv
    fi
    if ! install_onw; then
        echo "[setup] installation failed; rebuilding the Python environment once ..."
        rm -rf .venv
        python3 -m venv --system-site-packages .venv
        install_onw || exit 1
    fi
    echo ok > .venv/installed.txt
    echo ok > .venv/installed-tray.txt
fi

mkdir -p ~/.local/share/applications
cat > ~/.local/share/applications/onw-amd.desktop <<EOF
[Desktop Entry]
Type=Application
Name=onw AMD
Comment=LLMs on the AMD Ryzen AI NPU (俺のNPUがこんなに動くわけない)
Exec="$PWD/.venv/bin/onw" tray
Path=$PWD
Icon=$PWD/onw/web/icon.png
Terminal=false
Categories=Utility;
EOF
echo "[setup] application list: onw AMD"

if [ -n "$relogin" ]; then
    echo
    echo "[setup] The memlock limit was raised. Restart the PC once so that the NPU can be used,"
    echo "        then start onw AMD from the application list."
    exit 0
fi

echo "[start] onw AMD is starting in the top bar ..."
# --replace: a tray already running (an update by the one-line installer) restarts on the new files
nohup .venv/bin/onw tray --replace >/dev/null 2>&1 &