#!/usr/bin/env bash

set -eo pipefail

repo_root="$(cd "$(dirname "$0")/.." && pwd -P)"
compose_file="$repo_root/sandbox/open-mission/compose.yaml"
project="${PI_AGI_SANDBOX_PROJECT:-}"
[[ -n "$project" ]] || project="pi-agi-mission-$UID"
tmux_session="${PI_AGI_SANDBOX_TMUX_SESSION:-agi-open-mission}"
configured_agent_dir="${PI_AGI_SANDBOX_AGENT_DIR:-}"
default_agent_dir="$repo_root/sandbox/open-mission/agent-config"

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

usage() {
	cat <<'EOF'
Usage: npm run sandbox:mission:local -- <command> [argument]

Commands:
  preflight          Verify isolation, public egress, and package installation.
  start              Start the proxy + agent containers and local Pi in tmux.
  attach             Attach to the Pi TUI inside the agent container.
  run                Enable /agi and paste the real open-ended mission.
  status             Show owned services, TUI process, goal, wakes, and traces.
  monitor [seconds]  Reprint status until the goal is genuinely met.
  repair-proxy       Rebuild/restart only the proxy; preserve agent and TUI.
  nudge [message]    Continue an active idle goal after infrastructure recovery.
  recover-session [message]
                     Fresh Pi conversation; preserve goal and workspace.
  capture [file]     Print TUI scrollback or save it to a host file.
  export [dir]       Copy the preserved workspace to a host directory.
  stop               Gracefully stop only this project; keep its workspace.

Required for start:
  PI_AGI_SANDBOX_AGENT_DIR=/absolute/path/to/a/dedicated/pi-agent-config

Optional:
  PI_AGI_SANDBOX_PROJECT       Compose project (default pi-agi-mission-$UID)
  PI_AGI_SANDBOX_TMUX_SESSION  TUI session (default agi-open-mission)

Never use your normal ~/.pi/agent directory. No host browser/mail profile is
mounted. The dedicated config should contain only the model access needed by Pi.
EOF
}

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

validate_project() {
	[[ "$project" =~ ^[a-z0-9][a-z0-9_-]*$ ]] || die "unsafe Compose project name: $project"
	[[ "$tmux_session" =~ ^[A-Za-z0-9_.-]+$ ]] || die "unsafe tmux session name: $tmux_session"
}

prepare_compose_env() {
	local selected="$configured_agent_dir"
	[[ -n "$selected" ]] || selected="$default_agent_dir"
	[[ -d "$selected" ]] || die "agent config directory does not exist: $selected"
	export PI_AGI_SANDBOX_AGENT_DIR
	PI_AGI_SANDBOX_AGENT_DIR="$(cd "$selected" && pwd -P)"
}

compose() {
	docker compose --project-name "$project" --file "$compose_file" "$@"
}

container_id() {
	compose ps -q "$1"
}

verify_service() {
	local service="$1"
	local id stopped_id actual_project actual_service
	id="$(container_id "$service")"
	if [[ -z "$id" ]]; then
		stopped_id="$(compose ps -aq "$service")"
		if [[ -n "$stopped_id" ]]; then
			printf 'Service %s exited. Container state:\n' "$service" >&2
			docker inspect --format 'id={{.Id}} status={{.State.Status}} exit={{.State.ExitCode}} error={{.State.Error}}' "$stopped_id" >&2
			printf 'Container logs:\n' >&2
			compose logs --no-color --tail 100 "$service" >&2 || true
		fi
		die "service is not running: $service"
	fi
	actual_project="$(docker inspect --format '{{ index .Config.Labels "com.docker.compose.project" }}' "$id")"
	actual_service="$(docker inspect --format '{{ index .Config.Labels "com.docker.compose.service" }}' "$id")"
	[[ "$actual_project" == "$project" && "$actual_service" == "$service" ]] ||
		die "refusing unowned container $id ($actual_project/$actual_service)"
	printf '%s\n' "$id"
}

wait_for_proxy() {
	local id health
	id="$(verify_service proxy)"
	for _ in {1..60}; do
		health="$(docker inspect --format '{{ if .State.Health }}{{ .State.Health.Status }}{{ else }}none{{ end }}' "$id")"
		[[ "$health" == "healthy" ]] && return
		[[ "$health" == "unhealthy" ]] && die "egress proxy became unhealthy"
		sleep 0.5
	done
	die "egress proxy did not become healthy"
}

