#!/usr/bin/env bash
# Phase A + B install: process CPUAffinity drop-ins + boot IRQ affinity oneshot.
#
#   sudo bash install_cpu_affinity.sh
#   sudo OHD_CPU_ENABLE=0 bash install_cpu_affinity.sh   # remove pins / disable policy
#
# Re-run after deploy. Does not restart DyzineFrame (UI restart) — restarts
# ethergold units only when already enabled.

# CRLF heal without rewriting Dropbox tip (sed -i often fails / races there).
if grep -q $'\r' "${BASH_SOURCE[0]}" 2>/dev/null; then
  _tmp="$(mktemp)"
  tr -d '\r' <"${BASH_SOURCE[0]}" >"${_tmp}"
  chmod 0755 "${_tmp}"
  exec bash "${_tmp}" "$@"
fi

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=cpu_layout.sh
source "${SCRIPT_DIR}/cpu_layout.sh"
# shellcheck source=ethergold_field_env.sh
source "${SCRIPT_DIR}/ethergold_field_env.sh"

if [[ "${EUID}" -ne 0 ]]; then
  echo "FATAL: run with sudo" >&2
  exit 1
fi

IRQ_SCRIPT="${SCRIPT_DIR}/apply_irq_affinity.sh"
IRQ_LIB_DIR="/usr/local/lib/ohd"
IRQ_INSTALLED="${IRQ_LIB_DIR}/apply_irq_affinity.sh"
LAYOUT_INSTALLED="${IRQ_LIB_DIR}/cpu_layout.sh"
UNIT_IRQ="ohd-irq-affinity.service"
DROP_EG="/etc/systemd/system/${ETHERGOLD_UNIT}.d/cpu-affinity.conf"
DROP_COL="/etc/systemd/system/${ETHERGOLD_COLLECTOR_UNIT}.d/cpu-affinity.conf"

install_dropin() {
  local unit="$1" core="$2" path="$3"
  mkdir -p "$(dirname "${path}")"
  if ! ohd_cpu_lanes_ok; then
    rm -f "${path}"
    echo "cpu-affinity: skip drop-in ${unit} (lanes not ok nproc=$(ohd_cpu_count))"
    return 0
  fi
  cat >"${path}" <<EOF
# Installed by install_cpu_affinity.sh (Pi 4-lane OHD CPU policy).
# Re-run after changing OHD_CPU_* defaults in cpu_layout.sh.
[Service]
CPUAffinity=${core}
EOF
  echo "cpu-affinity: ${unit} CPUAffinity=${core} → ${path}"
}

_pin_pid() {
  local label="$1" pid="$2" mask="$3"
  [[ -n "${pid}" && -d "/proc/${pid}" ]] || return 0
  if taskset -pc "${mask}" "${pid}" >/dev/null 2>&1; then
    local list
    list="$(awk '/^Cpus_allowed_list:/{print $2}' "/proc/${pid}/status" 2>/dev/null || true)"
    echo "cpu-affinity: live-pin ${label} pid=${pid} → ${mask} (now ${list:-?})"
  else
    echo "cpu-affinity: WARN live-pin ${label} pid=${pid} failed" >&2
  fi
}

_live_rebind() {
  local mask="$1" which="${2:-all}"
  ohd_have_taskset || return 0
  local pid
  if [[ "${which}" == "all" || "${which}" == "eg" ]]; then
    for pid in $(pgrep -x pog422_uartd 2>/dev/null || true); do
      _pin_pid "uartd" "${pid}" "${mask}"
    done
    for pid in $(pgrep -f 'python3? .*(pog422_bridge\.py|run_ethergold|pog422_listen\.py)' 2>/dev/null || true); do
      _pin_pid "eg_python" "${pid}" "${mask}"
    done
  fi
  if [[ "${which}" == "all" || "${which}" == "df" ]]; then
    for pid in $(pgrep -f 'python3? .*dyzineframe_supervisor\.py' 2>/dev/null || true); do
      _pin_pid "df_supervisor" "${pid}" "${mask}"
    done
    for pid in $(pgrep -f 'python3? .*DyzineFrame\.py' 2>/dev/null || true); do
      _pin_pid "dyzineframe" "${pid}" "${mask}"
    done
  fi
}

