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

# Deterministic local provider-pool capacity smoke.
#
# This starts a loopback-only fake OpenAI-compatible backend, Control Plane,
# Edge, and Node. Two external model aliases share one capacity-1 provider.
# The fake backend holds its first request until this script creates a release
# sentinel, so the second alias must remain queued at Edge. The smoke then
# proves backend peak concurrency stays at one and the Control Plane status
# returns every normal/long counter to zero.

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
TMP_DIR="$(mktemp -d /tmp/iop-provider-capacity-smoke.XXXXXX)"
KEEP_TMP="${IOP_PROVIDER_CAPACITY_SMOKE_KEEP_TMP:-0}"
HTTP_CONNECT_TIMEOUT=1
HTTP_PROBE_TIMEOUT=2
HTTP_COMPLETION_TIMEOUT=50

FAKE_PID=""
CP_PID=""
EDGE_PID=""
NODE_PID=""
FIRST_PID=""
SECOND_PID=""

cleanup() {
    local rc=$?
    local pid
    for pid in "$FIRST_PID" "$SECOND_PID" "$NODE_PID" "$EDGE_PID" "$CP_PID" "$FAKE_PID"; do
        if [ -n "$pid" ] && kill -0 "$pid" 2>/dev/null; then
            kill "$pid" 2>/dev/null || true
        fi
    done
    for pid in "$FIRST_PID" "$SECOND_PID" "$NODE_PID" "$EDGE_PID" "$CP_PID" "$FAKE_PID"; do
        if [ -n "$pid" ]; then
            wait "$pid" 2>/dev/null || true
        fi
    done
    if [ "$rc" -ne 0 ]; then
        echo "[provider-capacity-smoke] FAIL evidence=$TMP_DIR"
        for log_file in "$TMP_DIR"/{fake,control-plane,edge,node}.log; do
            if [ -s "$log_file" ]; then
                echo "=== $(basename "$log_file") ==="
                tail -n 80 "$log_file"
            fi
        done
    fi
    if [ "$KEEP_TMP" = "1" ]; then
        echo "[provider-capacity-smoke] keeping evidence=$TMP_DIR"
    else
        rm -rf "$TMP_DIR"
    fi
}
trap cleanup EXIT

log() {
    echo "[provider-capacity-smoke] $*"
}

die() {
    log "ERROR: $*"
    exit 1
}

for required in curl jq go; do
    command -v "$required" >/dev/null 2>&1 || die "$required is required"
done

declare -a USED_PORTS=()
pick_port() {
    local base="$1"
    local attempt=0 candidate used
    while (( attempt < 100 )); do
        ((attempt += 1))
        candidate=$((base + RANDOM % 1000))
        used=0
        for port in "${USED_PORTS[@]}"; do
            [ "$port" = "$candidate" ] && used=1
        done
        [ "$used" -eq 0 ] || continue
        if ! (echo >/dev/tcp/127.0.0.1/"$candidate") 2>/dev/null; then
            USED_PORTS+=("$candidate")
            echo "$candidate"
            return 0
        fi
    done
    die "could not reserve a loopback port"
}

curl_with_timeout() {
    local max_time="$1"
    shift
    curl --connect-timeout "$HTTP_CONNECT_TIMEOUT" --max-time "$max_time" "$@"
}

curl_with_timeout_exec() {
    local max_time="$1"
    shift
    exec curl --connect-timeout "$HTTP_CONNECT_TIMEOUT" --max-time "$max_time" "$@"
}

curl_probe() {
    curl_with_timeout "$HTTP_PROBE_TIMEOUT" "$@"
}

wait_for_http() {
    local url="$1"
    local label="$2"
    local deadline=$((SECONDS + 30))
    while ! curl_probe -fsS "$url" >/dev/null 2>&1; do
        if (( SECONDS >= deadline )); then
            die "$label did not become ready: $url"
        fi
        sleep 0.2
    done
}