require_dedicated_agent_config() {
	[[ -n "$configured_agent_dir" ]] ||
		die "set PI_AGI_SANDBOX_AGENT_DIR to a dedicated Pi config; see sandbox/open-mission/agent-config/README.md"
	local resolved normal
	resolved="$(cd "$configured_agent_dir" && pwd -P)"
	normal="$(cd "$HOME/.pi/agent" 2>/dev/null && pwd -P || true)"
	[[ -z "$normal" || "$resolved" != "$normal" ]] ||
		die "refusing to mount the normal host Pi config; create a dedicated sandbox config"
	find "$resolved" -maxdepth 2 -type f ! -name README.md ! -name .gitignore -print -quit | grep -q . ||
		die "dedicated agent config is empty: $resolved"
}

start_services() {
	compose up -d --build proxy agent
	wait_for_proxy
	verify_service agent >/dev/null
}

verify_public_egress() {
	compose exec -T agent curl -fsS --max-time 20 https://example.com >/dev/null
	printf 'Public HTTPS through filtered proxy: OK\n'
}

verify_direct_egress_blocked() {
	if compose exec -T agent env \
		-u HTTP_PROXY -u HTTPS_PROXY -u http_proxy -u https_proxy -u ALL_PROXY -u all_proxy \
		curl -fsS --connect-timeout 3 --max-time 5 https://example.com >/dev/null 2>&1; then
		die "agent unexpectedly has direct internet egress"
	fi
	printf 'Direct agent egress: blocked\n'
}

verify_private_targets_blocked() {
	if compose exec -T agent curl -fsS --noproxy '' --proxy http://proxy:3128 --max-time 5 \
		http://169.254.169.254/latest/meta-data/ >/dev/null 2>&1; then
		die "proxy unexpectedly allowed cloud metadata address"
	fi
	if compose exec -T agent curl -fsS --noproxy '' --proxy http://proxy:3128 --max-time 5 \
		http://127.0.0.1/ >/dev/null 2>&1; then
		die "proxy unexpectedly allowed loopback target"
	fi
	printf 'Loopback/private/metadata proxy targets: blocked\n'
}

verify_install_capability() {
	compose exec -T agent sh -lc 'test "$(id -u)" = 0 && test -w /usr/local/bin && command -v apt-get >/dev/null' ||
		die "agent does not have container-root package installation capability"
	compose exec -T agent apt-get update -qq >/dev/null ||
		die "container package manager could not reach its repositories"
	printf 'Container-root package installation through filtered proxy: OK\n'
}

verify_container_isolation() {
	local id network
	id="$(verify_service agent)"
	if ! docker inspect "$id" | jq -e '
		.[0].HostConfig.ReadonlyRootfs == false and
		(.[0].Config.User == "" or .[0].Config.User == "root" or .[0].Config.User == "0") and
		(.[0].HostConfig.CapDrop | index("ALL") != null) and
		((.[0].HostConfig.CapAdd | map(sub("^CAP_"; "")) | sort) == (["CHOWN", "DAC_OVERRIDE", "FOWNER", "SETGID", "SETUID"] | sort)) and
		(.[0].HostConfig.SecurityOpt | any(. == "no-new-privileges:true" or . == "no-new-privileges")) and
		([.[0].Mounts[] | select(.Type == "bind")] | length == 1) and
		([.[0].Mounts[] | select(.Type == "bind")][0] | .Destination == "/config-ro" and .RW == false) and
		([.[0].Mounts[] | select(.Type == "volume")] | length == 1) and
		([.[0].Mounts[] | select(.Type == "volume")][0].Destination == "/workspace") and
		([.[0].Mounts[].Destination] | index("/var/run/docker.sock") == null)
	' >/dev/null; then
		docker inspect "$id" | jq '.[0] | {
			user: .Config.User,
			readonlyRootfs: .HostConfig.ReadonlyRootfs,
			capDrop: .HostConfig.CapDrop,
			capAdd: .HostConfig.CapAdd,
			securityOpt: .HostConfig.SecurityOpt,
			mounts: [.Mounts[] | {type: .Type, destination: .Destination, rw: .RW}]
		}' >&2
		die "agent mount or hardening policy differs from the expected sandbox"
	fi
	network="$(docker inspect --format '{{ range $name, $_ := .NetworkSettings.Networks }}{{ printf "%s\n" $name }}{{ end }}' "$id")"
	[[ "$network" == "${project}_sandbox" ]] || die "agent is attached to unexpected networks: $network"
	printf 'Host mounts/socket: absent; container root enabled; only package-install capabilities retained\n'
	printf 'Agent network: internal sandbox only\n'
}