if [[ "${OHD_CPU_ENABLE}" != "1" ]]; then
  rm -f "${DROP_EG}" "${DROP_COL}"
  systemctl disable --now "${UNIT_IRQ}" 2>/dev/null || true
  rm -f "/etc/systemd/system/${UNIT_IRQ}"
  systemctl daemon-reload
  if ohd_have_taskset; then
    _live_rebind "$(ohd_cpu_all_list)" all
  fi
  echo "cpu-affinity: DISABLED (OHD_CPU_ENABLE!=1 or run/cpu_affinity.disabled); drop-ins removed; live pids → all CPUs"
  unset -f _pin_pid _live_rebind
  exit 0
fi

if ! ohd_cpu_lanes_ok; then
  echo "cpu-affinity: WARN nproc=$(ohd_cpu_count) < 4 — no process pins applied" >&2
  exit 0
fi

install_dropin "${ETHERGOLD_UNIT}" "${OHD_CPU_EG}" "${DROP_EG}"
install_dropin "${ETHERGOLD_COLLECTOR_UNIT}" "${OHD_CPU_EG}" "${DROP_COL}"

# Install IRQ helper outside Dropbox so systemd never runs a CRLF/synced script.
mkdir -p "${IRQ_LIB_DIR}"
if [[ ! -f "${IRQ_SCRIPT}" ]]; then
  echo "cpu-affinity: WARN missing ${IRQ_SCRIPT} — skip IRQ oneshot" >&2
else
  # Strip CR on copy (Windows tip → Pi).
  tr -d '\r' <"${SCRIPT_DIR}/cpu_layout.sh" >"${LAYOUT_INSTALLED}"
  tr -d '\r' <"${IRQ_SCRIPT}" >"${IRQ_INSTALLED}"
  chmod 0755 "${IRQ_INSTALLED}"
  chmod 0644 "${LAYOUT_INSTALLED}"
fi

# Phase B oneshot — after udev; before ethergold when ordering works.
cat >"/etc/systemd/system/${UNIT_IRQ}" <<EOF
[Unit]
Description=OHD IRQ/RPS affinity (USB→EG core, net→NET core)
After=local-fs.target systemd-udev-settle.service
Before=ethergold.service ethergold-collector.service
DefaultDependencies=yes

[Service]
Type=oneshot
RemainAfterExit=yes
# Best-effort: pin failures must not fail the unit (avoids red "Job failed" on update.sh).
SuccessExitStatus=0
Environment=OHD_CPU_ENABLE=1
Environment=OHD_CPU_EG=${OHD_CPU_EG}
Environment=OHD_CPU_NET=${OHD_CPU_NET}
Environment=OHD_CPU_DF=${OHD_CPU_DF}
Environment=OHD_CPU_DF_SET=${OHD_CPU_DF_SET}
Environment=OHD_CPU_OS=${OHD_CPU_OS}
ExecStart=/bin/bash ${IRQ_INSTALLED}

[Install]
WantedBy=multi-user.target
EOF

systemctl daemon-reload
systemctl enable "${UNIT_IRQ}" >/dev/null 2>&1 || true
# Hide systemd's red "Job for … failed" — script always exits 0 now; still soft-guard.
if ! systemctl start "${UNIT_IRQ}" >/dev/null 2>&1; then
  echo "cpu-affinity: WARN ${UNIT_IRQ} start failed (non-fatal) — see journalctl -u ${UNIT_IRQ}" >&2
else
  echo "cpu-affinity: ${UNIT_IRQ} started (IRQ/RPS best-effort)"
fi

# Refresh units if present (prod installs).
for u in "${ETHERGOLD_UNIT}" "${ETHERGOLD_COLLECTOR_UNIT}"; do
  if systemctl is-enabled "${u}" >/dev/null 2>&1; then
    systemctl restart "${u}" || true
    echo "cpu-affinity: restarted ${u}"
  fi
done

