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

die() {
  echo "error: $*" >&2
  exit 2
}

usage() {
  cat <<'USAGE'
Usage:
  scripts/prepare/start-generation.sh --manifest PATH [options]

Starts distributed `tipoc prepare --ops generate` processes for every
generation node declared by a TiPOC manifest. The coordinator node is launched
first, then each non-coordinator generation worker is launched with the
coordinator endpoint.

Normal live-launch input:
  --manifest PATH
      The only required flag in the normal case. The manifest supplies the
      generation node ids, default node addresses, coordinator node, SSH user,
      SSH key/password refs, remote paths, and object_store URI context. This is
      sufficient when env refs are resolvable and manifest node addresses are
      SSH reachable from where this script runs.

Optional runtime overrides for execution-environment differences:
  --coordinator-endpoint-host HOST
      Optional private IP or DNS name reachable from every worker; workers use
      http://HOST:PORT. Defaults to the coordinator generation node address
      read from the manifest. Provide this when the manifest coordinator address
      is not the worker-reachable endpoint.
  SSH reachable hosts
      Node ids and default addresses are read from the manifest; do not provide
      the node list again. By default each node uses its manifest address for
      SSH. Use
      --node-hosts 'node-a=host-a,node-b=host-b' when the manifest has private
      worker/VPC addresses but the launcher runs outside the VPC and must SSH
      through public IPs, bastion-reachable names, or alternate DNS names.
  SSH identity
      SSH user and private-key/password refs are read from manifest init fields
      by default. If init.ssh_private_key_ref or scenario.init.ssh_private_key_ref
      uses env:NAME and NAME contains a private-key file path, the script uses
      that key. Use --ssh-user or --ssh-key only as execution-environment
      overrides. Password refs are not accepted as script arguments; configure
      noninteractive SSH outside this script when password-based access is the
      manifest contract. SSH secrets are not written to receipts.
  Remote paths
      Defaults come from manifest init.remote_attempt_root, then
      scenario.init.remote_attempt_root, then the manifest directory:
      <remote-root>/run-profile-manifest.json and <remote-root>/bin/tipoc.
      Use --remote-root, --remote-manifest, or --remote-bin if remote paths
      differ.
  Receipt path
      Optional --receipt PATH. Default:
      <local-attempt-root>/prepare/generation-launch-receipt.json.

S3/object-store context:
  The script reads object_store.uri from the manifest and records only sanitized
  prefix context for operator review. Query strings and fragments are stripped
  before receipt write. Prefix hygiene and cleanup are verified separately before
  launch. Do not pass S3/object-store arguments or credentials to this script.

Options:
  --manifest PATH
  --node-hosts 'node-a=host-a,node-b=host-b,...'
  --ssh-user USER
  --ssh-key PATH
  --remote-root PATH
  --remote-manifest PATH
  --remote-bin PATH
  --coordinator-endpoint-host HOST    default: coordinator manifest address
  --port N                         default: 18080
  --coordinator-wait-seconds N      default: 60
  --receipt PATH
  --dry-run
  --help
USAGE
}

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

is_positive_int() {
  [[ "${1:-}" =~ ^[1-9][0-9]*$ ]]
}

abs_path() {
  local path="$1"
  if [[ -d "$path" ]]; then
    (cd "$path" && pwd)
  else
    local dir base
    dir="$(dirname "$path")"
    base="$(basename "$path")"
    (cd "$dir" && printf '%s/%s\n' "$(pwd)" "$base")
  fi
}