verify_tmux() {
	verify_service agent >/dev/null
	compose exec -T agent tmux has-session -t "$tmux_session" 2>/dev/null ||
		die "Pi tmux session is not running; use start"
	local marker
	marker="$(compose exec -T agent tmux show-options -v -t "$tmux_session" @pi_agi_sandbox_project 2>/dev/null || true)"
	[[ "$marker" == "$project" ]] || die "refusing unowned tmux session (marker: ${marker:-missing})"
}

require_live_tmux() {
	verify_tmux
	local dead exit_status
	dead="$(compose exec -T agent tmux display-message -p -t "$tmux_session:0.0" '#{pane_dead}')"
	if [[ "$dead" == "1" ]]; then
		exit_status="$(compose exec -T agent tmux display-message -p -t "$tmux_session:0.0" '#{pane_dead_status}')"
		printf 'Pi TUI exited (status %s). Last scrollback:\n' "${exit_status:-unknown}" >&2
		compose exec -T agent tmux capture-pane -p -S -200 -t "$tmux_session:0.0" >&2 || true
		die "Pi tmux pane is not live"
	fi
}

start_pi() {
	verify_service agent >/dev/null
	if compose exec -T agent tmux has-session -t "$tmux_session" 2>/dev/null; then
		die "Pi tmux session already exists: $tmux_session"
	fi
	local launch attempt dead output
	launch='exec pi -ne -e /opt/pi-agi/src/index.ts --approve --session-dir /workspace/sessions --name agi-open-mission'
	compose exec -T agent tmux new-session -d -s "$tmux_session" -c /workspace 'exec sleep infinity'
	compose exec -T agent tmux set-option -t "$tmux_session" @pi_agi_sandbox_project "$project"
	compose exec -T agent tmux set-option -w -t "$tmux_session" history-limit 100000
	compose exec -T agent tmux set-option -w -t "$tmux_session" remain-on-exit on
	compose exec -T agent tmux respawn-pane -k -t "$tmux_session:0.0" "$launch"
	for attempt in {1..60}; do
		dead="$(compose exec -T agent tmux display-message -p -t "$tmux_session:0.0" '#{pane_dead}')"
		[[ "$dead" == "1" ]] && require_live_tmux
		output="$(compose exec -T agent tmux capture-pane -p -S -200 -t "$tmux_session:0.0")"
		[[ "$output" == *"Pi can explain its own features"* ]] && return
		sleep 0.5
	done
	printf '%s\n' "$output" >&2
	die "Pi TUI did not become ready within 30 seconds"
}

start() {
	require_dedicated_agent_config
	prepare_compose_env
	start_services
	verify_public_egress
	verify_direct_egress_blocked
	verify_private_targets_blocked
	verify_container_isolation
	start_pi
	printf '\nReal open-mission sandbox started.\n'
	printf '  project: %s\n' "$project"
	printf '  mail integration: none\n'
	printf '  host browser/profile: not mounted\n'
	printf '\nAttach: npm run sandbox:mission:local -- attach\n'
	printf 'Run mission: npm run sandbox:mission:local -- run\n'
}

repair_proxy() {
	prepare_compose_env
	verify_service agent >/dev/null
	compose up -d --build proxy
	wait_for_proxy
	verify_public_egress
	printf 'Repaired the owned egress proxy without restarting the agent or TUI.\n'
}

attach() {
	prepare_compose_env
	require_live_tmux
	exec docker compose --project-name "$project" --file "$compose_file" exec agent tmux attach-session -t "$tmux_session"
}

run_mission() {
	prepare_compose_env
	require_live_tmux
	compose exec -T agent tmux send-keys -t "$tmux_session:0.0" -l "/agi"
	compose exec -T agent tmux send-keys -t "$tmux_session:0.0" Enter
	local output
	for _ in {1..20}; do
		output="$(compose exec -T agent tmux capture-pane -p -S -100 -t "$tmux_session:0.0")"
		[[ "$output" == *"AGI mode enabled"* ]] && break
		sleep 0.5
	done
	[[ "$output" == *"AGI mode enabled"* ]] || die "AGI mode did not become ready"
	compose exec -T agent tmux load-buffer -b open-mission /workspace/mission.txt
	compose exec -T agent tmux paste-buffer -b open-mission -t "$tmux_session:0.0"
	compose exec -T agent tmux send-keys -t "$tmux_session:0.0" Enter
	printf 'Real mission submitted. It remains active until an actual reply is read.\n'
	printf 'Attach with: npm run sandbox:mission:local -- attach\n'
}

