#!/usr/bin/env bash

set -euo pipefail

repo_root="$(cd "$(dirname "$0")/.." && pwd -P)"
vps_host="${PI_AGI_VPS_HOST:-kortix-prod}"
vps_root="${PI_AGI_VPS_ROOT:-/root/pi-agi-vps-test}"
arena_root="${PI_AGI_ARENA_ROOT:-/opt/agent-harness-arena}"
configured_agent_dir="${PI_AGI_SANDBOX_AGENT_DIR:-${PI_CODING_AGENT_DIR:-$HOME/.pi/agent}}"
release_dir="${PI_AGI_RELEASE_DIR:-$repo_root/.release}"
ssh_args=(-o BatchMode=yes)

run_arena=0
publish_npm=0
skip_tests=0
skip_vps_preflight=0

die() {
	printf 'pi-agi release: %s\n' "$*" >&2
	exit 1
}

usage() {
	cat <<'EOF'
Usage: npm run release:all -- [options]

Builds one package from the current working tree, then uses it for the local Pi
installation, the production VPS sandbox, and the VPS Arena image.

Default workflow:
  1. npm run check, npm test, and npm run build
  2. Create a versioned tarball under .release/
  3. Install this checkout into local Pi and remove only npm:pi-agi
  4. Run the VPS sandbox preflight, deploy/restart it, and verify status
  5. Upload the same tarball to Arena, write exact provenance, preflight, and
     rebuild the pi-agi benchmark image

Options:
  --run-arena          Start the 30-minute Luna high pi-agi benchmark after rebuilding
  --publish-npm        Publish the tarball to npm (requires a clean checkout)
  --skip-tests         Skip npm run check/test/build
  --skip-vps-preflight Skip the separate VPS preflight before deployment
  -h, --help           Show this help

Environment overrides:
  PI_AGI_SANDBOX_AGENT_DIR  Local Pi model-config source
  PI_AGI_VPS_HOST           SSH alias (default kortix-prod)
  PI_AGI_VPS_ROOT           Sandbox deployment root
  PI_AGI_ARENA_ROOT         Arena checkout (default /opt/agent-harness-arena)
  PI_AGI_RELEASE_DIR        Persistent local artifact directory
EOF
}

while [[ $# -gt 0 ]]; do
	case "$1" in
		--run-arena) run_arena=1 ;;
		--publish-npm) publish_npm=1 ;;
		--skip-tests) skip_tests=1 ;;
		--skip-vps-preflight) skip_vps_preflight=1 ;;
		-h | --help) usage; exit 0 ;;
		*) usage; die "unknown option: $1" ;;
	esac
	shift
done

require_command() {
	command -v "$1" >/dev/null 2>&1 || die "required command not found: $1"
}

remote() {
	ssh "${ssh_args[@]}" "$vps_host" "$@"
}

validate_inputs() {
	[[ "$vps_host" =~ ^[A-Za-z0-9._-]+$ ]] || die "unsafe VPS host alias: $vps_host"
	[[ "$vps_root" =~ ^/root/[A-Za-z0-9._/-]+$ && "$vps_root" != "/root/" ]] ||
		die "unsafe VPS root: $vps_root"
	[[ "$arena_root" =~ ^/opt/[A-Za-z0-9._/-]+$ && "$arena_root" != "/opt/" ]] ||
		die "unsafe Arena root: $arena_root"
	[[ -d "$configured_agent_dir" ]] || die "Pi config directory does not exist: $configured_agent_dir"
	[[ -f "$configured_agent_dir/auth.json" || -f "$configured_agent_dir/models.json" ]] ||
		die "Pi config has neither auth.json nor models.json: $configured_agent_dir"
}

local_settings_path() {
	printf '%s/settings.json\n' "${PI_CODING_AGENT_DIR:-$HOME/.pi/agent}"
}

has_local_package() {
	local settings_path="$1" package_name="$2"
	[[ -f "$settings_path" ]] || return 1
	node -e '
		const fs = require("node:fs");
		const [path, wanted] = process.argv.slice(1);
		const parsed = JSON.parse(fs.readFileSync(path, "utf8"));
		process.exit(Array.isArray(parsed.packages) && parsed.packages.includes(wanted) ? 0 : 1);
	' "$settings_path" "$package_name"
}

