#!/bin/sh
set -eu
umask 077

usage() {
  echo "usage: $0 <service> <repo> <main> <image> <cluster> <arch> <work-dir> <mode> <context-json>" >&2
  exit 2
}

[ "$#" -eq 9 ] || usage

service=$1
repo=$2
main_path=$3
image=$4
cluster=$5
arch=$6
work_dir=$7
mode=$8
context_json=$9

repo=$(cd "$repo" && pwd)
script_dir=$(cd "$(dirname "$0")" && pwd)
mkdir -p "$work_dir"
work_dir=$(cd "$work_dir" && pwd)
source_manifest_tool=$script_dir/source-manifest.py
source_lock_tool=$script_dir/remote-source-lock.sh
source_manifest_metadata=$work_dir/source-manifest.json
source_archive=$work_dir/source.tar
source_archive_compressed=$work_dir/source.tar.gz
source_parts_dir=$work_dir/source-parts
GOPRIVATE=${GOPRIVATE:-github.com/tidbcloud/*}
GONOSUMDB=${GONOSUMDB:-github.com/tidbcloud/*}
export GOPRIVATE GONOSUMDB

eval "$(
  python3 - "$context_json" <<'PY'
import json
import pathlib
import shlex
import sys

ctx = json.loads(pathlib.Path(sys.argv[1]).read_text())
remote = (ctx.get("host") or {}).get("remote") or {}
github_app_client_id = ctx.get("github_app_client_id")
if not github_app_client_id:
    raise SystemExit("CP9_SESSION_CONTEXT.github_app_client_id is required for remote image builds")

def required(name):
    value = remote.get(name)
    if not value:
        raise SystemExit(f"CP9_SESSION_CONTEXT.host.remote.{name} is required for remote image builds")
    return str(value)

ssh_host = remote.get("public_dns") or remote.get("public_ip")
if not ssh_host:
    raise SystemExit("CP9_SESSION_CONTEXT.host.remote.public_dns or public_ip is required for remote image builds")

ssh_transport = str(remote.get("ssh_transport") or "direct")
if ssh_transport not in ("direct", "ssm"):
    raise SystemExit(f"unsupported CP9 remote SSH transport: {ssh_transport}")
ssh_proxy_command = ""
if ssh_transport == "ssm":
    ssh_proxy_command = shlex.join([
        "aws", "ssm", "start-session",
        "--profile", required("profile"),
        "--region", required("region"),
        "--target", required("instance_id"),
        "--document-name", "AWS-StartSSHSession",
        "--parameters", "portNumber=%p",
    ])

values = {
    "ssh_user": required("ssh_user"),
    "ssh_key": required("ssh_private_key_path"),
    "ssh_host": str(ssh_host),
    "ssh_proxy_command": ssh_proxy_command,
    "remote_state_dir": required("remote_state_dir"),
    "github_app_client_id": str(github_app_client_id),
}
for key, value in values.items():
    print(f"{key}={shlex.quote(value)}")
PY
)"

remote_root=$remote_state_dir/premium-build/$service
python3 "$source_manifest_tool" \
  --repo "$repo" \
  --metadata "$source_manifest_metadata" \
  --archive "$source_archive"
eval "$(
  python3 - "$source_manifest_metadata" <<'PY'
import json
import shlex
import sys
metadata = json.load(open(sys.argv[1], encoding="utf-8"))
values = {
    "remote_repo_rel": metadata["repo_relative"],
    "source_key": metadata["source_key"],
    "source_manifest_fingerprint": metadata["fingerprint"],
}
for key, value in values.items():
    print(f"{key}={shlex.quote(str(value))}")
PY
)"
source_stamp=$source_manifest_fingerprint
CP9_REMOTE_BUILD_STAMP_CONTRACT=cp9.remote-build/v1
stamp_value=$CP9_REMOTE_BUILD_STAMP_CONTRACT:$source_stamp
remote_shared_root=$remote_state_dir/premium-build
remote_source_parent=$remote_shared_root/source
remote_source_root=$remote_source_parent/$source_key
remote_source_lock=$remote_shared_root/source-locks/$source_key.lock
remote_source_lock_timeout=600
remote_source_lock_lease_seconds=300
remote_source_lock_heartbeat_seconds=30
remote_source_previous=$remote_source_root.previous
if [ "$remote_repo_rel" = "." ]; then
  remote_repo=$remote_source_root
else
  remote_repo=$remote_source_root/$remote_repo_rel
fi
remote_scripts=$remote_shared_root/scripts
remote_source_lock_helper=$remote_scripts/remote-source-lock.sh
remote_github_user_helper=$remote_scripts/github-user-module-download.sh
local_github_user_token_helper=$script_dir/github-app-user-token.sh
remote_work=$remote_root/work
ssh_target=$ssh_user@$ssh_host
remote_source_lock_heartbeat_pid=
remote_source_lock_acquired=
remote_source_lock_lost_file=$work_dir/remote-source-lock-lost-$service
remote_upload_parent=$remote_shared_root/source-uploads/$source_key
remote_upload_dir=
remote_source_lock_token=$(python3 - <<'PY'
import secrets

print(secrets.token_hex(16))
PY
)

ssh_common() {
  set -- \
    -o BatchMode=yes \
    -o StrictHostKeyChecking=no \
    -o UserKnownHostsFile=/dev/null \
    -o LogLevel=ERROR \
    -o ServerAliveInterval=30 \
    -o ServerAliveCountMax=3 \
    -o ControlMaster=no \
    -o ControlPath=none \
    -i "$ssh_key" \
    "$@"
  if [ -n "$ssh_proxy_command" ]; then
    set -- -o "$ssh_proxy_command" "$@"
  fi
  ssh "$@"
}

ssh_base() {
  ssh_common "$ssh_target" "$@"
}

remote_quote() {
  python3 - "$1" <<'PY'
import shlex
import sys

print(shlex.quote(sys.argv[1]))
PY
}

ssh_shell() {
  ssh_base "sh -lc $(remote_quote "$1")"
}

stop_remote_source_lock_heartbeat() {
  if [ -n "$remote_source_lock_heartbeat_pid" ]; then
    kill "$remote_source_lock_heartbeat_pid" >/dev/null 2>&1 || true
    wait "$remote_source_lock_heartbeat_pid" >/dev/null 2>&1 || true
    remote_source_lock_heartbeat_pid=
  fi
}

release_remote_source_lock() {
  stop_remote_source_lock_heartbeat
  if [ -n "$remote_source_lock_acquired" ]; then
    ssh_shell "$(remote_quote "$remote_source_lock_helper") release $(remote_quote "$remote_source_lock") $(remote_quote "$remote_source_lock_token")" || true
    remote_source_lock_acquired=
  fi
}

heartbeat_remote_source_lock() {
  while sleep "$remote_source_lock_heartbeat_seconds"; do
    if ! ssh_shell "$(remote_quote "$remote_source_lock_helper") heartbeat $(remote_quote "$remote_source_lock") $(remote_quote "$remote_source_lock_token") $(remote_quote "$remote_source_lock_lease_seconds")"; then
      : >"$remote_source_lock_lost_file"
      return 1
    fi
  done
}

cleanup_remote_build() {
  if [ -n "$remote_upload_dir" ]; then
    ssh_shell "rm -rf $(remote_quote "$remote_upload_dir")" || true
    remote_upload_dir=
  fi
  release_remote_source_lock
  rm -rf "$source_parts_dir" "$source_archive_compressed" "$remote_source_lock_lost_file"
}

. "$script_dir/remote-build-diagnostics.sh"
. "$script_dir/remote-build-driver.sh"
trap cleanup_remote_build EXIT HUP INT TERM

printf '%s\n' "[cp9 persona build] $service: sync source to remote $ssh_host"
ensure_remote_git
prepare_source="mkdir -p $(remote_quote "$remote_source_parent") $(remote_quote "$remote_shared_root/source-locks") $(remote_quote "$remote_scripts") $(remote_quote "$remote_work")"
ssh_shell "$prepare_source"
remote_source_lock_helper_tmp=$remote_source_lock_helper.tmp.$remote_source_lock_token
ssh_base "sh -lc $(remote_quote "cat > $(remote_quote "$remote_source_lock_helper_tmp") && chmod 755 $(remote_quote "$remote_source_lock_helper_tmp") && mv $(remote_quote "$remote_source_lock_helper_tmp") $(remote_quote "$remote_source_lock_helper")")" <"$source_lock_tool"

source_ready=$remote_source_root/.cp9-source-ready
source_stamp_file=$remote_source_root/.cp9-source-stamp
printf '%s\n' "[cp9 persona build] $service: wait for remote source sync lock $source_key"
remote_source_lock_owner_value="$service pid=$$ source_key=$source_key started=$(date -u +%Y-%m-%dT%H:%M:%SZ)"
rm -f "$remote_source_lock_lost_file"
ssh_shell "$(remote_quote "$remote_source_lock_helper") acquire $(remote_quote "$remote_source_lock") $(remote_quote "$remote_source_lock_token") $(remote_quote "$remote_source_lock_lease_seconds") $(remote_quote "$remote_source_lock_timeout") $(remote_quote "$remote_source_lock_owner_value")"
remote_source_lock_acquired=$remote_source_lock_token
heartbeat_remote_source_lock &
remote_source_lock_heartbeat_pid=$!
ssh_shell "if [ ! -e $(remote_quote "$remote_source_root") ] && [ -e $(remote_quote "$remote_source_previous") ]; then mv $(remote_quote "$remote_source_previous") $(remote_quote "$remote_source_root"); fi && rm -rf $(remote_quote "$remote_source_previous") $(remote_quote "$remote_upload_parent") && umask 077 && mkdir -p $(remote_quote "$remote_upload_parent")"
remote_source_stamp=$(ssh_shell "cat $(remote_quote "$source_stamp_file") 2>/dev/null || true")
if ssh_shell "test -f $(remote_quote "$source_ready")" && [ "$remote_source_stamp" = "$stamp_value" ]; then
  printf '%s\n' "[cp9 persona build] $service: reuse remote source on $ssh_host"
else
  rm -rf "$source_parts_dir" "$source_archive_compressed"
  mkdir -p "$source_parts_dir"
  gzip -c "$source_archive" >"$source_archive_compressed"
  if command -v sha256sum >/dev/null 2>&1; then
    source_archive_sha256=$(sha256sum "$source_archive_compressed" | awk '{print $1}')
  else
    source_archive_sha256=$(shasum -a 256 "$source_archive_compressed" | awk '{print $1}')
  fi
  split -b 1m -a 4 "$source_archive_compressed" "$source_parts_dir/part-"
  remote_upload_dir=$remote_upload_parent/$remote_source_lock_token
  remote_archive=$remote_upload_dir/source.tar.gz
  remote_extract_dir=$remote_upload_dir/source
  ssh_shell "umask 077 && mkdir -p $(remote_quote "$remote_upload_dir") $(remote_quote "$remote_extract_dir")"
  upload_pids=
  upload_count=0
  upload_failed=
  for source_part in "$source_parts_dir"/part-*; do
    remote_part=$remote_upload_dir/$(basename "$source_part")
    ssh_base "umask 077; cat > $(remote_quote "$remote_part")" <"$source_part" &
    upload_pids="$upload_pids $!"
    upload_count=$((upload_count + 1))
    if [ "$upload_count" -eq 4 ]; then
      for upload_pid in $upload_pids; do
        wait "$upload_pid" || upload_failed=true
      done
      [ -z "$upload_failed" ] || exit 1
      upload_pids=
      upload_count=0
    fi
  done
  for upload_pid in $upload_pids; do
    wait "$upload_pid" || upload_failed=true
  done
  [ -z "$upload_failed" ] || exit 1
  ssh_shell "umask 077 && cat $(remote_quote "$remote_upload_dir")/part-* > $(remote_quote "$remote_archive") && printf '%s  %s\n' $(remote_quote "$source_archive_sha256") $(remote_quote "$remote_archive") | sha256sum -c - && tar -C $(remote_quote "$remote_extract_dir") -xzf $(remote_quote "$remote_archive")"
  stop_remote_source_lock_heartbeat
  if [ -e "$remote_source_lock_lost_file" ]; then
    echo "remote source lock heartbeat failed; refusing to publish source" >&2
    exit 1
  fi
  ssh_shell "$(remote_quote "$remote_source_lock_helper") heartbeat $(remote_quote "$remote_source_lock") $(remote_quote "$remote_source_lock_token") $(remote_quote "$remote_source_lock_lease_seconds") && $(remote_quote "$remote_source_lock_helper") check $(remote_quote "$remote_source_lock") $(remote_quote "$remote_source_lock_token") && umask 077 && rm -rf $(remote_quote "$remote_source_previous") && if [ -e $(remote_quote "$remote_source_root") ]; then mv $(remote_quote "$remote_source_root") $(remote_quote "$remote_source_previous"); fi && if mv $(remote_quote "$remote_extract_dir") $(remote_quote "$remote_source_root"); then printf '%s\n' $(remote_quote "$stamp_value") > $(remote_quote "$source_stamp_file") && touch $(remote_quote "$source_ready") && rm -rf $(remote_quote "$remote_source_previous") $(remote_quote "$remote_upload_dir"); else status=\$?; mv $(remote_quote "$remote_source_previous") $(remote_quote "$remote_source_root") 2>/dev/null || true; exit \"\$status\"; fi"
  remote_upload_dir=
fi
release_remote_source_lock
tar -C "$script_dir/.." -cf - build-local-image.sh | ssh_base "tar -C $(remote_quote "$remote_scripts") -xf -"
ssh_base "umask 077; cat > $(remote_quote "$remote_github_user_helper")" <"$script_dir/github-user-module-download.sh"
ssh_shell "chmod 700 $(remote_quote "$remote_scripts/build-local-image.sh") $(remote_quote "$remote_github_user_helper")"

local_mod_root=$repo
if [ "$mode" != "account" ]; then
  main_dir=$(dirname "$main_path")
  if [ "$main_dir" != "." ]; then
    current=$repo/$main_dir
    while [ "$current" != "$repo" ] && [ "$current" != "/" ]; do
      if [ -f "$current/go.mod" ]; then
        local_mod_root=$current
        break
      fi
      current=$(dirname "$current")
    done
  fi
  if [ ! -f "$local_mod_root/go.mod" ]; then
    local_mod_root=$repo
  fi
fi
has_local_go_mod=false
if [ -f "$local_mod_root/go.mod" ]; then
  has_local_go_mod=true
fi

remote_env="GOPRIVATE=$(remote_quote "${GOPRIVATE:-github.com/tidbcloud/*}") GONOSUMDB=$(remote_quote "${GONOSUMDB:-github.com/tidbcloud/*}") GONOPROXY=$(remote_quote "github.com/tidbcloud/*") GOPROXY=$(remote_quote "https://proxy.golang.org,direct") CP9_PREMIUM_BUILD_CACHE_DIR=$(remote_quote "$remote_shared_root/cache")"
if [ "$has_local_go_mod" = "true" ]; then
  if [ "$local_mod_root" = "$repo" ]; then
    remote_mod_root=$remote_repo
  else
    remote_mod_root=$remote_repo/${local_mod_root#"$repo"/}
  fi
  printf '%s\n' "[cp9 persona build] $service: download private Go modules on EC2 with current GitHub user"
  if ! "$local_github_user_token_helper" "$github_app_client_id" |
    ssh_shell "$(remote_quote "$remote_github_user_helper") $(remote_quote "$remote_mod_root")"; then
    exit 1
  fi
else
  printf '%s\n' "[cp9 persona build] $service: skip Go module download; $local_mod_root has no go.mod"
fi

remote_cmd="rm -f $(remote_quote "$remote_work/image.stamp") && cd $(remote_quote "$remote_scripts") && $remote_env ./build-local-image.sh $(remote_quote "$service") $(remote_quote "$remote_repo") $(remote_quote "$main_path") $(remote_quote "$image") $(remote_quote "$cluster") $(remote_quote "$arch") $(remote_quote "$remote_work") $(remote_quote "$mode")"
printf '%s\n' "[cp9 persona build] $service: remote docker build/import on $ssh_host"
run_remote_build

stamp_tmp=$work_dir/image.stamp.tmp.$$
if ssh_base "cat $(remote_quote "$remote_work/image.stamp")" >"$stamp_tmp"; then
  mv "$stamp_tmp" "$work_dir/image.stamp"
  printf '%s\n' "[cp9 persona build] $service: copied remote image rollout stamp $work_dir/image.stamp"
else
  rm -f "$stamp_tmp"
  echo "remote image build did not produce $remote_work/image.stamp" >&2
  exit 1
fi