nudge() {
	prepare_compose_env
	require_live_tmux
	local current message
	current="$(goal_status | tr -d '\r' | head -n 1)"
	[[ "$current" == "active" ]] || die "refusing to nudge a non-active goal: ${current:-missing}"
	message="${1:-Continue the existing active goal from durable state. Infrastructure connectivity has been restored.}"
	compose exec -T agent tmux send-keys -t "$tmux_session:0.0" -l "$message"
	compose exec -T agent tmux send-keys -t "$tmux_session:0.0" Enter
	printf 'Nudged the active TUI without replacing its goal or mission prompt.\n'
}

recover_session() {
	prepare_compose_env
	require_live_tmux
	local current message
	current="$(goal_status | tr -d '\r' | head -n 1)"
	[[ "$current" == "active" ]] || die "refusing to recover a non-active goal: ${current:-missing}"
	message="${1:-Continue the existing active goal from durable state. Inspect the retained workspace and evidence before acting. Do not repeat already completed consequential actions.}"
	compose exec -T agent tmux send-keys -t "$tmux_session:0.0" -l "/new"
	compose exec -T agent tmux send-keys -t "$tmux_session:0.0" Enter
	sleep 1
	compose exec -T agent tmux send-keys -t "$tmux_session:0.0" -l "/agi"
	compose exec -T agent tmux send-keys -t "$tmux_session:0.0" Enter
	sleep 1
	compose exec -T agent tmux send-keys -t "$tmux_session:0.0" -l "$message"
	compose exec -T agent tmux send-keys -t "$tmux_session:0.0" Enter
	printf 'Started a fresh Pi conversation while preserving the active durable goal and workspace.\n'
}

goal_status() {
	compose exec -T agent sh -lc \
		"test -f /workspace/.pi/agi/goal.md && sed -n 's/^status:[[:space:]]*//p' /workspace/.pi/agi/goal.md | head -n 1" 2>/dev/null || true
}

status() {
	prepare_compose_env
	local agent_id proxy_id
	agent_id="$(verify_service agent)"
	proxy_id="$(verify_service proxy)"
	printf 'project: %s\nagent:   %s\nproxy:   %s\n' "$project" "$agent_id" "$proxy_id"
	if compose exec -T agent tmux has-session -t "$tmux_session" 2>/dev/null; then
		verify_tmux
		local pane_pid pane_command pane_path pane_dead pane_exit
		pane_pid="$(compose exec -T agent tmux display-message -p -t "$tmux_session:0.0" '#{pane_pid}')"
		pane_command="$(compose exec -T agent tmux display-message -p -t "$tmux_session:0.0" '#{pane_current_command}')"
		pane_path="$(compose exec -T agent tmux display-message -p -t "$tmux_session:0.0" '#{pane_current_path}')"
		pane_dead="$(compose exec -T agent tmux display-message -p -t "$tmux_session:0.0" '#{pane_dead}')"
		pane_exit="$(compose exec -T agent tmux display-message -p -t "$tmux_session:0.0" '#{pane_dead_status}')"
		printf 'pane pid: %s\ncommand:  %s\ncwd:      %s\ndead:     %s\nexit:     %s\n' \
			"$pane_pid" "$pane_command" "$pane_path" "$pane_dead" "${pane_exit:-n/a}"
		if [[ "$pane_dead" != "1" ]]; then
			compose exec -T agent ps -o pid,ppid,user,stat,lstart,cmd -p "$pane_pid"
		fi
	fi
	printf 'goal:     %s\n' "$(goal_status | tr -d '\r' | head -n 1)"
	compose exec -T agent sh -lc '
		wake_count="$(find /workspace/.pi/agi/.runtime/wake -type f 2>/dev/null | wc -l)"
		# A timed wake is consumed before the resulting turn arms its successor. A
		# routine check can take several seconds, so retry a zero for at most 2.5s.
		attempt=0
		while [ "$wake_count" -eq 0 ] && [ "$attempt" -lt 5 ]; do
			sleep 0.5
			wake_count="$(find /workspace/.pi/agi/.runtime/wake -type f 2>/dev/null | wc -l)"
			attempt=$((attempt + 1))
		done
		if [ "$wake_count" -eq 0 ]; then
			printf "wake files: 0 (none armed at snapshot; active turn may be settling)\n"
		else
			printf "wake files: %s\n" "$wake_count"
		fi
		printf "run files:  %s\n" "$(find /workspace/.pi/agi/runs -type f 2>/dev/null | wc -l)"
		printf "sessions:   %s\n" "$(find /workspace/sessions -type f 2>/dev/null | wc -l)"
	'
}