install_local() {
	local settings_path list_output
	settings_path="$(local_settings_path)"
	printf '\n==> Installing current checkout into local Pi\n'
	pi install "$repo_root"
	if has_local_package "$settings_path" "npm:pi-agi"; then
		pi remove npm:pi-agi
	fi
	list_output="$(pi list)"
	grep -F "$repo_root" <<<"$list_output" >/dev/null || die "local Pi does not resolve pi-agi to $repo_root"
	if has_local_package "$settings_path" "npm:pi-agi"; then
		die "npm:pi-agi remains in local Pi settings"
	fi
	printf '%s\n' "$list_output"
}

build_release() {
	local temp_dir="$1" pack_json pack_name packed_path version timestamp digest release_tag artifact
	printf '\n==> Packaging exact working tree\n' >&2
	pack_json="$(npm pack --pack-destination "$temp_dir" --json)"
	pack_name="$(node -e '
		let body = "";
		process.stdin.on("data", (chunk) => body += chunk);
		process.stdin.on("end", () => {
			const parsed = JSON.parse(body);
			if (!Array.isArray(parsed) || typeof parsed[0]?.filename !== "string") process.exit(1);
			process.stdout.write(parsed[0].filename);
		});
	' <<<"$pack_json")" || die "npm pack did not report an artifact"
	packed_path="$temp_dir/$pack_name"
	[[ -f "$packed_path" ]] || die "npm pack artifact is missing: $packed_path"
	version="$(node -p 'require("./package.json").version')"
	timestamp="$(date -u +%Y%m%dT%H%M%SZ | tr '[:upper:]' '[:lower:]')"
	digest="$(sha256sum "$packed_path" | awk '{print $1}')"
	release_tag="$timestamp-${digest:0:12}"
	mkdir -p "$release_dir"
	release_dir="$(cd "$release_dir" && pwd -P)"
	artifact="$release_dir/pi-agi-$version-$release_tag.tgz"
	install -m 0644 "$packed_path" "$artifact"
	printf 'Artifact: %s\nSHA256:   %s\n' "$artifact" "$digest" >&2
	printf '%s\t%s\t%s\t%s\n' "$artifact" "$digest" "$release_tag" "$version"
}

deploy_vps() {
	local artifact="$1"
	printf '\n==> Deploying production VPS sandbox\n'
	if [[ "$skip_vps_preflight" -eq 0 ]]; then
		PI_AGI_VPS_SOURCE_ARCHIVE="$artifact" \
		PI_AGI_SANDBOX_AGENT_DIR="$configured_agent_dir" \
		PI_AGI_VPS_HOST="$vps_host" \
		PI_AGI_VPS_ROOT="$vps_root" \
			bash "$repo_root/scripts/agi-vps-sandbox.sh" preflight
	fi
	PI_AGI_VPS_SOURCE_ARCHIVE="$artifact" \
	PI_AGI_SANDBOX_AGENT_DIR="$configured_agent_dir" \
	PI_AGI_VPS_HOST="$vps_host" \
	PI_AGI_VPS_ROOT="$vps_root" \
		bash "$repo_root/scripts/agi-vps-sandbox.sh" start
	PI_AGI_VPS_HOST="$vps_host" PI_AGI_VPS_ROOT="$vps_root" \
		bash "$repo_root/scripts/agi-vps-sandbox.sh" status
}

update_arena_config() {
	local release_version="$1" arena_image="$2"
	ARENA_RELEASE_VERSION="$release_version" ARENA_RELEASE_IMAGE="$arena_image" \
		ssh "${ssh_args[@]}" "$vps_host" \
		"ARENA_RELEASE_VERSION='$release_version' ARENA_RELEASE_IMAGE='$arena_image' '$arena_root/.venv/bin/python' - '$arena_root/config/experiment.yaml' '$arena_root/config/experiment.local.yaml'" <<'PY'
from __future__ import annotations

import os
import sys
from pathlib import Path

import yaml

base = Path(sys.argv[1])
target = Path(sys.argv[2])
source = target if target.exists() else base
payload = yaml.safe_load(source.read_text())
payload["budget_seconds"] = 1800
models = payload.get("models")
if isinstance(models, dict):
    payload["default_model"] = "gpt-5.6-luna"
    models["gpt-5.6-luna"] = {
        "display_name": "GPT-5.6 Luna",
        "models_json": "~/.pi/agent/models.json",
        "provider": "gateway",
        "model": "gpt-5.6-luna",
        "agent_model_names": {"pi-agi": "gateway/gpt-5.6-luna"},
    }
else:
    model_source = payload.setdefault("model_source", {})
    model_source["provider"] = "gateway"
    model_source["model"] = "gpt-5.6-luna"
agent = payload["agents"]["pi-agi"]
if isinstance(agent, str):
    profiles = payload.get("profiles") or {}
    profile_dir = Path(profiles.get("agents", "profiles/agents"))
    if not profile_dir.is_absolute():
        profile_dir = base.parent.parent / profile_dir
    agent = yaml.safe_load((profile_dir / f"{agent}.yaml").read_text())
    payload["agents"]["pi-agi"] = agent
agent["version"] = os.environ["ARENA_RELEASE_VERSION"]
agent["model_name"] = "gateway/gpt-5.6-luna"
agent["effort"] = "high"
agent.setdefault("kwargs", {})["thinking"] = "high"
agent["runtime"]["image"] = os.environ["ARENA_RELEASE_IMAGE"]
temporary = target.with_suffix(target.suffix + ".tmp")
temporary.write_text(yaml.safe_dump(payload, sort_keys=False))
temporary.replace(target)
PY
}