# Live process pins (EG 3.0.2.13+):
# Bench bridges/DF often predate taskset launchers — systemd only covers ethergold
# units. Re-bind running pids so update.sh takes effect without a full reboot.
# DyzineFrame: pin supervisor *and* DyzineFrame.py to DF *set* (default 0-2,
# not a single core — fireworks starved on 3.0.2.13 exclusive-2 pin).
# Skip DF (and use all CPUs) when run/cpu_affinity.disabled is present.
if ohd_have_taskset && ohd_cpu_lanes_ok; then
  _live_rebind "${OHD_CPU_EG}" eg
  _live_rebind "${OHD_CPU_DF_SET}" df
fi
unset -f _pin_pid _live_rebind

# Ensure DF tree has cpu_affinity.sh so next *new* runit/supervisor start is pinned.
_DF_HOME=""
for _d in /home/topbox/DyzineFrame "${HOME}/DyzineFrame"; do
  if [[ -d "${_d}" ]]; then
    _DF_HOME="${_d}"
    break
  fi
done
if [[ -n "${_DF_HOME}" && ! -f "${_DF_HOME}/cpu_affinity.sh" ]]; then
  cat >"${_DF_HOME}/cpu_affinity.sh" <<'EOF'
# Installed by EtherGold install_cpu_affinity.sh — DF set excludes EG core.
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
  echo "FATAL: source cpu_affinity.sh" >&2
  exit 1
fi
_DF_CPU_LAYOUTS=(
  "${HOME}/EtherGold/prod/cpu_layout.sh"
  "${HOME}/EtherGold/dev/cpu_layout.sh"
)
_OHD_LAYOUT_SOURCED=0
for _layout in "${_DF_CPU_LAYOUTS[@]}"; do
  if [[ -f "${_layout}" ]]; then
    # shellcheck source=/dev/null
    source "${_layout}"
    _OHD_LAYOUT_SOURCED=1
    break
  fi
done
if [[ "${_OHD_LAYOUT_SOURCED}" != "1" ]]; then
  OHD_CPU_ENABLE="${OHD_CPU_ENABLE:-1}"
  OHD_CPU_DF="${OHD_CPU_DF:-2}"
  OHD_CPU_DF_SET="${OHD_CPU_DF_SET:-0-2}"
  ohd_cpu_count() { nproc 2>/dev/null || echo 1; }
  ohd_cpu_lanes_ok() {
    [[ "${OHD_CPU_ENABLE}" == "1" ]] || return 1
    [[ "$(ohd_cpu_count)" -ge 4 ]] || return 1
  }
  ohd_have_taskset() { command -v taskset >/dev/null 2>&1; }
  ohd_df_exec() {
    if ohd_cpu_lanes_ok && ohd_have_taskset; then
      exec taskset -c "${OHD_CPU_DF_SET}" "$@"
    else
      exec "$@"
    fi
  }
  ohd_df_run() {
    if ohd_cpu_lanes_ok && ohd_have_taskset; then
      taskset -c "${OHD_CPU_DF_SET}" "$@"
    else
      "$@"
    fi
  }
fi
EOF
  chown topbox:topbox "${_DF_HOME}/cpu_affinity.sh" 2>/dev/null || true
  echo "cpu-affinity: wrote ${_DF_HOME}/cpu_affinity.sh (was missing)"
fi
unset _d _DF_HOME

echo "cpu-affinity: layout os=${OHD_CPU_OS} net=${OHD_CPU_NET} df_pref=${OHD_CPU_DF} df_set=${OHD_CPU_DF_SET} eg=${OHD_CPU_EG} enable=${OHD_CPU_ENABLE}"
echo "  disable: touch ~/DyzineFrame/run/cpu_affinity.disabled  (or HDMI Supervisor CPU pin → Disable)"
echo "  verify:  taskset -pc \$(pgrep -n pog422_uartd)  # expect cpu ${OHD_CPU_EG} when enabled"
echo "  verify:  systemctl show ethergold -p CPUAffinity"
echo "  DF: live pin set=${OHD_CPU_DF_SET} unless cpu_affinity.disabled (all cores)"