monitor() {
	prepare_compose_env
	local interval="${1:-30}" current
	[[ "$interval" =~ ^[1-9][0-9]*$ ]] || die "monitor interval must be a positive integer"
	while true; do
		printf '\n[%s]\n' "$(date -u +%Y-%m-%dT%H:%M:%SZ)"
		status
		current="$(goal_status | tr -d '\r' | head -n 1)"
		if [[ "$current" == "met" ]]; then
			printf 'Goal reports met. Inspect the trace and real reply evidence before accepting completion.\n'
			return
		fi
		sleep "$interval"
	done
}

capture() {
	prepare_compose_env
	verify_tmux
	local output="${1:-}"
	if [[ -z "$output" ]]; then
		compose exec -T agent tmux capture-pane -p -S - -t "$tmux_session:0.0"
		return
	fi
	local parent
	parent="$(dirname "$output")"
	[[ -d "$parent" ]] || die "output directory does not exist: $parent"
	compose exec -T agent tmux capture-pane -p -S - -t "$tmux_session:0.0" >"$output"
	printf 'Saved TUI scrollback to %s\n' "$output"
}

export_workspace() {
	prepare_compose_env
	local id destination
	id="$(verify_service agent)"
	destination="${1:-/tmp/pi-agi-mission-export-$UID}"
	mkdir -p "$destination"
	destination="$(cd "$destination" && pwd -P)"
	docker cp "$id:/workspace/." "$destination/"
	printf 'Exported sandbox workspace to %s\n' "$destination"
}

stop() {
	prepare_compose_env
	verify_service proxy >/dev/null
	verify_service agent >/dev/null
	if compose exec -T agent tmux has-session -t "$tmux_session" 2>/dev/null; then
		verify_tmux
		local pane_pid pane_command pane_path pane_owner pane_parent
		pane_pid="$(compose exec -T agent tmux display-message -p -t "$tmux_session:0.0" '#{pane_pid}')"
		pane_command="$(compose exec -T agent tmux display-message -p -t "$tmux_session:0.0" '#{pane_current_command}')"
		pane_path="$(compose exec -T agent tmux display-message -p -t "$tmux_session:0.0" '#{pane_current_path}')"
		pane_owner="$(compose exec -T agent ps -o user= -p "$pane_pid" | tr -d ' ')"
		pane_parent="$(compose exec -T agent ps -o ppid= -p "$pane_pid" | tr -d ' ')"
		printf 'Stopping owned TUI: pid=%s ppid=%s owner=%s command=%s cwd=%s\n' \
			"$pane_pid" "$pane_parent" "$pane_owner" "$pane_command" "$pane_path"
		compose exec -T agent ps -o pid,ppid,user,stat,lstart,cmd -p "$pane_pid"
		compose exec -T agent tmux send-keys -t "$tmux_session:0.0" C-d || true
		sleep 0.5
	fi
	compose down --remove-orphans
	printf 'Stopped exact Compose project %s. Workspace volume was retained.\n' "$project"
}

preflight() {
	local original_project="$project"
	project="pi-agi-mission-preflight-$UID-$$"
	configured_agent_dir=""
	prepare_compose_env
	trap 'docker compose --project-name "$project" --file "$compose_file" down --volumes --rmi local --remove-orphans >/dev/null 2>&1 || true' EXIT
	start_services
	verify_public_egress
	verify_direct_egress_blocked
	verify_private_targets_blocked
	verify_install_capability
	verify_container_isolation
	printf 'Generic installable-tool internet sandbox preflight: PASS\n'
	docker compose --project-name "$project" --file "$compose_file" down --volumes --rmi local --remove-orphans >/dev/null
	trap - EXIT
	project="$original_project"
	printf 'Removed exact preflight containers, networks, volume, and local images.\n'
}

require_command docker
require_command jq
require_command curl
validate_project

action="${1:-help}"
if [[ $# -gt 0 ]]; then shift; fi

case "$action" in
	preflight) preflight ;;
	start) start ;;
	attach) attach ;;
	run) run_mission ;;
	nudge) nudge "${1:-}" ;;
	recover-session) recover_session "${1:-}" ;;
	status) status ;;
	monitor) monitor "${1:-}" ;;
	repair-proxy) repair_proxy ;;
	capture) capture "${1:-}" ;;
	export) export_workspace "${1:-}" ;;
	stop) stop ;;
	help | -h | --help) usage ;;
	*) usage; die "unknown command: $action" ;;
esac