#!/bin/bash
# autopkgtest for lemond (lemonade-server)
#
# Starts lemond on a random free port, verifies HTTP endpoints,
# then shuts the daemon down.  No inference backends or GPU hardware
# are required.
#
# Exit codes: 0 = all checks passed, non-zero = one or more failed.

set -euo pipefail

###############################################################################
# Pre-flight checks
###############################################################################

echo "=== autopkgtest: lemond ==="
echo "Date: $(date -u)"

# Verify python3-requests is available
if ! python3 -c "import requests" 2>/dev/null; then
    echo "FAIL: python3-requests is not installed"
    echo "Install with: apt-get install -y python3-requests"
    exit 1
fi

# Find lemond binary
LEMOND=""
for candidate in \
    "$(dpkg -L lemonade-server 2>/dev/null | grep /usr/bin/lemond)" \
    "$(which lemond 2>/dev/null)" \
    "/usr/bin/lemond"; do
    if [ -n "$candidate" ] && [ -x "$candidate" ]; then
        LEMOND="$candidate"
        break
    fi
done

if [ -z "$LEMOND" ]; then
    echo "FAIL: lemond binary not found"
    echo "Is lemonade-server installed?"
    echo "Installed files:"
    dpkg -L lemonade-server 2>/dev/null || true
    echo "PATH: $PATH"
    exit 1
fi

echo "lemond found at: $LEMOND"

###############################################################################
# Start lemond on a free port
###############################################################################

