#!/bin/bash
set -euo pipefail

# ==============================
# {{ project_name }} Model Provisioner
# Generated by ABI-Core scaffolding
# ==============================

# Configuration
PROJECT_DIR="{{ project_dir }}"
MODEL_NAME="{{ model_name | default('qwen2.5:3b') }}"
EMBEDDING_MODEL="{{ embedding_model | default('nomic-embed-text:v1.5') }}"
MODEL_SERVING="{{ model_serving | default('distributed') }}"

READY_WAIT_SECS="${READY_WAIT_SECS:-2}"
PULL_POLL_SECS="${PULL_POLL_SECS:-10}"
TIMEOUT_SECS="${TIMEOUT_SECS:-900}"

ts() { date '+%H:%M:%S'; }

echo "🚀 {{ project_name }} Model Provisioner"
echo "Mode: $MODEL_SERVING"
echo "LLM: $MODEL_NAME"
echo "Embeddings: $EMBEDDING_MODEL"
echo "========================================"

# Start a Docker Compose service
start_service() {
  local service="$1"
  echo "[$(ts)] 🚀 Starting service: $service..."
  if docker-compose up -d "$service" 2>/dev/null; then
    echo "[$(ts)] ✅ Service $service started"
    sleep 3  # Give service time to initialize
    return 0
  else
    echo "[$(ts)] ⚠️  Warning: Could not start $service (may already be running)"
    return 0  # Don't fail if already running
  fi
}

# Wait for Ollama to be ready
wait_ollama() {
  local host="$1"
  local service_name="$2"
  echo "[$(ts)] 🚦 Waiting for Ollama in $service_name..."
  local start=$(date +%s)
  
  until curl -fsS "$host/api/tags" >/dev/null 2>&1; do
    local now=$(date +%s)
    if [ $((now-start)) -ge "$TIMEOUT_SECS" ]; then
      echo "[$(ts)] ❌ Timeout waiting for Ollama in $service_name"
      return 1
    fi
    sleep "$READY_WAIT_SECS"
  done
  
  echo "[$(ts)] ✅ Ollama ready in $service_name"
  return 0
}

# Check if model exists
has_model() {
  local host="$1"
  local model="$2"
  curl -fsS "$host/api/tags" 2>/dev/null | grep -q "\"name\":\"$model\""
}

# Pull model
pull_model() {
  local host="$1"
  local model="$2"
  local service_name="$3"
  
  echo "[$(ts)] ⬇️  Pulling '$model' in $service_name..."
  
  # Start pull
  curl -fsS -X POST "$host/api/pull" \
    -H "Content-Type: application/json" \
    -d "{\"name\":\"$model\"}" >/dev/null 2>&1 || true
  
  # Wait for completion
  local start=$(date +%s)
  until has_model "$host" "$model"; do
    local now=$(date +%s)
    if [ $((now-start)) -ge "$TIMEOUT_SECS" ]; then
      echo "[$(ts)] ❌ Timeout pulling '$model' in $service_name"
      return 1
    fi
    echo "[$(ts)] ⏱  Downloading '$model'..."
    sleep "$PULL_POLL_SECS"
  done
  
  echo "[$(ts)] ✅ Model '$model' ready in $service_name"
  return 0
}

# Ensure model is available
ensure_model() {
  local host="$1"
  local model="$2"
  local service_name="$3"
  
  if ! wait_ollama "$host" "$service_name"; then
    return 1
  fi
  
  if has_model "$host" "$model"; then
    echo "[$(ts)] ✅ '$model' already available in $service_name"
    return 0
  else
    pull_model "$host" "$model" "$service_name"
    return $?
  fi
}

# Main provisioning logic
if [ "$MODEL_SERVING" = "centralized" ]; then
  echo ""
  echo "📦 CENTRALIZED MODE"
  echo "-------------------"
  
  # Start centralized Ollama service
  start_service "${PROJECT_DIR}-ollama"
  
  # Pull LLM and embeddings from centralized ollama
  OLLAMA_HOST="http://localhost:11434"
  
  if ! ensure_model "$OLLAMA_HOST" "$MODEL_NAME" "${PROJECT_DIR}-ollama"; then
    echo "❌ Failed to provision LLM model"
    exit 1
  fi
  
  if ! ensure_model "$OLLAMA_HOST" "$EMBEDDING_MODEL" "${PROJECT_DIR}-ollama"; then
    echo "❌ Failed to provision embedding model"
    exit 1
  fi
  
else
  echo ""
  echo "📦 DISTRIBUTED MODE"
  echo "-------------------"
  
  # Pull LLM from each agent's ollama
  {% if agents %}
  {% for agent_name, agent_config in agents.items() %}
  echo ""
  echo "Agent: {{ agent_name }}"
  
  # Start agent service (which includes Ollama via START_OLLAMA=true)
  start_service "${PROJECT_DIR}-{{ agent_name }}"
  
  AGENT_HOST="http://localhost:{{ agent_config.ollama_port | default('11434') }}"
  if ! ensure_model "$AGENT_HOST" "$MODEL_NAME" "{{ agent_name }}"; then
    echo "⚠️  Warning: Failed to provision model for {{ agent_name }}"
  fi
  {% endfor %}
  {% else %}
  echo "ℹ️  No agents found. Create agents first with: abi-core add agent <name>"
  {% endif %}
  
  # Start main Ollama service for embeddings
  start_service "${PROJECT_DIR}-ollama"
  
  # Pull embeddings from main ollama service
  echo ""
  echo "Embeddings (main ollama service):"
  OLLAMA_HOST="http://localhost:11434"
  if ! ensure_model "$OLLAMA_HOST" "$EMBEDDING_MODEL" "${PROJECT_DIR}-ollama"; then
    echo "❌ Failed to provision embedding model"
    exit 1
  fi
fi

echo ""
echo "========================================"
echo "🎉 Model provisioning completed!"
echo "========================================"
echo "LLM: $MODEL_NAME ✅"
echo "Embeddings: $EMBEDDING_MODEL ✅"
echo ""
