#!/usr/bin/env bash
# shellcheck disable=SC2140,SC2154

pimox_ssh() {
    local host="$1"
    local user="$2"
    local key_path="$3"

    shift 3
    ssh -i "${key_path}" -o BatchMode=yes -o ConnectTimeout=10 -o StrictHostKeyChecking=accept-new "${user}@${host}" "$@"
}

pimox_guest_ipv4() {
    local guest_json
    local host="$1"
    local user="$2"
    local key_path="$3"
    local vmid="$4"
    local ip_prefix="$5"
    local qm_bin="${6:-${LAB_PIMOX_QM_BIN:-/usr/sbin/qm}}"

    guest_json="$(pimox_ssh "${host}" "${user}" "${key_path}" "sudo '${qm_bin}' guest cmd '${vmid}' network-get-interfaces" 2>/dev/null || true)"
    if [[ -z "${guest_json}" ]]; then
        return 1
    fi

    GUEST_JSON="${guest_json}" python3 - "${ip_prefix}" <<'PY'
import json
import os
import sys

prefix = sys.argv[1]
try:
    interfaces = json.loads(os.environ.get("GUEST_JSON", ""))
except Exception:
    sys.exit(1)

for iface in interfaces or []:
    for address in iface.get("ip-addresses") or []:
        if address.get("ip-address-type") != "ipv4":
            continue
        ip = address.get("ip-address", "")
        if not ip or ip.startswith(("127.", "169.254.")):
            continue
        if prefix and not ip.startswith(prefix):
            continue
        print(ip)
        sys.exit(0)
sys.exit(1)
PY
}

pimox_worker_vm_debug() {
    local host="$1"
    local user="$2"
    local key_path="$3"
    local vmid="$4"
    local qm_bin="$5"
    local script
    local encoded_script

    script="$(cat <<'EOF'
set +e
echo "-- addresses --"
ip -br addr
echo "-- routes --"
ip route
echo "-- units --"
systemctl is-active ssh sshd 2>/dev/null || true
echo "-- listeners --"
ss -ltnp 2>/dev/null | grep :22 || ss -ltn 2>/dev/null | grep :22 || true
echo "-- sshd config test --"
sudo sshd -t 2>&1 || true
echo "-- recent ssh logs --"
journalctl -u ssh -u sshd --no-pager -n 40 2>/dev/null || true
EOF
)"
    encoded_script="$(printf '%s' "${script}" | base64 | tr -d '\n')"
    pimox_ssh "${host}" "${user}" "${key_path}" "set +e
echo 'Pimox VM ${vmid} status:'
sudo '${qm_bin}' status '${vmid}'
echo 'Pimox VM ${vmid} config summary:'
sudo '${qm_bin}' config '${vmid}' | grep -E '^(agent|boot|net0|scsi0|virtio0|sata0|ide0|ide2|efidisk0):' || true
echo 'Pimox VM ${vmid} guest-agent network-get-interfaces:'
sudo '${qm_bin}' guest cmd '${vmid}' network-get-interfaces || true
echo 'Pimox VM ${vmid} guest SSH diagnostics:'
sudo '${qm_bin}' guest exec '${vmid}' -- bash -lc \"printf '%s' '${encoded_script}' | base64 -d | sudo bash\" || true" >&2 || true
}

pimox_worker_guest_agent_recovery_hint() {
    local host="$1"
    local user="$2"
    local vmid="$3"
    local qm_bin="$4"

    cat >&2 <<EOF

QEMU guest agent did not become available for Pimox VM ${vmid}.
Fast checks:
  ssh ${user}@${host} 'sudo ${qm_bin} status ${vmid}; sudo ${qm_bin} config ${vmid}; sudo ${qm_bin} guest cmd ${vmid} network-get-interfaces'

Recovery options:
  1. If the clone is still booting slowly, rerun the Pimox stage with a longer agent wait:
     LAB_RPI_SERVICES_DEPLOY=false LAB_PIMOX_GUEST_AGENT_CONFIG_TIMEOUT_SECONDS=600 ./jeannie up
  2. If the same VM still reports "QEMU guest agent is not running", recreate the worker clone:
     LAB_RPI_SERVICES_DEPLOY=false LAB_PIMOX_WORKER_REPLACE_EXISTING=true ./jeannie up
  3. If a recreated worker still has no guest agent, rebuild the template and worker:
     LAB_RPI_SERVICES_DEPLOY=false LAB_PIMOX_TEMPLATE_REPLACE_EXISTING=true LAB_PIMOX_WORKER_REPLACE_EXISTING=true ./jeannie up
EOF
}