abs_path_allow_missing() {
  local path="$1"
  if [[ "$path" == /* ]]; then
    printf '%s\n' "$path"
  else
    printf '%s/%s\n' "$(pwd)" "$path"
  fi
}

trim_slash() {
  local value="$1"
  value="${value%/}"
  printf '%s\n' "$value"
}

sanitize_uri_context() {
  local value="$1"
  value="${value%%\#*}"
  value="${value%%\?*}"
  printf '%s\n' "$value"
}

parse_node_hosts_json() {
  local list="$1"
  local json='{}'
  local seen=','
  local item node host
  [[ -n "$list" ]] || {
    printf '%s\n' "$json"
    return
  }
  IFS=',' read -ra items <<< "$list"
  for item in "${items[@]}"; do
    [[ -n "$item" ]] || continue
    node="${item%%=*}"
    host="${item#*=}"
    [[ "$node" != "$item" ]] || die "node host override $item must use node-id=host"
    [[ -n "$node" && -n "$host" ]] || die "invalid node host override $item"
    [[ "$seen" != *",$node,"* ]] || die "duplicate --node-hosts entry for $node"
    seen="${seen}${node},"
    json="$(jq -c --arg node "$node" --arg host "$host" '. + {($node): $host}' <<<"$json")"
  done
  printf '%s\n' "$json"
}

shell_single_quote() {
  local value="$1"
  printf "'%s'" "${value//\'/\'\"\'\"\'}"
}

ssh_target() {
  local host="$1"
  if [[ -n "$ssh_user" ]]; then
    printf '%s@%s\n' "$ssh_user" "$host"
  else
    printf '%s\n' "$host"
  fi
}

run_remote_preflight() {
  local host="$1"
  local target
  target="$(ssh_target "$host")"
  ssh "${ssh_opts[@]}" "$target" bash -s -- "$remote_bin" "$remote_manifest" <<'REMOTE'
set -euo pipefail
bin="$1"
manifest="$2"
[[ -x "$bin" ]]
[[ -s "$manifest" ]]
REMOTE
}

run_remote_launch() {
  local host="$1"
  local node_id="$2"
  local role="$3"
  local log_path="$4"
  local pid_path="$5"
  local target endpoint
  target="$(ssh_target "$host")"
  endpoint="http://${coordinator_endpoint_host}:${port}"
  if [[ "$role" == "coordinator" ]]; then
    ssh "${ssh_opts[@]}" "$target" bash -s -- \
      "$remote_root" "$remote_bin" "$remote_manifest" "$node_id" "$port" "$log_path" "$pid_path" <<'REMOTE'
set -euo pipefail
root="$1"
bin="$2"
manifest="$3"
node_id="$4"
port="$5"
log_path="$6"
pid_path="$7"
mkdir -p "$(dirname "$log_path")" "$(dirname "$pid_path")" "$root/prepare"
nohup env \
  TIPOC_PREPARE_GENERATION_NODE_ID="$node_id" \
  TIPOC_PREPARE_GENERATION_COORDINATOR_LISTEN="0.0.0.0:$port" \
  TIPOC_PREPARE_GENERATION_LEDGER_PATH="$root/prepare/generation-ledger.json" \
  "$bin" prepare --manifest "$manifest" --ops generate \
  >"$log_path" 2>&1 &
pid="$!"
printf '%s\n' "$pid" >"$pid_path"
printf '%s\n' "$pid"
REMOTE
  else
    ssh "${ssh_opts[@]}" "$target" bash -s -- \
      "$remote_root" "$remote_bin" "$remote_manifest" "$node_id" "$endpoint" "$log_path" "$pid_path" <<'REMOTE'
set -euo pipefail
root="$1"
bin="$2"
manifest="$3"
node_id="$4"
endpoint="$5"
log_path="$6"
pid_path="$7"
mkdir -p "$(dirname "$log_path")" "$(dirname "$pid_path")" "$root/prepare"
nohup env \
  TIPOC_PREPARE_GENERATION_NODE_ID="$node_id" \
  TIPOC_PREPARE_GENERATION_COORDINATOR_ENDPOINT="$endpoint" \
  "$bin" prepare --manifest "$manifest" --ops generate \
  >"$log_path" 2>&1 &
pid="$!"
printf '%s\n' "$pid" >"$pid_path"
printf '%s\n' "$pid"
REMOTE
  fi
}

wait_remote_coordinator_ready() {
  local host="$1"
  local target
  target="$(ssh_target "$host")"
  ssh "${ssh_opts[@]}" "$target" bash -s -- "$port" "$coordinator_wait_seconds" <<'REMOTE'
set -euo pipefail
port="$1"
deadline="$2"
end=$((SECONDS + deadline))
while (( SECONDS <= end )); do
  if command -v nc >/dev/null 2>&1; then
    nc -z 127.0.0.1 "$port" >/dev/null 2>&1 && exit 0
  elif command -v bash >/dev/null 2>&1; then
    bash -c ":</dev/tcp/127.0.0.1/$port" >/dev/null 2>&1 && exit 0
  fi
  sleep 1
done
echo "coordinator did not listen on 127.0.0.1:$port within ${deadline}s" >&2
exit 1
REMOTE
}

manifest=""
node_hosts_arg=""
ssh_user_override=""
ssh_key_override=""
remote_root_override=""
remote_manifest_override=""
remote_bin_override=""
coordinator_endpoint_host=""
port="18080"
coordinator_wait_seconds="60"
receipt_override=""
dry_run=false

while [[ $# -gt 0 ]]; do
  case "$1" in
    --manifest)
      [[ $# -ge 2 ]] || die "--manifest requires PATH"
      manifest="$2"
      shift 2
      ;;
    --node-hosts)
      [[ $# -ge 2 ]] || die "--node-hosts requires a value"
      node_hosts_arg="$2"
      shift 2
      ;;
    --ssh-user)
      [[ $# -ge 2 ]] || die "--ssh-user requires USER"
      ssh_user_override="$2"
      shift 2
      ;;
    --ssh-key)
      [[ $# -ge 2 ]] || die "--ssh-key requires PATH"
      ssh_key_override="$2"
      shift 2
      ;;
    --remote-root)
      [[ $# -ge 2 ]] || die "--remote-root requires PATH"
      remote_root_override="$2"
      shift 2
      ;;
    --remote-manifest)
      [[ $# -ge 2 ]] || die "--remote-manifest requires PATH"
      remote_manifest_override="$2"
      shift 2
      ;;
    --remote-bin)
      [[ $# -ge 2 ]] || die "--remote-bin requires PATH"
      remote_bin_override="$2"
      shift 2
      ;;
    --coordinator-endpoint-host)
      [[ $# -ge 2 ]] || die "--coordinator-endpoint-host requires HOST"
      coordinator_endpoint_host="$2"
      shift 2
      ;;
    --port)
      [[ $# -ge 2 ]] || die "--port requires N"
      port="$2"
      shift 2
      ;;
    --coordinator-wait-seconds)
      [[ $# -ge 2 ]] || die "--coordinator-wait-seconds requires N"
      coordinator_wait_seconds="$2"
      shift 2
      ;;
    --receipt)
      [[ $# -ge 2 ]] || die "--receipt requires PATH"
      receipt_override="$2"
      shift 2
      ;;
    --dry-run)
      dry_run=true
      shift
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      die "unknown argument: $1"
      ;;
  esac
done

require_command jq

[[ -n "$manifest" ]] || die "--manifest is required"
[[ -s "$manifest" ]] || die "manifest does not exist or is empty: $manifest"
is_positive_int "$port" || die "--port must be a positive integer"
is_positive_int "$coordinator_wait_seconds" || die "--coordinator-wait-seconds must be a positive integer"
manifest="$(abs_path "$manifest")"
manifest_dir="$(dirname "$manifest")"
node_hosts_json="$(parse_node_hosts_json "$node_hosts_arg")"

generation_nodes_json="$(jq -c '
  def normalize_node:
    if type == "object" then
      {
        node_id: ((.node_id // .id // .name // .node // .key // "") | tostring),
        address: ((.address // .host // .hostname // .ip // "") | tostring)
      }
    else empty end;
  def normalize_nodes($nodes):
    if ($nodes | type) == "array" then
      [$nodes[] | normalize_node]
    elif ($nodes | type) == "object" then
      [($nodes | to_entries[] | (if (.value | type) == "object" then (.value + {key: .key}) else {key: .key, address: (.value | tostring)} end) | normalize_node)]
    else [] end;
  (normalize_nodes(.generation.nodes)) as $top
  | (normalize_nodes(.scenario.generation.nodes)) as $scenario
  | (if ($top | length) > 0 then $top else $scenario end)
  | map(select(.node_id != ""))
' "$manifest")"

node_count="$(jq 'length' <<<"$generation_nodes_json")"
(( node_count > 0 )) || die "manifest contains no generation nodes"

duplicate_ids="$(jq -r '
  def ids($nodes):
    if ($nodes | type) == "array" then
      [$nodes[] | select(type == "object") | (.node_id // .id // .name // .node // empty) | tostring]
    elif ($nodes | type) == "object" then
      [$nodes | to_entries[] | (.value.node_id // .value.id // .value.name // .value.node // .key // empty) | tostring]
    else [] end;
  (ids(.generation.nodes)) as $top
  | (ids(.scenario.generation.nodes)) as $scenario
  | (if ($top | length) > 0 then $top else $scenario end)
  | map(select(. != ""))
  | group_by(.)
  | map(select(length > 1) | .[0])
  | join(",")
' "$manifest")"
[[ -z "$duplicate_ids" ]] || die "duplicate generation node ids: $duplicate_ids"

coordinator_node_id="$(jq -r '
  .init.coordinator_node_id
  // .scenario.init.coordinator_node_id
  // empty
' "$manifest")"
if [[ -z "$coordinator_node_id" ]]; then
  coordinator_node_id="$(jq -r '.[0].node_id' <<<"$generation_nodes_json")"
fi
[[ -n "$coordinator_node_id" ]] || die "cannot determine coordinator node id"
jq -e --arg id "$coordinator_node_id" 'any(.[]; .node_id == $id)' <<<"$generation_nodes_json" >/dev/null \
  || die "coordinator node $coordinator_node_id is not in generation nodes"

remote_root="$(jq -r '
  .init.remote_attempt_root
  // .scenario.init.remote_attempt_root
  // empty
' "$manifest")"
[[ -n "$remote_root_override" ]] && remote_root="$remote_root_override"
[[ -n "$remote_root" ]] || remote_root="$manifest_dir"
remote_root="$(trim_slash "$remote_root")"

remote_manifest="${remote_root}/run-profile-manifest.json"
[[ -n "$remote_manifest_override" ]] && remote_manifest="$remote_manifest_override"
remote_bin="${remote_root}/bin/tipoc"
[[ -n "$remote_bin_override" ]] && remote_bin="$remote_bin_override"

ssh_user="$(jq -r '
  .init.ssh_user
  // .scenario.init.ssh_user
  // empty
' "$manifest")"
[[ -n "$ssh_user_override" ]] && ssh_user="$ssh_user_override"

ssh_key="$ssh_key_override"
if [[ -z "$ssh_key" ]]; then
  ssh_key_ref="$(jq -r '
    .init.ssh_private_key_ref
    // .scenario.init.ssh_private_key_ref
    // empty
  ' "$manifest")"
  if [[ "$ssh_key_ref" == env:* ]]; then
    env_name="${ssh_key_ref#env:}"
    ssh_key="${!env_name:-}"
  fi
fi

if [[ -n "$ssh_key" ]]; then
  [[ -f "$ssh_key" ]] || die "SSH key file not found: $ssh_key"
fi

ssh_opts=(-o BatchMode=yes)
if [[ -n "$ssh_key" ]]; then
  ssh_opts+=(-i "$ssh_key")
fi

local_attempt_root="$(jq -r '
  .init.local_attempt_root
  // .init.attempt_root
  // .scenario.init.local_attempt_root
  // .scenario.init.attempt_root
  // .evidence_root
  // empty
' "$manifest")"
[[ -n "$local_attempt_root" ]] || local_attempt_root="$manifest_dir"
if [[ "$local_attempt_root" != /* ]]; then
  local_attempt_root="$(abs_path_allow_missing "$local_attempt_root")"
fi

receipt_path="$receipt_override"
[[ -n "$receipt_path" ]] || receipt_path="${local_attempt_root}/prepare/generation-launch-receipt.json"
if [[ "$receipt_path" != /* ]]; then
  receipt_path="$(abs_path_allow_missing "$receipt_path")"
fi

object_store_uri="$(jq -r '.object_store.uri // .scenario.object_store.uri // empty' "$manifest")"
object_store_uri_sanitized="$(sanitize_uri_context "$object_store_uri")"

plan_nodes_json="$(jq -c \
  --argjson overrides "$node_hosts_json" \
  --arg coordinator "$coordinator_node_id" \
  --arg remote_root "$remote_root" \
  '
  map(
    . as $node
    | ($overrides[$node.node_id] // $node.address) as $host
    | {
        node_id: $node.node_id,
        host: ($host | tostring),
        role: (if $node.node_id == $coordinator then "coordinator" else "worker" end),
        log_path: ($remote_root + "/logs/prepare/" + (if $node.node_id == $coordinator then "coordinator" else $node.node_id end) + ".log"),
        pid_path: ($remote_root + "/logs/prepare/" + $node.node_id + ".pid"),
        status: "planned"
      }
  )
  ' <<<"$generation_nodes_json")"

missing_hosts="$(jq -r 'map(select(.host == "" or .host == "null") | .node_id) | join(",")' <<<"$plan_nodes_json")"
[[ -z "$missing_hosts" ]] || die "missing SSH host for generation node(s): $missing_hosts"

if [[ -z "$coordinator_endpoint_host" ]]; then
  coordinator_endpoint_host="$(jq -r --arg id "$coordinator_node_id" '.[] | select(.node_id == $id) | .address' <<<"$generation_nodes_json")"
fi
[[ -n "$coordinator_endpoint_host" && "$coordinator_endpoint_host" != "null" ]] \
  || die "missing coordinator endpoint host; set --coordinator-endpoint-host"

receipt_base="$(jq -n \
  --arg status "planned" \
  --arg manifest "$manifest" \
  --arg remote_root "$remote_root" \
  --arg remote_manifest "$remote_manifest" \
  --arg remote_bin "$remote_bin" \
  --arg coordinator_node_id "$coordinator_node_id" \
  --arg coordinator_endpoint "http://${coordinator_endpoint_host}:${port}" \
  --arg coordinator_wait_seconds "$coordinator_wait_seconds" \
  --arg object_store_uri "$object_store_uri_sanitized" \
  --argjson nodes "$plan_nodes_json" \
  '{
    schema_version: "tipoc.prepare.generation_launch_receipt.v1",
    status: $status,
    local_manifest: $manifest,
    remote_root: $remote_root,
    remote_manifest: $remote_manifest,
    remote_bin: $remote_bin,
    coordinator_node_id: $coordinator_node_id,
    coordinator_endpoint: $coordinator_endpoint,
    coordinator_wait_seconds: ($coordinator_wait_seconds | tonumber),
    object_store_uri: $object_store_uri,
    s3_prefix_hygiene: "derived_from_manifest_object_store_uri_verify_separately",
    nodes: $nodes
  }')"

mkdir -p "$(dirname "$receipt_path")"

if [[ "$dry_run" == true ]]; then
  jq '.status = "dry_run" | .nodes |= map(.status = "planned_dry_run")' <<<"$receipt_base" >"$receipt_path"
  cat "$receipt_path"
  exit 0
fi

while IFS= read -r row; do
  host="$(jq -r '.host' <<<"$row")"
  run_remote_preflight "$host"
done < <(jq -c '.[] | select(.role == "coordinator")' <<<"$plan_nodes_json")
while IFS= read -r row; do
  host="$(jq -r '.host' <<<"$row")"
  run_remote_preflight "$host"
done < <(jq -c '.[] | select(.role != "coordinator")' <<<"$plan_nodes_json")

launched_nodes='[]'
coordinator_host=""
while IFS= read -r row; do
  node_id="$(jq -r '.node_id' <<<"$row")"
  host="$(jq -r '.host' <<<"$row")"
  role="$(jq -r '.role' <<<"$row")"
  log_path="$(jq -r '.log_path' <<<"$row")"
  pid_path="$(jq -r '.pid_path' <<<"$row")"
  pid="$(run_remote_launch "$host" "$node_id" "$role" "$log_path" "$pid_path")"
  coordinator_host="$host"
  launched_nodes="$(jq -c --argjson row "$row" --arg pid "$pid" '. + [($row + {status: "launched", pid: $pid})]' <<<"$launched_nodes")"
done < <(jq -c '.[] | select(.role == "coordinator")' <<<"$plan_nodes_json")

[[ -n "$coordinator_host" ]] || die "coordinator launch did not produce a host"
wait_remote_coordinator_ready "$coordinator_host"

while IFS= read -r row; do
  node_id="$(jq -r '.node_id' <<<"$row")"
  host="$(jq -r '.host' <<<"$row")"
  role="$(jq -r '.role' <<<"$row")"
  log_path="$(jq -r '.log_path' <<<"$row")"
  pid_path="$(jq -r '.pid_path' <<<"$row")"
  pid="$(run_remote_launch "$host" "$node_id" "$role" "$log_path" "$pid_path")"
  launched_nodes="$(jq -c --argjson row "$row" --arg pid "$pid" '. + [($row + {status: "launched", pid: $pid})]' <<<"$launched_nodes")"
done < <(jq -c '.[] | select(.role != "coordinator")' <<<"$plan_nodes_json")

jq --argjson nodes "$launched_nodes" '.status = "launched" | .nodes = $nodes' <<<"$receipt_base" >"$receipt_path"
cat "$receipt_path"