#!/bin/bash
# MyAI privileged helper - called via pkexec by the Flask backend.
# All output (stdout + stderr) is streamed to the UI via NDJSON.
#
# We run ollama under a systemd unit (myai-ollama.service) that we ship
# ourselves. That unit needs `User=ollama` to resolve, which means the
# `ollama` system user must exist. Ollama's install.sh tries to create
# it with `useradd -U` which fails if the `ollama` group already exists
# from a previous attempt (and crucially: useradd then does NOT create
# the user). We pre-create the user ourselves to bypass that footgun.

export LANG=C.UTF-8
export LC_ALL=C.UTF-8
export PATH=/usr/sbin:/usr/bin:/sbin:/bin

OLLAMA_API=http://127.0.0.1:11434/api/tags
MYAI_UNIT=myai-ollama.service
OLLAMA_HOME=/usr/share/ollama

action="${1:-}"

api_responds() {
    curl -fsS -o /dev/null --max-time 2 "$OLLAMA_API" 2>/dev/null
}

wait_for_api() {
    local max=${1:-15}
    local i=0
    while [ "$i" -lt "$max" ]; do
        api_responds && return 0
        sleep 1
        i=$((i + 1))
    done
    return 1
}

ensure_ollama_user() {
    # Idempotent. Handles three states:
    #   (a) user + group both exist -> nothing to do
    #   (b) group exists, user missing -> create user with that gid
    #   (c) neither exists -> create both
    if id -u ollama >/dev/null 2>&1; then
        return 0
    fi
    echo ">>> Creating ollama system user"
    if ! getent group ollama >/dev/null 2>&1; then
        groupadd --system ollama || { echo "!!! groupadd failed"; return 1; }
    fi
    useradd --system \
            --shell /usr/sbin/nologin \
            --home-dir "$OLLAMA_HOME" \
            --no-create-home \
            --gid ollama \
            ollama \
        || { echo "!!! useradd failed"; return 1; }
    mkdir -p "$OLLAMA_HOME"
    chown ollama:ollama "$OLLAMA_HOME"
    chmod 755 "$OLLAMA_HOME"
    echo ">>> ollama user created"
}

stop_conflicting() {
    # If upstream's ollama.service exists, stop+disable it so it doesn't
    # fight ours for port 11434.
    if systemctl list-unit-files ollama.service >/dev/null 2>&1; then
        systemctl disable --now ollama.service >/dev/null 2>&1 || true
    fi
    # And any stray ollama processes that install.sh might have spawned.
    pkill -x ollama 2>/dev/null
    sleep 1
}

bring_up_myai_unit() {
    systemctl daemon-reload >/dev/null 2>&1 || true
    # Reset the failed/restart-loop state from prior attempts so systemd
    # actually tries again instead of "auto-restart"ing into the cached
    # exit code.
    systemctl reset-failed "$MYAI_UNIT" 2>/dev/null || true
    echo ">>> Enabling and starting $MYAI_UNIT"
    if ! systemctl enable --now "$MYAI_UNIT" 2>&1; then
        echo "!!! systemctl enable --now failed"
        systemctl --no-pager --lines=20 status "$MYAI_UNIT" 2>&1 || true
        return 1
    fi
    if wait_for_api 15; then
        echo ">>> Ollama API is responding on port 11434"
        return 0
    fi
    echo "!!! Ollama API did not respond"
    systemctl --no-pager --lines=20 status "$MYAI_UNIT" 2>&1 || true
    return 1
}

install_ollama() {
    echo ">>> Preparing ollama system user"
    ensure_ollama_user || return 1

    echo ">>> Running Ollama installer"
    # install.sh's useradd will see our user already exists and skip the
    # broken path. Other steps (download, group setup, binary install)
    # proceed normally.
    curl -fsSL https://ollama.com/install.sh | sh

    echo ">>> Verifying installation"
    if ! command -v ollama >/dev/null 2>&1; then
        echo "!!! ollama binary not found after install"
        return 1
    fi
    if ! id ollama >/dev/null 2>&1; then
        echo "!!! ollama user disappeared (this shouldn't happen)"
        return 1
    fi

    stop_conflicting
    bring_up_myai_unit
}

start_ollama() {
    # Defensive: if the user is missing on a fresh start (e.g. user wiped
    # /etc/passwd somehow, or Ollama purged), put it back.
    ensure_ollama_user || return 1

    if ! command -v ollama >/dev/null 2>&1; then
        echo "!!! ollama binary not installed — run install first"
        return 1
    fi

    stop_conflicting
    bring_up_myai_unit
}

case "$action" in
    install-ollama)
        install_ollama
        ;;
    start-ollama)
        start_ollama
        ;;
    *)
        echo "Unknown action: $action" >&2
        exit 1
        ;;
esac