FAKE_PORT="$(pick_port 41000)"
CP_HTTP_PORT="$(pick_port 29000)"
CP_CLIENT_PORT="$(pick_port 30000)"
CP_EDGE_PORT="$(pick_port 31000)"
EDGE_NODE_PORT="$(pick_port 32000)"
EDGE_BOOTSTRAP_PORT="$(pick_port 33000)"
EDGE_OPENAI_PORT="$(pick_port 34000)"

RELEASE_FILE="$TMP_DIR/release"
FAKE_SOURCE="$TMP_DIR/fake_provider.go"
FAKE_BIN="$TMP_DIR/fake-provider"
CP_BIN="$TMP_DIR/control-plane"
EDGE_BIN="$TMP_DIR/iop-edge"
NODE_BIN="$TMP_DIR/iop-node"
CP_CONFIG="$TMP_DIR/control-plane.yaml"
EDGE_CONFIG="$TMP_DIR/edge.yaml"
NODE_CONFIG="$TMP_DIR/node.yaml"
STATUS_URL="http://127.0.0.1:$CP_HTTP_PORT/edges/ornith-edge/status"
OPENAI_URL="http://127.0.0.1:$EDGE_OPENAI_PORT/v1/chat/completions"
FAKE_STATS_URL="http://127.0.0.1:$FAKE_PORT/stats"

cat > "$FAKE_SOURCE" <<'EOF'
package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"sync"
	"time"
)

type chatRequest struct {
	Model  string `json:"model"`
	Stream bool   `json:"stream"`
}

type stats struct {
	mu     sync.Mutex
	active int
	peak   int
	calls  int
}

func main() {
	if len(os.Args) != 4 {
		log.Fatal("usage: fake-provider <listen-addr> <release-file> <served-model>")
	}
	listenAddr, releaseFile, servedModel := os.Args[1], os.Args[2], os.Args[3]
	var state stats
	mux := http.NewServeMux()
	mux.HandleFunc("/hang", func(_ http.ResponseWriter, r *http.Request) {
		select {
		case <-r.Context().Done():
			return
		case <-time.After(10 * time.Second):
			return
		}
	})
	mux.HandleFunc("/v1/models", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = fmt.Fprintf(w, `{"object":"list","data":[{"id":%q}]}`, servedModel)
	})
	mux.HandleFunc("/stats", func(w http.ResponseWriter, _ *http.Request) {
		state.mu.Lock()
		defer state.mu.Unlock()
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]int{"active": state.active, "peak": state.peak, "calls": state.calls})
	})
	mux.HandleFunc("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		var request chatRequest
		if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
			http.Error(w, "invalid JSON", http.StatusBadRequest)
			return
		}
		if request.Model != servedModel {
			http.Error(w, "unexpected model", http.StatusBadRequest)
			return
		}

		state.mu.Lock()
		state.calls++
		call := state.calls
		state.active++
		if state.active > state.peak {
			state.peak = state.active
		}
		state.mu.Unlock()
		defer func() {
			state.mu.Lock()
			state.active--
			state.mu.Unlock()
		}()

		if call == 1 {
			deadline := time.Now().Add(45 * time.Second)
			for {
				if _, err := os.Stat(releaseFile); err == nil {
					break
				}
				if time.Now().After(deadline) {
					http.Error(w, "release sentinel timeout", http.StatusGatewayTimeout)
					return
				}
				time.Sleep(25 * time.Millisecond)
			}
		}

		if request.Stream {
			w.Header().Set("Content-Type", "text/event-stream")
			_, _ = fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ORNITH_CAPACITY_OK\"},\"finish_reason\":\"stop\"}]}\n\n")
			_, _ = fmt.Fprint(w, "data: [DONE]\n\n")
			if flush, ok := w.(http.Flusher); ok {
				flush.Flush()
			}
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = fmt.Fprintf(w, `{"id":"ornith-%d","object":"chat.completion","model":%q,"choices":[{"index":0,"message":{"role":"assistant","content":"ORNITH_CAPACITY_OK"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`, call, servedModel)
	})
	log.Fatal(http.ListenAndServe(listenAddr, mux))
}
EOF

