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

PRIMARY_BASE="${PRIMARY_BASE:-qwen3-coder:30b}"
PRIMARY_NAME="${PRIMARY_NAME:-junie-qwen3-coder-30b-a6000}"

ALT_BASE="${ALT_BASE:-qwen2.5-coder:32b}"
ALT_NAME="${ALT_NAME:-junie-qwen25-coder-32b-a6000}"

FAST_BASE="${FAST_BASE:-qwen2.5-coder:7b}"
FAST_NAME="${FAST_NAME:-junie-qwen25-coder-7b-fast}"

PRIMARY_CTX="${PRIMARY_CTX:-65536}"
ALT_CTX="${ALT_CTX:-32768}"
FAST_CTX="${FAST_CTX:-16384}"

WORKDIR="${WORKDIR:-$HOME/ollama-junie-a6000}"
JUNIE_HOME="${JUNIE_HOME:-$HOME/.junie}"
JUNIE_PROFILE="${JUNIE_PROFILE:-local-ollama-a6000}"

mkdir -p "$WORKDIR" "$JUNIE_HOME/models"

echo "==> Checking Ollama"
if ! command -v ollama >/dev/null 2>&1; then
  echo "==> Installing Ollama"
  curl -fsSL https://ollama.com/install.sh | sh
fi

echo "==> Configuring Ollama systemd for A6000 long-context use"
if systemctl list-unit-files | grep -q '^ollama\.service'; then
  sudo mkdir -p /etc/systemd/system/ollama.service.d
  sudo tee /etc/systemd/system/ollama.service.d/a6000.conf >/dev/null <<EOF
[Service]
Environment="OLLAMA_HOST=127.0.0.1:11434"
Environment="OLLAMA_FLASH_ATTENTION=1"
Environment="OLLAMA_KV_CACHE_TYPE=q8_0"
Environment="OLLAMA_KEEP_ALIVE=24h"
Environment="OLLAMA_NUM_PARALLEL=1"
Environment="OLLAMA_MAX_LOADED_MODELS=1"
EOF

  sudo systemctl daemon-reload
  sudo systemctl enable ollama >/dev/null
  sudo systemctl restart ollama
else
  echo "WARN: ollama.service niet gevonden. Start Ollama handmatig met:"
  echo 'OLLAMA_HOST=127.0.0.1:11434 OLLAMA_FLASH_ATTENTION=1 OLLAMA_KV_CACHE_TYPE=q8_0 OLLAMA_KEEP_ALIVE=24h ollama serve'
fi

echo "==> Pulling base models"
ollama pull "$PRIMARY_BASE"
ollama pull "$ALT_BASE"
ollama pull "$FAST_BASE"

cat > "$WORKDIR/Modelfile.${PRIMARY_NAME}" <<EOF
FROM ${PRIMARY_BASE}

PARAMETER num_ctx ${PRIMARY_CTX}
PARAMETER num_predict 4096
PARAMETER temperature 0.3
PARAMETER top_p 0.9
PARAMETER repeat_penalty 1.05

SYSTEM """
You are a senior software-engineering coding agent inside a JetBrains/Junie workflow.

Rules:
- Be precise, deterministic, and practical.
- Prefer small, safe, reviewable changes.
- Read surrounding code before proposing edits.
- Preserve existing architecture and style.
- For code changes, explain intent briefly and produce clean patches.
- Do not invent files, APIs, commands, or project structure.
- When uncertain, inspect first instead of guessing.
- Keep output concise unless a deeper design answer is needed.
"""
EOF

cat > "$WORKDIR/Modelfile.${ALT_NAME}" <<EOF
FROM ${ALT_BASE}

PARAMETER num_ctx ${ALT_CTX}
PARAMETER num_predict 4096
PARAMETER temperature 0.25
PARAMETER top_p 0.9
PARAMETER repeat_penalty 1.05