pimox_guest_exec_exitcode() {
    local guest_exec_json="$1"

    GUEST_EXEC_JSON="${guest_exec_json}" python3 - <<'PY'
import json
import os
import sys

try:
    document = json.loads(os.environ.get("GUEST_EXEC_JSON", ""))
except Exception:
    sys.exit(2)

try:
    print(int(document.get("exitcode", 1)))
except Exception:
    sys.exit(2)
PY
}

pimox_worker_guest_ssh_repair() {
    local host="$1"
    local user="$2"
    local key_path="$3"
    local vmid="$4"
    local qm_bin="$5"
    local script
    local encoded_script
    local repair_output
    local repair_exitcode

    script="$(cat <<'EOF'
set -eu
ssh_unit=""
if systemctl list-unit-files ssh.service 2>/dev/null | grep -q '^ssh[.]service'; then
  ssh_unit=ssh.service
elif systemctl list-unit-files sshd.service 2>/dev/null | grep -q '^sshd[.]service'; then
  ssh_unit=sshd.service
else
  echo "Neither ssh.service nor sshd.service exists in the guest" >&2
  exit 1
fi

sudo ssh-keygen -A
sudo install -d -m 0755 /run/sshd
sudo mkdir -p /etc/ssh/sshd_config.d
sudo tee /etc/ssh/sshd_config.d/99-homelab-worker-listen.conf >/dev/null <<'SSHD_CONFIG'
Port 22
ListenAddress 0.0.0.0
PubkeyAuthentication yes
PasswordAuthentication no
KbdInteractiveAuthentication no
PermitRootLogin no
UsePAM yes
SSHD_CONFIG

sudo sshd -t
sudo systemctl unmask ssh.service ssh.socket sshd.service sshd.socket >/dev/null 2>&1 || true
sudo systemctl disable --now ssh.socket sshd.socket >/dev/null 2>&1 || true
sudo systemctl enable "$ssh_unit" >/dev/null
sudo systemctl restart "$ssh_unit"
sleep 2
sudo systemctl is-active "$ssh_unit"
if command -v ss >/dev/null 2>&1; then
  ss -ltnp 2>/dev/null || ss -ltn
  ss -ltn | awk '$4 == "0.0.0.0:22" { found = 1 } END { exit found ? 0 : 1 }'
fi
EOF
)"
    encoded_script="$(printf '%s' "${script}" | base64 | tr -d '\n')"
    if ! repair_output="$(pimox_ssh "${host}" "${user}" "${key_path}" "sudo '${qm_bin}' guest exec '${vmid}' -- bash -lc \"printf '%s' '${encoded_script}' | base64 -d | sudo bash\"" 2>&1)"; then
        echo "Could not run guest SSH repair through qemu-guest-agent for VM ${vmid}." >&2
        printf '%s\n' "${repair_output}" | sed 's/^/    guest-ssh-repair: /' >&2
        return 1
    fi

    printf '%s\n' "${repair_output}" | sed 's/^/    guest-ssh-repair: /' >&2
    if repair_exitcode="$(pimox_guest_exec_exitcode "${repair_output}")" && [[ "${repair_exitcode}" == "0" ]]; then
        return 0
    fi
    echo "Guest SSH repair failed inside Pimox VM ${vmid}." >&2
    if [[ -n "${repair_exitcode:-}" ]]; then
        echo "Guest exit code: ${repair_exitcode}" >&2
    fi
    return 1
}