preflight_arena_pi_agi() {
	ssh "${ssh_args[@]}" "$vps_host" "cd '$arena_root' && .venv/bin/python -" <<'PY'
from arena.config import load_config
from arena.preflight import run_preflight

checks = run_preflight(load_config(), {"pi-agi"})
for check in checks:
    marker = "PASS" if check.ok else "FAIL"
    print(f"{marker:4}  {check.name:24} {check.detail}")
raise SystemExit(0 if all(check.ok or not check.blocking for check in checks) else 1)
PY
}

deploy_arena() {
	local artifact="$1" digest="$2" release_tag="$3" version="$4"
	local remote_upload release_version arena_image remote_digest
	remote_upload="$arena_root/vendor/.pi-agi-current-$release_tag.tgz"
	release_version="$version+$release_tag"
	arena_image="agent-harness-arena/sokoban-pi-agi:$version-$release_tag"

	printf '\n==> Refreshing VPS Arena bundle\n'
	remote "test -d '$arena_root' && test -x '$arena_root/.venv/bin/arena' && install -d -m 0755 '$arena_root/vendor' '$arena_root/.deploy-backups'"
	scp "${ssh_args[@]}" "$artifact" "$vps_host:$remote_upload"
	remote_digest="$(remote "sha256sum '$remote_upload'" | awk '{print $1}')"
	[[ "$remote_digest" == "$digest" ]] || die "Arena upload checksum mismatch"
	remote "set -eu
		backup='$arena_root/.deploy-backups/pi-agi-$release_tag'
		install -d -m 0755 \"\$backup\"
		if [ -f '$arena_root/vendor/pi-agi-current.tgz' ]; then cp -p '$arena_root/vendor/pi-agi-current.tgz' \"\$backup/\"; fi
		if [ -f '$arena_root/config/experiment.local.yaml' ]; then cp -p '$arena_root/config/experiment.local.yaml' \"\$backup/\"; fi
		install -m 0644 '$remote_upload' '$arena_root/vendor/pi-agi-current.tgz'
		rm -f -- '$remote_upload'"
	update_arena_config "$release_version" "$arena_image"
	remote "cd '$arena_root' && .venv/bin/arena build-images pi-agi"
	ARENA_RELEASE_TAG="$version-$release_tag" \
		ssh "${ssh_args[@]}" "$vps_host" \
		"ARENA_RELEASE_TAG='$version-$release_tag' '$arena_root/.venv/bin/python' - '$arena_root'" <<'PY'
from __future__ import annotations

import os
import subprocess
import sys
from pathlib import Path

from arena.config import load_config
from arena.images import build_task_image, task_image_tag
from arena.tasks.definition import resolve_task

root = Path(sys.argv[1])
config = load_config()
runtime = config.agents["pi-agi"].runtime
release_tag = os.environ["ARENA_RELEASE_TAG"]
for task_id in ("super-mario", "snake_maze_campaign", "2048"):
    task = resolve_task(config, task_id)
    if build_task_image(task) != 0:
        raise SystemExit(f"failed to build task image for {task_id}")
    target = f"agent-harness-arena/{task_id}-pi-agi:{release_tag}"
    command = [
        "docker",
        "build",
        "--file",
        str(runtime.dockerfile),
        "--tag",
        target,
    ]
    build_args = {**runtime.build_args, "BASE_IMAGE": task_image_tag(task)}
    for key, value in sorted(build_args.items()):
        command.extend(["--build-arg", f"{key}={value}"])
    command.append(str(runtime.build_context))
    subprocess.run(command, cwd=root, check=True)
PY
	preflight_arena_pi_agi
	remote "set -eu
		cd '$arena_root'
		test \"\$(sha256sum vendor/pi-agi-current.tgz | awk '{print \$1}')\" = '$digest'
		grep -F '$release_version' config/experiment.local.yaml >/dev/null
		grep -F '$arena_image' config/experiment.local.yaml >/dev/null
		grep -F 'model: gpt-5.6-luna' config/experiment.local.yaml >/dev/null
		grep -F 'model_name: gateway/gpt-5.6-luna' config/experiment.local.yaml >/dev/null
		grep -F 'budget_seconds: 1800' config/experiment.local.yaml >/dev/null
		grep -F 'effort: high' config/experiment.local.yaml >/dev/null
		grep -F 'thinking: high' config/experiment.local.yaml >/dev/null
		docker image inspect '$arena_image' >/dev/null
		docker image inspect 'agent-harness-arena/super-mario-pi-agi:$version-$release_tag' >/dev/null
		docker image inspect 'agent-harness-arena/snake_maze_campaign-pi-agi:$version-$release_tag' >/dev/null
		docker image inspect 'agent-harness-arena/2048-pi-agi:$version-$release_tag' >/dev/null"
	printf 'Arena image: %s\nViewer: http://arena.178.104.6.186.sslip.io:8788/\n' "$arena_image"

	if [[ "$run_arena" -eq 1 ]]; then
		printf '\n==> Starting paid pi-agi Arena benchmark\n'
		start_arena
	else
		printf 'Benchmark not started. Run with --run-arena when intended.\n'
	fi
}

