#!/usr/bin/env bash
# imajin-diffusion service runner
# Usage: ./run [command]

set -euo pipefail
SERVICE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SERVICE_DIR"

LOG_FILE="/tmp/imajin-diffusion.log"
PID_FILE="/tmp/imajin-diffusion.pid"
PORT=8002
APP="src.api.main:app"

_pid() {
  ss -tlnp "sport = :${PORT}" 2>/dev/null | grep -oP 'pid=\K[0-9]+' | head -1 || true
}

cmd="${1:-status}"

case "$cmd" in
  start)
    existing=$(_pid)
    if [[ -n "$existing" ]]; then
      echo "Already running (PID $existing)"
      exit 0
    fi
    nohup python3 -m uvicorn "$APP" --host 0.0.0.0 --port "$PORT" > "$LOG_FILE" 2>&1 &
    echo $! > "$PID_FILE"
    echo "Started PID $! — logs: $LOG_FILE"
    ;;
  stop)
    pid=$(_pid)
    if [[ -n "$pid" ]]; then
      kill "$pid" && echo "Stopped PID $pid" || echo "Failed to stop"
      rm -f "$PID_FILE"
    else
      echo "Not running"
    fi
    ;;
  restart)
    "$0" stop
    sleep 2
    "$0" start
    ;;
  dev)
    exec python3 -m uvicorn "$APP" --host 0.0.0.0 --port "$PORT" --reload
    ;;
  logs)
    exec tail -f "$LOG_FILE"
    ;;
  status)
    pid=$(_pid)
    if [[ -n "$pid" ]]; then
      echo "Running (PID $pid)"
      curl -s "http://localhost:${PORT}/health" && echo ""
    else
      echo "Stopped"
    fi
    ;;
  leases)
    curl -s http://localhost:8210/v1/leases | python3 -c "
import sys, json
data = json.load(sys.stdin)
for l in data['leases']:
    print(f\"  GPU {l['gpu_index']} {l['vram_mb']:>6}MB  {l['service_name']:<45} {l['lease_id'][:8]}\")
print(f\"Total: {len(data['leases'])} leases\")
"
    ;;
  release-stale)
    # Release leases whose service is no longer running
    curl -s http://localhost:8210/v1/leases | python3 -c "
import sys, json, subprocess, re
data = json.load(sys.stdin)
for l in data['leases']:
    lease_id = l['lease_id']
    service = l['service_name']
    print(f'Releasing {service} ({lease_id[:8]})...')
    import urllib.request
    req = urllib.request.Request(
        f'http://localhost:8210/v1/leases/{lease_id}/release',
        method='POST'
    )
    with urllib.request.urlopen(req) as r:
        print(' ', r.read().decode())
" 2>/dev/null || echo "No leases to release"
    ;;
  *)
    echo "Usage: ./run [command]"
    echo "  start          Start service in background (port $PORT)"
    echo "  stop           Stop background service"
    echo "  restart        Restart background service"
    echo "  dev            Start with --reload (hot-reload, foreground)"
    echo "  logs           Tail service logs"
    echo "  status         Show running status + health"
    echo "  leases         Show GPU leases from model-boss"
    echo "  release-stale  Release all GPU leases (use when leases are stuck)"
    exit 1
    ;;
esac
