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

MODEL_NAME="${MODEL_NAME:-qwen2.5-coder:3b}"
TRIAGE_MODEL="${TRIAGE_MODEL:-$MODEL_NAME}"
ACTIONS_MODEL="${ACTIONS_MODEL:-$MODEL_NAME}"
WRITER_MODEL="${WRITER_MODEL:-$MODEL_NAME}"

OLLAMA_URL="${OLLAMA_URL:-http://127.0.0.1:11434}"
export OLLAMA_URL MODEL_NAME TRIAGE_MODEL ACTIONS_MODEL WRITER_MODEL

echo "[start.sh] Starting Ollama..."
ollama serve > /tmp/ollama.log 2>&1 &
OLLAMA_PID=$!

# Wait until Ollama is ready
python3 - << 'PY'
import os, time, urllib.request
base = os.environ.get("OLLAMA_URL","http://127.0.0.1:11434").rstrip("/")
url = base + "/api/tags"
for _ in range(180):
    try:
        urllib.request.urlopen(url, timeout=1).read()
        print("Ollama is up")
        raise SystemExit(0)
    except Exception:
        time.sleep(1)
raise SystemExit("Ollama did not start")
PY

# Unique model list
MODELS=("$MODEL_NAME" "$TRIAGE_MODEL" "$ACTIONS_MODEL" "$WRITER_MODEL")
UNIQ=()
for m in "${MODELS[@]}"; do
  seen=false
  for u in "${UNIQ[@]}"; do
    if [ "$u" = "$m" ]; then seen=true; fi
  done
  if [ "$seen" = false ]; then UNIQ+=("$m"); fi
done

echo "[start.sh] Ensuring models exist..."
for m in "${UNIQ[@]}"; do
  echo "[start.sh] model: $m"
  if ! ollama show "$m" >/dev/null 2>&1; then
    echo "[start.sh] pulling $m..."
    ollama pull "$m"
  else
    echo "[start.sh] already present"
  fi
done

echo "[start.sh] Warming up models..."
python3 - << 'PY'
import os, json, urllib.request
base = os.environ.get("OLLAMA_URL","http://127.0.0.1:11434").rstrip("/")
ctx = int(os.environ.get("OLLAMA_CONTEXT_LENGTH","4096"))

models = []
for k in ("MODEL_NAME","TRIAGE_MODEL","ACTIONS_MODEL","WRITER_MODEL"):
    v = os.environ.get(k)
    if v and v not in models:
        models.append(v)

for m in models:
    payload = {
      "model": m,
      "prompt": "ping",
      "stream": False,
      "options": {"num_ctx": ctx, "num_predict": 16}
    }
    req = urllib.request.Request(
        base + "/api/generate",
        data=json.dumps(payload).encode(),
        headers={"Content-Type":"application/json"},
    )
    urllib.request.urlopen(req, timeout=1800).read()
    print("Warmup OK:", m)
PY

echo "[start.sh] Launching app..."
exec python3 /app/app.py