log "building loopback binaries"
go build -o "$FAKE_BIN" "$FAKE_SOURCE"
go build -o "$CP_BIN" "$REPO_ROOT/apps/control-plane/cmd/control-plane"
go build -o "$EDGE_BIN" "$REPO_ROOT/apps/edge/cmd/edge"
go build -o "$NODE_BIN" "$REPO_ROOT/apps/node/cmd/node"

cat > "$CP_CONFIG" <<EOF
server:
  listen: "127.0.0.1:$CP_HTTP_PORT"
  wire_listen: "127.0.0.1:$CP_CLIENT_PORT"
  edge_wire_listen: "127.0.0.1:$CP_EDGE_PORT"
logging:
  level: "error"
  pretty: false
EOF

cat > "$EDGE_CONFIG" <<EOF
edge:
  id: "ornith-edge"
  name: "Ornith Capacity Smoke"
server:
  listen: "127.0.0.1:$EDGE_NODE_PORT"
bootstrap:
  listen: "127.0.0.1:$EDGE_BOOTSTRAP_PORT"
  artifact_dir: "$TMP_DIR/artifacts"
logging:
  level: "error"
metrics:
  port: 0
refresh:
  enabled: false
  listen: "127.0.0.1:0"
openai:
  enabled: true
  listen: "127.0.0.1:$EDGE_OPENAI_PORT"
  timeout_sec: 45
  strict_output: false
a2a:
  enabled: false
  listen: "127.0.0.1:0"
control_plane:
  enabled: true
  wire_addr: "127.0.0.1:$CP_EDGE_PORT"
  reconnect_interval_sec: 1
long_context_threshold_tokens: 100000
provider_pool:
  max_queue: 2
  queue_timeout_ms: 30000
models:
  - id: "ornith:35b"
    context_window_tokens: 4096
    providers:
      ornith-provider: "ornith-served"
  - id: "ornith-fast"
    context_window_tokens: 4096
    providers:
      ornith-provider: "ornith-served"
nodes:
  - id: "ornith-node"
    alias: "ornith-node"
    token: "capacity-smoke-token"
    providers:
      - id: "ornith-provider"
        type: "openai_api"
        category: "api"
        provider: "vllm"
        endpoint: "http://127.0.0.1:$FAKE_PORT"
        models: ["ornith-served"]
        health: "available"
        capacity: 1
        long_context_capacity: 1
        total_context_tokens: 4096
        request_timeout_ms: 45000
EOF

cat > "$NODE_CONFIG" <<EOF
transport:
  edge_addr: "127.0.0.1:$EDGE_NODE_PORT"
  token: "capacity-smoke-token"
reconnect:
  interval_sec: 1
  max_attempts: 0
logging:
  level: "error"
metrics:
  port: 0
node:
  id: "ornith-node"
  alias: "ornith-node"
EOF

"$FAKE_BIN" "127.0.0.1:$FAKE_PORT" "$RELEASE_FILE" "ornith-served" >"$TMP_DIR/fake.log" 2>&1 &
FAKE_PID=$!
wait_for_http "http://127.0.0.1:$FAKE_PORT/v1/models" "fake provider"
hang_started=$SECONDS
if curl_probe -fsS "http://127.0.0.1:$FAKE_PORT/hang" >"$TMP_DIR/hang-probe.out" 2>"$TMP_DIR/hang-probe.err"; then
    die "bounded probe unexpectedly succeeded against hanging peer"
fi
hang_elapsed=$((SECONDS - hang_started))
if (( hang_elapsed > HTTP_PROBE_TIMEOUT + 2 )); then
    die "bounded probe exceeded ${HTTP_PROBE_TIMEOUT}s deadline: ${hang_elapsed}s"