start_arena() {
	local stamp session log
	stamp="$(date -u +%Y%m%dT%H%M%SZ | tr '[:upper:]' '[:lower:]')"
	session="arena-sokoban-pi-agi-$stamp"
	log="$arena_root/runs/$session.log"
	ssh "${ssh_args[@]}" "$vps_host" bash -s -- "$arena_root" "$session" "$log" <<'SH'
set -eu
arena_root="$1"
session="$2"
log="$3"
command -v tmux >/dev/null
install -d -m 0755 "$arena_root/runs"
if ps -eo args= | grep -F "$arena_root/.venv/bin/python .venv/bin/arena run pi-agi" | grep -v grep >/dev/null; then
	printf 'Refusing to overlap an active pi-agi Arena benchmark.\n' >&2
	exit 1
fi
tmux new-session -d -s "$session" -c "$arena_root" \
	"env PATH=\"$arena_root/.venv/bin:$PATH\" PYTHONUNBUFFERED=1 .venv/bin/arena run pi-agi > \"$log\" 2>&1"
i=0
while [ "$i" -lt 20 ]; do
	if grep -q '^Starting ' "$log" 2>/dev/null; then break; fi
	if ! tmux has-session -t "$session" 2>/dev/null; then break; fi
	i=$((i + 1))
	sleep 1
done
grep -m1 '^Starting ' "$log" 2>/dev/null || true
if ! tmux has-session -t "$session" 2>/dev/null; then
	printf 'Arena session exited during startup; log follows:\n' >&2
	tail -80 "$log" >&2 || true
	exit 1
fi
SH
	printf 'Arena tmux: %s\nArena log:  %s:%s\n' "$session" "$vps_host" "$log"
}

for command in npm node pi ssh scp tar sha256sum awk grep install; do
	require_command "$command"
done
validate_inputs

cd "$repo_root"
if [[ "$skip_tests" -eq 0 ]]; then
	printf '==> Validating checkout\n'
	npm run check
	npm test
	npm run build
fi

temp_dir="$(mktemp -d)"
trap 'rm -rf -- "$temp_dir"' EXIT
IFS=$'\t' read -r artifact digest release_tag version < <(build_release "$temp_dir")

install_local

if [[ "$publish_npm" -eq 1 ]]; then
	[[ -z "$(git status --porcelain)" ]] || die "--publish-npm requires a clean checkout"
	printf '\n==> Publishing %s to npm\n' "$artifact"
	npm publish "$artifact"
	[[ "$(npm view "pi-agi@$version" version)" == "$version" ]] || die "npm registry did not report pi-agi@$version"
fi

deploy_vps "$artifact"
deploy_arena "$artifact" "$digest" "$release_tag" "$version"

printf '\nRelease complete.\nArtifact: %s\nRelease:  %s+%s\n' "$artifact" "$version" "$release_tag"