pimox_generated_mac() {
    local vmid="$1"

    printf '02:68:10:%02x:%02x:%02x\n' \
        $(((vmid >> 16) & 255)) \
        $(((vmid >> 8) & 255)) \
        $((vmid & 255))
}

pimox_worker_static_ip() {
    local index="$1"
    local static_ips="$2"
    local current_index=1
    local ip

    static_ips="${static_ips//,/ }"
    for ip in ${static_ips}; do
        if ((current_index == index)); then
            printf '%s\n' "${ip}"
            return 0
        fi
        current_index=$((current_index + 1))
    done

    return 1
}

pimox_worker_net0_config() {
    local vmid="$1"
    local bridge="$2"
    local mac_mode="${LAB_PIMOX_WORKER_MAC_MODE:-deterministic}"
    local mac

    case "${mac_mode}" in
        auto)
            printf 'virtio,bridge=%s\n' "${bridge}"
            ;;
        deterministic)
            mac="$(pimox_generated_mac "${vmid}")"
            printf 'virtio=%s,bridge=%s\n' "${mac}" "${bridge}"
            ;;
        *)
            echo "LAB_PIMOX_WORKER_MAC_MODE must be 'auto' or 'deterministic'." >&2
            return 1
            ;;
    esac
}

pimox_worker_cpu_affinity() {
    local index="$1"
    local affinities="$2"
    local worker_cores="$3"
    local affinity
    local affinity_index=1
    local cpu_count

    for affinity in ${affinities}; do
        if ((affinity_index == index)); then
            if ! cpu_count="$(cpuset_cpu_count "${affinity}")"; then
                echo "Invalid Pimox worker CPU affinity '${affinity}'. Use CPU IDs or ranges, such as 4-5." >&2
                exit 1
            fi
            if ((cpu_count != worker_cores)); then
                echo "Pimox worker index ${index} uses ${worker_cores} cores but affinity '${affinity}' contains ${cpu_count} CPUs." >&2
                exit 1
            fi
            printf '%s\n' "${affinity}"
            return 0
        fi
        affinity_index=$((affinity_index + 1))
    done

    echo "No LAB_PIMOX_WORKER_CPU_AFFINITIES entry exists for Pimox worker index ${index}." >&2
    exit 1
}