SYSTEM """
You are a senior coding assistant optimized for Java, Kotlin, Python, TypeScript, SQL, Linux, Docker, and architecture work.

Rules:
- Give direct, technically correct answers.
- Prefer minimal diffs and commands that can be copied.
- Do not hallucinate APIs or project files.
- Ask only when blocked; otherwise make a reasonable local change plan.
"""
EOF

cat > "$WORKDIR/Modelfile.${FAST_NAME}" <<EOF
FROM ${FAST_BASE}

PARAMETER num_ctx ${FAST_CTX}
PARAMETER num_predict 2048
PARAMETER temperature 0.2
PARAMETER top_p 0.9

SYSTEM """
You are a fast helper model for summarizing, routing, and classifying coding context.
Be concise and preserve exact technical details.
"""
EOF

echo "==> Creating tuned Ollama models"
ollama create "$PRIMARY_NAME" -f "$WORKDIR/Modelfile.${PRIMARY_NAME}"
ollama create "$ALT_NAME" -f "$WORKDIR/Modelfile.${ALT_NAME}"
ollama create "$FAST_NAME" -f "$WORKDIR/Modelfile.${FAST_NAME}"

echo "==> Creating Junie CLI custom model profile"
cat > "$JUNIE_HOME/models/${JUNIE_PROFILE}.json" <<EOF
{
  "id": "${PRIMARY_NAME}",
  "baseUrl": "http://localhost:11434/v1/chat/completions",
  "apiType": "OpenAICompletion",
  "temperature": 0.3,
  "extraHeaders": {
    "X-Custom-Source": "Junie-Ollama-A6000"
  },
  "primaryModel": {
    "id": "${PRIMARY_NAME}",
    "temperature": 0.3
  },
  "fasterModel": {
    "id": "${FAST_NAME}",
    "temperature": 0.2
  }
}
EOF

echo "==> Setting Junie CLI default model profile"
if command -v python3 >/dev/null 2>&1; then
  CONFIG="$JUNIE_HOME/config.json" python3 - <<'PY'
import json, os, pathlib, shutil

path = pathlib.Path(os.environ["CONFIG"])
profile = "local-ollama-a6000"

if path.exists():
    backup = path.with_suffix(".json.bak")
    shutil.copy2(path, backup)
    try:
        data = json.loads(path.read_text())
    except Exception:
        data = {}
else:
    data = {}

data["model"] = profile
data["model-default-locations"] = True

path.write_text(json.dumps(data, indent=2) + "\n")
PY
else
  echo "WARN: python3 ontbreekt; zet handmatig in ~/.junie/config.json:"
  echo '{ "model": "local-ollama-a6000", "model-default-locations": true }'
fi

echo "==> Preloading primary model"
curl -s http://127.0.0.1:11434/api/generate \
  -H "Content-Type: application/json" \
  -d "{\"model\":\"${PRIMARY_NAME}\",\"prompt\":\"\",\"keep_alive\":-1}" >/dev/null || true

echo "==> Testing OpenAI-compatible endpoint"
curl -s http://127.0.0.1:11434/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d "{
    \"model\": \"${PRIMARY_NAME}\",
    \"messages\": [{\"role\":\"user\",\"content\":\"Return exactly: OK\"}],
    \"stream\": false
  }" | sed 's/,/,\n/g' | head -40

echo
echo "Done."
echo
echo "Models:"
ollama list | grep -E "${PRIMARY_NAME}|${ALT_NAME}|${FAST_NAME}" || true
echo
echo "Check GPU load with:"
echo "  ollama ps"
echo "  watch -n 1 nvidia-smi"
echo
echo "Junie CLI:"
echo "  cd /path/to/project"
echo "  junie --model ${JUNIE_PROFILE}"
echo
echo "IntelliJ AI Assistant:"
echo "  Settings > Tools > AI Assistant > Providers & API keys > Ollama"
echo "  Host: http://localhost:11434"
echo "  Model: ${PRIMARY_NAME}"