fi
log "metrics=disabled bounded_probe_timeout_sec=$hang_elapsed"

"$CP_BIN" serve --config "$CP_CONFIG" >"$TMP_DIR/control-plane.log" 2>&1 &
CP_PID=$!
wait_for_http "http://127.0.0.1:$CP_HTTP_PORT/healthz" "Control Plane"

"$EDGE_BIN" serve --config "$EDGE_CONFIG" >"$TMP_DIR/edge.log" 2>&1 &
EDGE_PID=$!

edge_deadline=$((SECONDS + 30))
while ! curl_probe -fsS "http://127.0.0.1:$CP_HTTP_PORT/edges/ornith-edge" >/dev/null 2>&1; do
    if (( SECONDS >= edge_deadline )); then
        die "Edge did not enroll with Control Plane"
    fi
    sleep 0.2
done

"$NODE_BIN" serve --config "$NODE_CONFIG" >"$TMP_DIR/node.log" 2>&1 &
NODE_PID=$!

connected_deadline=$((SECONDS + 30))
while true; do
    if curl_probe -fsS "$STATUS_URL" >"$TMP_DIR/status-connected.json" 2>/dev/null \
        && jq -e '.nodes | any(.[]; .node_id == "ornith-node" and .connected == true)' "$TMP_DIR/status-connected.json" >/dev/null; then
        break
    fi
    if (( SECONDS >= connected_deadline )); then
        die "Node did not become dispatch-ready"
    fi
    sleep 0.2
done

curl_probe -fsS "http://127.0.0.1:$EDGE_OPENAI_PORT/v1/models" >"$TMP_DIR/models.json"
jq -e '[.data[].id] | index("ornith:35b") and index("ornith-fast")' "$TMP_DIR/models.json" >/dev/null \
    || die "OpenAI model catalog did not expose both aliases"

fire_alias() {
    local alias="$1"
    local prefix="$2"
    jq -n --arg model "$alias" '{model:$model, stream:false, messages:[{role:"user", content:"capacity smoke"}]}' >"$prefix.request.json"
    (
        curl_with_timeout_exec "$HTTP_COMPLETION_TIMEOUT" \
            -sS -o "$prefix.response.json" -w '%{http_code}' \
            -H 'Content-Type: application/json' \
            --data-binary "@$prefix.request.json" \
            "$OPENAI_URL" >"$prefix.code" 2>"$prefix.curl.err"
    ) &
    REQUEST_PID=$!
}

wait_for_request() {
    local pid="$1"
    local prefix="$2"
    local request_rc
    if wait "$pid"; then
        request_rc=0
    else
        request_rc=$?
    fi
    printf '%s\n' "$request_rc" >"$prefix.rc"
}

fire_alias "ornith:35b" "$TMP_DIR/ornith-35b"
FIRST_PID="$REQUEST_PID"
backend_deadline=$((SECONDS + 20))
while ! curl_probe -fsS "$FAKE_STATS_URL" >"$TMP_DIR/backend-before-second.json" 2>/dev/null \
    || ! jq -e '.calls == 1 and .active == 1 and .peak == 1' "$TMP_DIR/backend-before-second.json" >/dev/null; do
    if (( SECONDS >= backend_deadline )); then
        die "first alias did not reach the held fake backend"
    fi
    sleep 0.1
done

fire_alias "ornith-fast" "$TMP_DIR/ornith-fast"
SECOND_PID="$REQUEST_PID"
queue_deadline=$((SECONDS + 20))
while true; do
    if curl_probe -fsS "$STATUS_URL" >"$TMP_DIR/status-queued.json" 2>/dev/null \
        && jq -e '
            [.nodes[].provider_snapshots[]? | select(.id == "ornith-provider")] as $providers
            | $providers | length == 1
            and $providers[0].capacity == 1
            and $providers[0].in_flight == 1
            and $providers[0].queued >= 1
            and $providers[0].long_context_capacity == 1
            and $providers[0].long_in_flight == 0
            and $providers[0].long_queued == 0
        ' "$TMP_DIR/status-queued.json" >/dev/null; then
        break
    fi
    if (( SECONDS >= queue_deadline )); then
        die "Control Plane status did not observe the shared-provider queue"
    fi
    sleep 0.1