pimox_start_vm_with_cpuset() {
    local host="$1"
    local user="$2"
    local key_path="$3"
    local qm_bin="$4"
    local vmid="$5"
    local cpuset="$6"

    if [[ -z "${cpuset}" ]]; then
        pimox_ssh "${host}" "${user}" "${key_path}" "sudo '${qm_bin}' start '${vmid}'"
        return
    fi
    if ! cpuset_cpu_count "${cpuset}" >/dev/null; then
        echo "Invalid Pimox KVM CPU set '${cpuset}'. Use CPU IDs or ranges, such as 4-7." >&2
        exit 1
    fi

    # Pimox 7 daemonizes qm start before KVM is spawned. Launch qm's generated
    # command under taskset so PSCI brings secondary guest CPUs online reliably.
    pimox_ssh "${host}" "${user}" "${key_path}" "set -eu
vmid='${vmid}'
qm_bin='${qm_bin}'
cpuset='${cpuset}'

if ! command -v taskset >/dev/null 2>&1; then
  echo 'taskset is required to start this Pimox VM with a KVM CPU set' >&2
  exit 1
fi
if ! taskset -c \"\$cpuset\" true >/dev/null 2>&1; then
  echo \"Pimox KVM CPU set \$cpuset is not valid on this host\" >&2
  exit 1
fi

start_script=\$(mktemp \"/tmp/homelab-qemu-\${vmid}.XXXXXX\")
cleanup() {
  rm -f \"\$start_script\"
}
trap cleanup 0

sudo \"\$qm_bin\" showcmd \"\$vmid\" --pretty >\"\$start_script\"
if ! awk -v vmid=\"\$vmid\" '
  {
    for (i = 1; i <= NF; i++) {
      if ($i == "/usr/bin/kvm" || $i == "qemu-system-aarch64" || $i ~ /qemu-system-aarch64$/) kvm = 1
      if ($i == "-id" && (i + 1) <= NF && $(i + 1) == vmid) id = 1
    }
  }
  END { exit (kvm && id) ? 0 : 1 }
' \"\$start_script\"; then
  echo \"qm showcmd for VM \$vmid did not produce the expected KVM command\" >&2
  sed 's/^/    showcmd: /' \"\$start_script\" >&2
  exit 1
fi
smp_count=\$(awk '{ for (i = 1; i <= NF; i++) if ($i == "-smp") count++ } END { print count + 0 }' \"\$start_script\")
if [ \"\$smp_count\" -ne 1 ]; then
  echo \"Expected one native -smp argument for VM \$vmid, found \$smp_count. Remove stale qm args before starting.\" >&2
  exit 1
fi

sudo taskset -c \"\$cpuset\" bash \"\$start_script\""
}

pimox_shutdown_vm_gracefully() {
    local host="$1"
    local user="$2"
    local key_path="$3"
    local qm_bin="$4"
    local vmid="$5"

    pimox_ssh "${host}" "${user}" "${key_path}" "set -eu
sudo '${qm_bin}' shutdown '${vmid}' --timeout 120 || true
elapsed=0
while [ \"\$elapsed\" -lt 300 ]; do
  if sudo '${qm_bin}' status '${vmid}' | grep -q 'status: stopped'; then
    exit 0
  fi
  sleep 5
  elapsed=\$((elapsed + 5))
done
echo 'VM ${vmid} did not stop gracefully within 300 seconds.' >&2
exit 1"
}

pimox_guest_cpu_count() {
    local guest_ip="$1"
    local guest_user="$2"
    local guest_key_path="$3"
    local known_hosts_file="${REPO_ROOT}/.lab/pimox-worker-known_hosts"

    ssh -i "${guest_key_path}" \
        -o BatchMode=yes \
        -o ConnectTimeout=10 \
        -o IdentitiesOnly=yes \
        -o StrictHostKeyChecking=accept-new \
        -o UserKnownHostsFile="${known_hosts_file}" \
        "${guest_user}@${guest_ip}" \
        'nproc --all'
}

cluster_worker_var_file_has_workers() {
    local var_file="$1"

    [[ -s "${var_file}" ]] || return 1
    python3 - "${var_file}" <<'PY'
import json
import sys

with open(sys.argv[1], encoding="utf-8") as handle:
    document = json.load(handle)
sys.exit(0 if document.get("worker_nodes") else 1)
PY
}

cluster_worker_targets() {
    local worker_ssh_targets="${WORKER_SSH_TARGETS-}"
    local var_file="${REPO_ROOT}/.lab/cluster-workers.auto.tfvars.json"

    CLUSTER_WORKER_TARGETS=()
    read -r -a CLUSTER_WORKER_TARGETS <<< "${worker_ssh_targets}"

    if [[ -s "${var_file}" ]] && command -v python3 >/dev/null 2>&1; then
        while IFS= read -r target; do
            [[ -n "${target}" ]] || continue
            if [[ ! " ${CLUSTER_WORKER_TARGETS[*]-} " =~ [[:space:]]${target}[[:space:]] ]]; then
                CLUSTER_WORKER_TARGETS+=("${target}")
            fi
        done < <(python3 - "${var_file}" <<'PY'
import json
import sys

var_file = sys.argv[1]
with open(var_file, encoding="utf-8") as handle:
    document = json.load(handle)

for key, node in sorted((document.get("worker_nodes") or {}).items()):
    host = node.get("host")
    user = node.get("user")
    if host and user:
        print(f"{user}@{host}")
PY
)
    fi
}

pimox_worker_count_default() {
    local spec_file="${REPO_ROOT}/.lab/pimox-workers.tsv"
    local var_file="${REPO_ROOT}/.lab/cluster-workers.auto.tfvars.json"
    local worker_key_prefix="${LAB_PIMOX_WORKER_KEY_PREFIX:-pimox}"
    local worker_node_prefix="${LAB_PIMOX_WORKER_NODE_PREFIX:-pimox-worker}"
    local max_count=0
    local count
    local state_count

    if [[ -s "${spec_file}" ]]; then
        count="$(awk -F'\t' -v key_prefix="${worker_key_prefix}" -v node_prefix="${worker_node_prefix}-" '
            $1 ~ "^" key_prefix "[0-9]+$" {
                worker_index = substr($1, length(key_prefix) + 1) + 0
                if (worker_index > max) max = worker_index
            }
            $4 ~ "^" node_prefix "[0-9]+$" {
                worker_index = substr($4, length(node_prefix) + 1) + 0
                if (worker_index > max) max = worker_index
            }
            END { print max + 0 }
        ' "${spec_file}")"
        if [[ "${count}" =~ ^[0-9]+$ && "${count}" -gt "${max_count}" ]]; then
            max_count="${count}"
        fi
    fi

    if [[ -s "${var_file}" ]] && command -v python3 >/dev/null 2>&1; then
        count="$(python3 - "${var_file}" "${worker_key_prefix}" "${worker_node_prefix}" <<'PY'
import json
import re
import sys

var_file, worker_key_prefix, worker_node_prefix = sys.argv[1:4]
with open(var_file, encoding="utf-8") as handle:
    document = json.load(handle)
nodes = document.get("worker_nodes") or {}
highest = 0
key_pattern = re.compile(rf"^{re.escape(worker_key_prefix)}(\d+)$")
node_pattern = re.compile(rf"^{re.escape(worker_node_prefix)}-(\d+)$")
for key, node in nodes.items():
    for candidate in (key, str(node.get("node_name", ""))):
        match = key_pattern.match(candidate) or node_pattern.match(candidate)
        if match:
            highest = max(highest, int(match.group(1)))
print(highest)
PY
)"
        if [[ "${count}" =~ ^[0-9]+$ && "${count}" -gt "${max_count}" ]]; then
            max_count="${count}"
        fi
    fi

    state_count="$(tofu -chdir="${REPO_ROOT}/bootstrap/cluster" state show null_resource.worker_nodes_required 2>/dev/null |
        awk -F'"' '/"worker_count"[[:space:]]*=/ { print $4; found = 1 } END { exit found ? 0 : 1 }' || true)"
    if [[ "${state_count}" =~ ^[0-9]+$ && "${state_count}" -gt "${max_count}" ]]; then
        max_count="${state_count}"
    fi

    count="${LAB_PIMOX_DEFAULT_WORKER_COUNT:-1}"
    if [[ "${count}" =~ ^[0-9]+$ && "${count}" -gt "${max_count}" ]]; then
        max_count="${count}"
    fi

    printf '%s\n' "${max_count}"
}

pimox_worker_count_effective() {
    local configured="${LAB_PIMOX_WORKER_COUNT:-}"
    local detected

    detected="$(pimox_worker_count_default)"
    if [[ "${configured}" =~ ^[0-9]+$ && "${configured}" -gt "${detected}" ]]; then
        printf '%s\n' "${configured}"
        return 0
    fi
    printf '%s\n' "${detected}"
}

cluster_control_plane_tracked() {
    tofu_state_has_resource "bootstrap/cluster" "null_resource.kubeadm_control_plane"
}

cluster_admin_kubeconfig_present() {
    [[ -s "${KUBECONFIG_PATH}" ]] ||
        sudo test -s /etc/kubernetes/admin.conf 2>/dev/null ||
        [[ -n "${KUBECONFIG_PATH:-}" && -s "${KUBECONFIG_PATH}" ]]
}