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

[[ $# -eq 3 ]] || {
  echo "usage: test-native-final-resident-candidate.sh <candidate-image> <recovery-token-file> <receipt-output>" >&2
  exit 64
}

candidate=$(readlink -f "$1")
token_file=$(readlink -f "$2")
receipt=$(readlink -m "$3")
source_root=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
native_root=${source_root}/native/x86_64-bios
test_root=$(mktemp -d)
peer_pid=
cleanup() {
  if [[ -n ${peer_pid} ]]; then
    kill "${peer_pid}" 2>/dev/null || true
  fi
  if [[ ${GHOS_KEEP_TEST_ROOT:-0} != 1 ]]; then
    rm -rf "${test_root}"
  else
    echo "GHOS_TEST_ROOT=${test_root}" >&2
  fi
}
trap cleanup EXIT

[[ $(stat -c %s "${candidate}") -eq 14848 ]]
[[ $(tr -d '\r\n ' <"${token_file}" | wc -c) -eq 32 ]]
python3 - "${token_file}" "${test_root}/wrong-token.hex" <<'PY'
import pathlib
import sys

token = bytearray.fromhex(pathlib.Path(sys.argv[1]).read_text().strip())
token[0] ^= 0x01
pathlib.Path(sys.argv[2]).write_text(token.hex() + "\n", encoding="ascii")
PY
truncate -s 700M "${test_root}/disk.img"
nasm -f bin -dCANDIDATE_LBA=105 -dPROOF_LBA=134 \
  "${native_root}/physical-test-mbr.asm" -o "${test_root}/mbr.bin"
dd if="${test_root}/mbr.bin" of="${test_root}/disk.img" \
  bs=512 seek=0 conv=notrunc status=none
dd if="${candidate}" of="${test_root}/disk.img" \
  bs=512 seek=105 conv=notrunc status=none

peer_port=$((33000 + ${BASHPID:-$$} % 1000))
qemu_port=$((peer_port + 1))
python3 "${source_root}/scripts/qemu-native-net-peer.py" \
  --listen-port "${peer_port}" \
  --qemu-port "${qemu_port}" \
  --receipt "${test_root}/peer.hldp" \
  --final-resident \
  --recovery-token-file "${token_file}" >"${test_root}/peer.log" 2>&1 &
peer_pid=$!
set +e
timeout 40 qemu-system-x86_64 \
  -machine pc,accel=tcg \
  -m 64M \
  -drive "if=none,id=ghboot,format=raw,file=${test_root}/disk.img" \
  -device virtio-blk-pci,drive=ghboot,disable-modern=on,bootindex=0 \
  -netdev "dgram,id=ghnet,local.type=inet,local.host=127.0.0.1,local.port=${qemu_port},remote.type=inet,remote.host=127.0.0.1,remote.port=${peer_port}" \
  -device virtio-net-pci,netdev=ghnet,disable-modern=on,mac=52:54:00:26:71:98 \
  -display none \
  -monitor none \
  -serial "file:${test_root}/serial.log" \
  -device isa-debug-exit,iobase=0xf4,iosize=0x04
qemu_status=$?
set -e
[[ ${qemu_status} -eq 33 ]]
wait "${peer_pid}"
peer_pid=

grep -q '^final_resident_control: true$' "${test_root}/peer.hldp"
grep -q '^resident_login_reply_count: 3$' "${test_root}/peer.hldp"
grep -q '^recovery_reply_verified: true$' "${test_root}/peer.hldp"
grep -q '^recovery_selected_by_native: false$' "${test_root}/peer.hldp"
grep -q '^GHOS_NATIVE_RECOVERY_BEACON=WRITE_READ_VERIFIED' "${test_root}/serial.log"
grep -q '^GHOS_DISK_PROOF_OBSERVED_AFTER_RESET=LBA134' "${test_root}/serial.log"

python3 - "${test_root}/disk.img" <<'PY'
import pathlib
import sys

with pathlib.Path(sys.argv[1]).open("rb") as disk:
    def sector(lba: int, count: int = 1) -> bytes:
        disk.seek(lba * 512)
        return disk.read(count * 512)

    proof = sector(134)
    assert proof[0] == 0xA7
    assert proof[1:].startswith(b"GHOS_NATIVE_LONG64_DISK_PROOF\0")
    assert proof[99:102] == bytes([1, 1, 1])
    assert proof[102:105] == bytes([1, 1, 1])
    assert proof[105] == 6
    assert sector(139, 2).startswith(b"# GRUB Environment Block\n")
    assert sector(141).startswith(b"GHOS_GHCIP_INDEX_V1\n")
PY

truncate -s 700M "${test_root}/wrong-capability-disk.img"
dd if="${test_root}/mbr.bin" of="${test_root}/wrong-capability-disk.img" \
  bs=512 seek=0 conv=notrunc status=none
dd if="${candidate}" of="${test_root}/wrong-capability-disk.img" \
  bs=512 seek=105 conv=notrunc status=none
wrong_peer_port=$((peer_port + 2))
wrong_qemu_port=$((qemu_port + 2))
python3 "${source_root}/scripts/qemu-native-net-peer.py" \
  --listen-port "${wrong_peer_port}" \
  --qemu-port "${wrong_qemu_port}" \
  --receipt "${test_root}/wrong-peer.hldp" \
  --final-resident \
  --recovery-token-file "${test_root}/wrong-token.hex" \
  >"${test_root}/wrong-peer.log" 2>&1 &
peer_pid=$!
set +e
timeout 5 qemu-system-x86_64 \
  -machine pc,accel=tcg \
  -m 64M \
  -drive "if=none,id=ghboot,format=raw,file=${test_root}/wrong-capability-disk.img" \
  -device virtio-blk-pci,drive=ghboot,disable-modern=on,bootindex=0 \
  -netdev "dgram,id=ghnet,local.type=inet,local.host=127.0.0.1,local.port=${wrong_qemu_port},remote.type=inet,remote.host=127.0.0.1,remote.port=${wrong_peer_port}" \
  -device virtio-net-pci,netdev=ghnet,disable-modern=on,mac=52:54:00:26:71:98 \
  -display none \
  -monitor none \
  -serial "file:${test_root}/wrong-capability-serial.log" \
  -device isa-debug-exit,iobase=0xf4,iosize=0x04
wrong_qemu_status=$?
set -e
[[ ${wrong_qemu_status} -eq 124 ]]
wait "${peer_pid}"
peer_pid=
python3 - "${test_root}/wrong-capability-disk.img" <<'PY'
import pathlib
import sys

with pathlib.Path(sys.argv[1]).open("rb") as disk:
    disk.seek(134 * 512)
    proof = disk.read(512)
    disk.seek(139 * 512)
    recovery = disk.read(1024)
assert proof[0] == 0xA7
assert recovery == bytes(1024)
PY

observed_at=$(date --iso-8601=seconds)
image_sha=$(sha256sum "${candidate}" | awk '{print $1}')
cat >"${receipt}" <<EOF
schema: guanghu.jd-native-final-resident-qemu-test/v1
receipt_id: GH-OS-JD-FD-PRIMARY-001-NATIVE-FINAL-RESIDENT-CONTROL-QEMU-001
status: PASS_100
observed_at: ${observed_at}
node_id: JD-FD-PRIMARY
candidate:
  lba_start: 105
  sector_count: 29
  sha256: ${image_sha}
resident_runtime:
  initial_login_replies: 3
  code_commit_reply_verified: true
  branch_move_reply_verified: true
  sustained_login_replies_before_recovery: 2
recovery_control:
  selected_by_native_runtime: false
  per_deployment_capability_required: true
  protected_relay_capability_verified: true
  wrong_capability_ignored_while_native_remained_resident: true
  recovery_beacon_write_read: PASS_100
boundary:
  qemu_capability: 100
  physical_server_capability: 0
  final_native_residency_proven: false
next_action: PHYSICAL_ONE_TIME_FINAL_RESIDENCY_AND_PROTECTED_RECOVERY_GATE
EOF