done

touch "$RELEASE_FILE"
wait_for_request "$FIRST_PID" "$TMP_DIR/ornith-35b"
wait_for_request "$SECOND_PID" "$TMP_DIR/ornith-fast"

for prefix in "$TMP_DIR/ornith-35b" "$TMP_DIR/ornith-fast"; do
    [ "$(cat "$prefix.rc")" = "0" ] || die "$(basename "$prefix") curl failed"
    [ "$(cat "$prefix.code")" = "200" ] || die "$(basename "$prefix") returned HTTP $(cat "$prefix.code")"
    grep -q 'ORNITH_CAPACITY_OK' "$prefix.response.json" || die "$(basename "$prefix") response lacked success marker"
done

curl_probe -fsS "$FAKE_STATS_URL" >"$TMP_DIR/backend-final.json"
jq -e '.calls == 2 and .active == 0 and .peak == 1' "$TMP_DIR/backend-final.json" >/dev/null \
    || die "fake backend stats did not recover with peak=1"

final_status_is_online_and_recovered() {
    local snapshot="$1"
    jq -e '
        [.nodes[] | select(.node_id == "ornith-node")] as $nodes
        | [$nodes[0].provider_snapshots[]? | select(.id == "ornith-provider")] as $providers
        | ($nodes | length == 1)
          and ($nodes[0].connected == true)
          and ($providers | length == 1)
          and ($providers[0].status == "available")
          and ($providers[0].health == "available")
          and ($providers[0].capacity == 1)
          and ($providers[0].long_context_capacity == 1)
          and ($providers[0].in_flight == 0)
          and ($providers[0].queued == 0)
          and ($providers[0].long_in_flight == 0)
          and ($providers[0].long_queued == 0)
    ' "$snapshot" >/dev/null
}

cat >"$TMP_DIR/status-offline.json" <<'EOF'
{
  "nodes": [
    {
      "node_id": "ornith-node",
      "connected": false,
      "provider_snapshots": [
        {
          "id": "ornith-provider",
          "status": "unavailable",
          "health": "offline",
          "capacity": 0,
          "in_flight": 0,
          "queued": 0,
          "long_context_capacity": 0,
          "long_in_flight": 0,
          "long_queued": 0
        }
      ]
    }
  ]
}
EOF
if final_status_is_online_and_recovered "$TMP_DIR/status-offline.json"; then
    die "final recovery predicate accepted an offline snapshot"
fi
log "offline_snapshot_rejected=true"

final_deadline=$((SECONDS + 20))
while true; do
    if curl_probe -fsS "$STATUS_URL" >"$TMP_DIR/status-final.json" 2>/dev/null \
        && final_status_is_online_and_recovered "$TMP_DIR/status-final.json"; then
        break
    fi
    if (( SECONDS >= final_deadline )); then
        die "provider counters did not recover to zero"
    fi
    sleep 0.1
done

log "aliases=ornith:35b,ornith-fast queue_observed=true"
log "backend=$(jq -c '{calls,active,peak}' "$TMP_DIR/backend-final.json")"
log "final_provider=$(jq -c '[.nodes[] | select(.node_id == "ornith-node") | {node_id,connected,providers:[.provider_snapshots[]? | select(.id == "ornith-provider") | {id,status,health,capacity,in_flight,queued,long_context_capacity,long_in_flight,long_queued}]}]' "$TMP_DIR/status-final.json")"
log "PASS evidence=$TMP_DIR"