# Allocate a random free port
PORT=""
for _attempt in $(seq 1 5); do
    PORT=$(python3 -c "
import socket
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.bind(('127.0.0.1', 0))
print(s.getsockname()[1])
s.close()
")
    # Verify the port is actually free
    if ! python3 -c "
import socket
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
    s.bind(('127.0.0.1', $PORT))
except OSError:
    import sys; sys.exit(1)
s.close()
" 2>/dev/null; then
        continue
    fi
    break
done

if [ -z "$PORT" ]; then
    echo "FAIL: could not allocate a free port"
    exit 1
fi

echo "Using port: $PORT"

# Create a temporary cache directory
CACHE_DIR=$(mktemp -d lemonade-autopkgtest-XXXXXX)
LOG_FILE="$CACHE_DIR/lemond.log"

echo "Cache directory: $CACHE_DIR"

# Start lemond
python3 -c "
import subprocess, sys, os
proc = subprocess.Popen(
    ['${LEMOND}', '${CACHE_DIR}', '--port', '${PORT}'],
    stdout=open('${LOG_FILE}', 'w'),
    stderr=subprocess.STDOUT,
    env=os.environ.copy(),
)
print(proc.pid)
" > "$CACHE_DIR/pid"

LEMOND_PID=$(cat "$CACHE_DIR/pid")
echo "lemond PID: $LEMOND_PID"

###############################################################################
# Cleanup function
###############################################################################

cleanup() {
    echo ""
    echo "=== Cleaning up ==="
    # Kill lemond if still running
    if kill -0 "$LEMOND_PID" 2>/dev/null; then
        echo "Stopping lemond (PID $LEMOND_PID)..."
        kill "$LEMOND_PID" 2>/dev/null || true
        # Wait up to 10 seconds
        for _i in $(seq 1 20); do
            kill -0 "$LEMOND_PID" 2>/dev/null || break
            sleep 0.5
        done
        # Force kill if still running
        if kill -0 "$LEMOND_PID" 2>/dev/null; then
            echo "Force killing lemond..."
            kill -9 "$LEMOND_PID" 2>/dev/null || true
        fi
    fi
    rm -rf "$CACHE_DIR"
    echo "Cleanup done."
}

trap cleanup EXIT

###############################################################################
# Wait for lemond to become healthy
###############################################################################

echo ""
echo "=== Waiting for lemond to become healthy ==="

HEALTHY=false
for _i in $(seq 1 20); do
    if ! kill -0 "$LEMOND_PID" 2>/dev/null; then
        echo "FAIL: lemond exited before becoming healthy"
        echo "Last 30 lines of log:"
        tail -30 "$LOG_FILE" 2>/dev/null || true
        exit 1
    fi

    HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" \
        --max-time 1 "http://127.0.0.1:${PORT}/api/v1/health" 2>/dev/null || echo "000")

    if [ "$HTTP_CODE" = "200" ]; then
        HEALTHY=true
        break
    fi
    sleep 0.5
done

if [ "$HEALTHY" = false ]; then
    echo "FAIL: lemond did not become healthy within 10 seconds"
    echo "Last 50 lines of log:"
    tail -50 "$LOG_FILE" 2>/dev/null || true
    exit 1
fi

echo "lemond is healthy."

###############################################################################
# Run tests
###############################################################################

echo ""
echo "=== Running endpoint tests ==="
echo ""

PASS=0
FAIL=0
TOTAL=0

check() {
    local desc="$1"
    local cond="$2"
    TOTAL=$((TOTAL + 1))
    if eval "$cond"; then
        PASS=$((PASS + 1))
        echo "  PASS [$TOTAL] $desc"
    else
        FAIL=$((FAIL + 1))
        echo "  FAIL [$TOTAL] $desc"
    fi
}

BASE="http://127.0.0.1:${PORT}"

# 1. Endpoint registration (all four prefix families)
echo "--- Endpoint registration ---"
for ep in \
    "chat/completions" "completions" "embeddings" "models" \
    "models/check-updates" "responses" "pull" "registry/search" \
    "pull/variants" "routing/validate" "delete" "load" \
    "unload" "health" "stats" "system-info" "rerank" \
    "reranking" "reranker" "audio/transcriptions" \
    "images/generations" "install" "uninstall"; do
    for ver in v0 v1; do
        CODE=$(curl -s -o /dev/null -w "%{http_code}" \
            --max-time 1 -X HEAD "${BASE}/api/${ver}/${ep}" 2>/dev/null || true)
        # Ensure CODE is just the HTTP status
        CODE="${CODE%%[^0-9]*}"
        [ -z "$CODE" ] && CODE="000"
        check "HEAD /api/${ver}/${ep} != 404 (got $CODE)" "[ \"$CODE\" != \"404\" ]"
    done
done

# POST-only routes
for ep in "models/register"; do
    for ver in v0 v1; do
        CODE=$(curl -s -o /dev/null -w "%{http_code}" \
            --max-time 1 -X POST -H "Content-Type: application/json" \
            -d '{}' "${BASE}/api/${ver}/${ep}" 2>/dev/null || true)
        CODE="${CODE%%[^0-9]*}"
        [ -z "$CODE" ] && CODE="000"
        check "POST /api/${ver}/${ep} != 404 (got $CODE)" "[ \"$CODE\" != \"404\" ]"
    done
done

# 2. Unknown route returns 404
echo ""
echo "--- 404 handling ---"
CODE=$(curl -s -o /dev/null -w "%{http_code}" \
    --max-time 1 "${BASE}/api/v1/nonexistent-route" 2>/dev/null || true)
CODE="${CODE%%[^0-9]*}"
[ -z "$CODE" ] && CODE="000"
check "GET /api/v1/nonexistent-route returns 404" "[ \"$CODE\" = \"404\" ]"

# 3. Retired endpoints are gone
echo ""
echo "--- Retired endpoints ---"
for ep in "params" "log-level" "test"; do
    for prefix in "/api/v0" "/api/v1" "/v0" "/v1"; do
        CODE=$(curl -s -o /dev/null -w "%{http_code}" \
            --max-time 1 "${BASE}${prefix}/${ep}" 2>/dev/null || true)
        CODE="${CODE%%[^0-9]*}"
        [ -z "$CODE" ] && CODE="000"
        check "404 for retired ${prefix}/${ep} (got $CODE)" "[ \"$CODE\" = \"404\" ]"
    done
done
CODE=$(curl -s -o /dev/null -w "%{http_code}" \
    --max-time 1 "${BASE}/status" 2>/dev/null || true)
CODE="${CODE%%[^0-9]*}"
[ -z "$CODE" ] && CODE="000"
check "404 for retired /status (got $CODE)" "[ \"$CODE\" = \"404\" ]"

# 4. Legacy status page
echo ""
echo "--- Legacy status page ---"
CODE=$(curl -s -o /dev/null -w "%{http_code}" \
    --max-time 1 "${BASE}/api/v1" 2>/dev/null || true)
CODE="${CODE%%[^0-9]*}"
[ -z "$CODE" ] && CODE="000"
check "GET /api/v1 returns 200" "[ \"$CODE\" = \"200\" ]"

# 5. /live endpoint
echo ""
echo "--- /live endpoint ---"
CODE=$(curl -s -o /dev/null -w "%{http_code}" \
    --max-time 1 "${BASE}/live" 2>/dev/null || true)
CODE="${CODE%%[^0-9]*}"
[ -z "$CODE" ] && CODE="000"
check "GET /live returns 200" "[ \"$CODE\" = \"200\" ]"

# 6. /health endpoint
echo ""
echo "--- /health endpoint ---"
CODE=$(curl -s -o /dev/null -w "%{http_code}" \
    --max-time 1 "${BASE}/api/v1/health" 2>/dev/null || true)
CODE="${CODE%%[^0-9]*}"
[ -z "$CODE" ] && CODE="000"
check "GET /api/v1/health returns 200" "[ \"$CODE\" = \"200\" ]"

if [ "$CODE" = "200" ]; then
    HEALTH_BODY=$(curl -s --max-time 1 "${BASE}/api/v1/health" 2>/dev/null)
    check "/health has status=ok" "$(echo "$HEALTH_BODY" | python3 -c "import sys,json; d=json.load(sys.stdin); sys.exit(0 if d.get('status')=='ok' else 1)")"
    check "/health has all_models_loaded" "$(echo "$HEALTH_BODY" | python3 -c "import sys,json; d=json.load(sys.stdin); sys.exit(0 if 'all_models_loaded' in d else 1)")"
    check "/health has max_models" "$(echo "$HEALTH_BODY" | python3 -c "import sys,json; d=json.load(sys.stdin); sys.exit(0 if 'max_models' in d else 1)")"
    if echo "$HEALTH_BODY" | python3 -c "import sys,json; d=json.load(sys.stdin); m=d.get('max_models',{}); sys.exit(0 if all(k in m for k in ['llm','embedding','reranking']) else 1)"; then
        check "/health max_models has llm, embedding, reranking" "true"
    else
        check "/health max_models has llm, embedding, reranking" "false"
    fi
fi

# 7. /metrics endpoint
echo ""
echo "--- /metrics endpoint ---"
CODE=$(curl -s -o /dev/null -w "%{http_code}" \
    --max-time 1 "${BASE}/metrics" 2>/dev/null || true)
CODE="${CODE%%[^0-9]*}"
[ -z "$CODE" ] && CODE="000"
check "GET /metrics returns 200" "[ \"$CODE\" = \"200\" ]"
if [ "$CODE" = "200" ]; then
    METRICS_BODY=$(curl -s --max-time 1 "${BASE}/metrics" 2>/dev/null)
    # Check content-type header using GET (HEAD may not return it)
    CT=$(curl -s -D - --max-time 1 "${BASE}/metrics" 2>/dev/null | grep -i "^Content-Type:" | head -1)
    if echo "$CT" | grep -qi 'text/plain'; then
        check "/metrics content-type includes text/plain" "true"
    else
        check "/metrics content-type includes text/plain" "false"
    fi
    # Check for lemonade_server_up metric
    if echo "$METRICS_BODY" | grep -q 'lemonade_server_up'; then
        check "/metrics has lemonade_server_up" "true"
    else
        check "/metrics has lemonade_server_up" "false"
    fi
fi

# 8. /models endpoint
echo ""
echo "--- /models endpoint ---"
CODE=$(curl -s -o /dev/null -w "%{http_code}" \
    --max-time 1 "${BASE}/api/v1/models" 2>/dev/null || true)
CODE="${CODE%%[^0-9]*}"
[ -z "$CODE" ] && CODE="000"
check "GET /api/v1/models returns 200" "[ \"$CODE\" = \"200\" ]"

# 9. /stats endpoint
echo ""
echo "--- /stats endpoint ---"
CODE=$(curl -s -o /dev/null -w "%{http_code}" \
    --max-time 1 "${BASE}/api/v1/stats" 2>/dev/null || true)
CODE="${CODE%%[^0-9]*}"
[ -z "$CODE" ] && CODE="000"
check "GET /api/v1/stats returns 200" "[ \"$CODE\" = \"200\" ]"

# 10. /system-info endpoint
echo ""
echo "--- /system-info endpoint ---"
CODE=$(curl -s -o /dev/null -w "%{http_code}" \
    --max-time 1 "${BASE}/api/v1/system-info" 2>/dev/null || true)
CODE="${CODE%%[^0-9]*}"
[ -z "$CODE" ] && CODE="000"
check "GET /api/v1/system-info returns 200" "[ \"$CODE\" = \"200\" ]"

# 11. POST /unload (nothing loaded — 200 or 404 OK)
echo ""
echo "--- /unload edge cases ---"
CODE=$(curl -s -o /dev/null -w "%{http_code}" \
    --max-time 1 -X POST -H "Content-Type: application/json" \
    -d '{}' "${BASE}/api/v1/unload" 2>/dev/null || true)
CODE="${CODE%%[^0-9]*}"
[ -z "$CODE" ] && CODE="000"
check "POST /api/v1/unload returns 200 or 404 (got $CODE)" "[ \"$CODE\" = \"200\" ] || [ \"$CODE\" = \"404\" ]"

# 12. POST /delete with unknown model
CODE=$(curl -s -o /dev/null -w "%{http_code}" \
    --max-time 1 -X POST -H "Content-Type: application/json" \
    -d '{"model_name":"nonexistent-model-abc-123"}' "${BASE}/api/v1/delete" 2>/dev/null || true)
CODE="${CODE%%[^0-9]*}"
[ -z "$CODE" ] && CODE="000"
check "POST /api/v1/delete for unknown model < 500 (got $CODE)" "[ \"$CODE\" -lt 500 ]"

# 13. Registry search validation
echo ""
echo "--- Registry search validation ---"
CODE=$(curl -s -o /dev/null -w "%{http_code}" \
    --max-time 1 "${BASE}/api/v1/registry/search" 2>/dev/null || true)
CODE="${CODE%%[^0-9]*}"
[ -z "$CODE" ] && CODE="000"
check "registry/search without query returns 400" "[ \"$CODE\" = \"400\" ]"

CODE=$(curl -s -o /dev/null -w "%{http_code}" \
    --max-time 1 "${BASE}/api/v1/registry/search?query=test&source=unknown" 2>/dev/null || true)
CODE="${CODE%%[^0-9]*}"
[ -z "$CODE" ] && CODE="000"
check "registry/search with bad source returns 400 (got $CODE)" "[ \"$CODE\" = \"400\" ]"

###############################################################################
# Summary
###############################################################################

echo ""
echo "============================="
echo "Results: $PASS passed, $FAIL failed out of $TOTAL checks"
echo "============================="

if [ "$FAIL" -gt 0 ]; then
    echo "Some checks failed."
    exit 1
fi

echo "All checks passed."
exit 0
