#!/usr/bin/env bash
# One-click setup for GenLovers' local uncensored AI companion tutorial.
#
# Installs Ollama if it isn't already present, pulls the abliterated Qwen3.5
# model, and confirms it responds. Everything runs on this machine only —
# nothing is uploaded anywhere.
#
# Usage: bash setup-local-uncensored-companion.sh
# (or just double-click it if your file manager runs .sh files on open)

set -euo pipefail

OLLAMA_URL="${OLLAMA_URL:-http://localhost:11434}"
MODEL_TAG="huihui-qwen3.5-0.8b-abliterated"
HF_REPO="mradermacher/Huihui-Qwen3.5-0.8B-abliterated-GGUF"
HF_QUANT="Q4_K_M"
HF_GGUF_FILE="Huihui-Qwen3.5-0.8B-abliterated.${HF_QUANT}.gguf"

info()  { printf '\n\033[1;36m==>\033[0m %s\n' "$1"; }
ok()    { printf '\033[1;32m  ok:\033[0m %s\n' "$1"; }
warn()  { printf '\033[1;33m  warn:\033[0m %s\n' "$1"; }
fail()  { printf '\033[1;31m  error:\033[0m %s\n' "$1"; exit 1; }

# ---------------------------------------------------------------------------
# 1. Install Ollama if missing
# ---------------------------------------------------------------------------
info "Checking for Ollama..."
if command -v ollama >/dev/null 2>&1; then
  ok "Ollama already installed ($(ollama --version 2>/dev/null | head -n1))."
else
  info "Ollama not found. Installing (this needs your password on some systems)..."
  if [ "$(uname -s)" = "Darwin" ] && command -v brew >/dev/null 2>&1; then
    brew install ollama || curl -fsSL https://ollama.com/install.sh | sh
  else
    curl -fsSL https://ollama.com/install.sh | sh
  fi
  command -v ollama >/dev/null 2>&1 || fail "Ollama install did not complete. Install it manually from https://ollama.com/download, then re-run this script."
  ok "Ollama installed."
fi

# ---------------------------------------------------------------------------
# 2. Make sure the Ollama server is running
# ---------------------------------------------------------------------------
info "Making sure the Ollama server is running..."
if curl -fsS "${OLLAMA_URL}/api/version" >/dev/null 2>&1; then
  ok "Ollama server already running at ${OLLAMA_URL}."
else
  if [ "$(uname -s)" = "Darwin" ]; then
    open -a Ollama 2>/dev/null || true
  fi
  if command -v systemctl >/dev/null 2>&1 && systemctl list-unit-files 2>/dev/null | grep -q '^ollama\.service'; then
    sudo systemctl start ollama 2>/dev/null || true
  fi
  if ! curl -fsS "${OLLAMA_URL}/api/version" >/dev/null 2>&1; then
    nohup ollama serve >/tmp/ollama-serve.log 2>&1 &
  fi
  for i in $(seq 1 15); do
    curl -fsS "${OLLAMA_URL}/api/version" >/dev/null 2>&1 && break
    sleep 1
  done
  curl -fsS "${OLLAMA_URL}/api/version" >/dev/null 2>&1 || fail "Ollama server did not start. Start it manually (ollama serve) and re-run this script."
  ok "Ollama server is up."
fi

# ---------------------------------------------------------------------------
# 3. Pull the abliterated model
# ---------------------------------------------------------------------------
info "Checking for the model (${MODEL_TAG})..."
if ollama list 2>/dev/null | awk '{print $1}' | grep -qx "${MODEL_TAG}"; then
  ok "Model already pulled."
else
  info "Pulling ${MODEL_TAG} from Hugging Face (hf.co/${HF_REPO}:${HF_QUANT}, ~528 MB). This can take a few minutes."
  if ollama pull "hf.co/${HF_REPO}:${HF_QUANT}"; then
    # Ollama's hf.co puller tags the model as the source path, not our short
    # name. Alias it so the rest of this script and the companion chat page
    # can use one stable tag regardless of the pull method.
    PULLED_TAG=$(ollama list 2>/dev/null | awk 'NR>1 {print $1}' | grep -i "huihui-qwen3.5-0.8b-abliterated" | head -n1)
    if [ -n "${PULLED_TAG:-}" ] && [ "${PULLED_TAG}" != "${MODEL_TAG}" ]; then
      ollama cp "${PULLED_TAG}" "${MODEL_TAG}" || true
    fi
    ok "Model pulled."
  else
    warn "Direct pull failed (this Hugging Face endpoint is known to be flaky on the manifest fetch). Falling back to a direct download."
    TMP_DIR=$(mktemp -d)
    trap 'rm -rf "${TMP_DIR}"' EXIT
    GGUF_URL="https://huggingface.co/${HF_REPO}/resolve/main/${HF_GGUF_FILE}"
    info "Downloading ${HF_GGUF_FILE} directly..."
    curl -fL --progress-bar -o "${TMP_DIR}/model.gguf" "${GGUF_URL}" || fail "Direct download also failed. Check your internet connection and try again."
    printf 'FROM %s/model.gguf\n' "${TMP_DIR}" > "${TMP_DIR}/Modelfile"
    ollama create "${MODEL_TAG}" -f "${TMP_DIR}/Modelfile" || fail "ollama create failed on the downloaded file."
    ok "Model built from the direct download."
  fi
fi

# ---------------------------------------------------------------------------
# 4. Warm-up / verification prompt
# ---------------------------------------------------------------------------
info "Sending a test prompt to confirm everything works..."
RESPONSE=$(curl -fsS "${OLLAMA_URL}/api/generate" \
  -d "{\"model\": \"${MODEL_TAG}\", \"prompt\": \"Say hello in one short sentence.\", \"think\": false, \"stream\": false}" \
  2>/dev/null) || fail "The test prompt request failed. Is the Ollama server still running?"

REPLY=$(printf '%s' "${RESPONSE}" | python3 -c 'import json,sys; print(json.load(sys.stdin).get("response","").strip())' 2>/dev/null || true)
if [ -n "${REPLY}" ]; then
  ok "Model responded: \"${REPLY}\""
else
  warn "Got a response but couldn't parse it. Raw output: ${RESPONSE}"
fi

echo ""
echo "Setup complete. Model tag: ${MODEL_TAG}"
echo "Next: open the companion chat page (local-companion-chat.html) and enter this model tag on its setup screen."