#!/usr/bin/env bash
# passthrough.func — host-side passthrough logic (VAAPI & NVIDIA) for LXC
# This file ONLY touches host config (/etc/pve/lxc/<CTID>.conf) and whiptail.
# Inside-CT package setup lives in tools.func (hwaccel_setup_in_ct / nvidia_setup_in_ct).

# --------------------------- Common helpers -----------------------------------

_whiptail_dims() {
  local n="$1" L="$2"
  local maxW=$((L + 8))
  ((maxW < 70)) && maxW=70
  ((maxW > 100)) && maxW=100
  local H=$((10 + n * 2))
  ((H > 22)) && H=22
  echo "$H $maxW"
}

# ------------------------------ USB -------------------------------------------

usb_handle_passthrough() {
  local CTID="$1" CTTYPE="$2"
  local LXC_CONFIG="/etc/pve/lxc/${CTID}.conf"

  [[ "$CTTYPE" != "0" ]] && return 0 # USB passthrough only for privileged CTs

  cat <<EOF >>"$LXC_CONFIG"
# USB passthrough
lxc.cgroup2.devices.allow: a
lxc.cap.drop:
lxc.cgroup2.devices.allow: c 188:* rwm
lxc.cgroup2.devices.allow: c 189:* rwm
lxc.mount.entry: /dev/serial/by-id  dev/serial/by-id  none bind,optional,create=dir
lxc.mount.entry: /dev/ttyUSB0       dev/ttyUSB0       none bind,optional,create=file
lxc.mount.entry: /dev/ttyUSB1       dev/ttyUSB1       none bind,optional,create=file
lxc.mount.entry: /dev/ttyACM0       dev/ttyACM0       none bind,optional,create=file
lxc.mount.entry: /dev/ttyACM1       dev/ttyACM1       none bind,optional,create=file
EOF
}

# ------------------------------ VAAPI -----------------------------------------

_vaapi_gid() {
  local g="$1" gid
  gid="$(getent group "$g" | cut -d: -f3)"
  if [[ -z "$gid" ]]; then
    case "$g" in
    video) echo 44 ;;
    render) echo 104 ;;
    *) echo 44 ;;
    esac
  else
    echo "$gid"
  fi
}

_vaapi_pairs() {
  local seen=() by real id idx card pci pci_info name
  shopt -s nullglob
  for by in /dev/dri/by-path/*-render /dev/dri/renderD*; do
    [[ -e "$by" ]] || continue
    real="$(readlink -f "$by" || true)"
    [[ -e "$real" ]] || continue
    id="$(basename "$real")"
    [[ " ${seen[*]} " == *" $id "* ]] && continue
    seen+=("$id")

    idx="${id#renderD}"
    if [[ "$idx" =~ ^[0-9]+$ ]]; then
      idx=$((idx - 128))
    else
      idx=0
    fi
    card="/dev/dri/card${idx}"
    [[ -e "$card" ]] || card=""

    if [[ "$by" == *"/by-path/"* ]]; then
      pci="$(basename "$by" | sed -E 's/^pci-([0-9a-fA-F:.]+)-render/\1/')"
      pci_info="$(lspci -nn 2>/dev/null | grep -i "${pci#0000:}" || true)"
      name="${pci_info#*: }"
      [[ -z "$name" ]] && name="GPU ${pci}"
    else
      name="DRM $(basename "$real")"
    fi

    label="$(basename "$real")"
    [[ -n "$card" ]] && label+=" + $(basename "$card")"
    label+=" – ${name}"

    printf "%s:%s\t%s\n" "$real" "$card" "$label"
  done
  shopt -u nullglob
}

vaapi_select_and_apply() {
  local CTID="$1" CTTYPE="$2"
  local LXC_CONFIG="/etc/pve/lxc/${CTID}.conf"

  local pairs=() items=() maxlen=0 n h w
  mapfile -t pairs < <(_vaapi_pairs)

  if ((${#pairs[@]} == 0)); then
    msg_warn "No VAAPI devices detected – skipping."
    return 0
  fi

  for p in "${pairs[@]}"; do
    local devs="${p%%$'\t'*}"
    local label="${p#*$'\t'}"
    items+=("$devs" "$label" "OFF")
    ((${#label} > maxlen)) && maxlen=${#label}
  done
  n=$((${#items[@]} / 3))
  read -r h w < <(_whiptail_dims "$n" "$maxlen")

  if [[ "$CTTYPE" == "0" ]]; then
    whiptail --title "VAAPI passthrough" --msgbox "\
VAAPI passthrough will be enabled.

• Privileged CT: full DRM access
• You may need to install drivers inside the CT (intel-media-driver, vainfo)." 12 "$w"
  else
    whiptail --title "VAAPI passthrough (unprivileged)" --msgbox "\
Unprivileged CT: VAAPI may be limited.

If it fails, consider a privileged CT. You may still need drivers inside the CT." 12 "$w"
  fi

  local SELECTED
  SELECTED="$(
    whiptail --title "VAAPI Device Selection" \
      --checklist "Select GPU / VAAPI device(s) to passthrough:" "$h" "$w" "$((n > 6 ? 6 : n))" \
      "${items[@]}" 3>&1 1>&2 2>&3
  )" || {
    msg_warn "VAAPI selection cancelled."
    return 0
  }

  [[ -z "$SELECTED" ]] && {
    msg_warn "No devices selected – skipping."
    return 0
  }

  local DID_MOUNT_DRI=0 idx=0
  for dev in $SELECTED; do
    dev="${dev%\"}"
    dev="${dev#\"}"
    IFS=":" read -r path card <<<"$dev"
    for d in "$path" "$card"; do
      [[ -n "$d" && -e "$d" ]] || continue
      if [[ "$CTTYPE" == "0" ]]; then
        if [[ "$DID_MOUNT_DRI" -eq 0 && -d /dev/dri ]]; then
          echo "lxc.mount.entry: /dev/dri /dev/dri none bind,optional,create=dir" >>"$LXC_CONFIG"
          DID_MOUNT_DRI=1
        fi
        if mm=$(stat -c '%t:%T' "$d" 2>/dev/null | awk -F: '{printf "%d:%d","0x"$1,"0x"$2}'); then
          echo "lxc.cgroup2.devices.allow: c $mm rwm" >>"$LXC_CONFIG"
          echo "lxc.mount.entry: $d $d none bind,optional,create=file" >>"$LXC_CONFIG"
        else
          msg_warn "Could not stat $d – skipping."
        fi
      else
        local gid
        if [[ "$d" =~ renderD ]]; then gid="$(_vaapi_gid render)"; else gid="$(_vaapi_gid video)"; fi
        echo "dev${idx}: $d,gid=${gid}" >>"$LXC_CONFIG"
        idx=$((idx + 1))
      fi
    done
  done

  # fallback for card0/card1 flip
  if [[ "$CTTYPE" == "0" ]]; then
    cat <<'EOF' >>"$LXC_CONFIG"
# VAAPI fallback: bind /dev/dri and allow 226:* to survive node flips
lxc.mount.entry: /dev/dri /dev/dri none bind,optional,create=dir
lxc.cgroup2.devices.allow: c 226:* rwm
EOF
  fi
}

vaapi_handle_passthrough() {
  local CTID="$1" CTTYPE="$2" APP="$3"

  local VAAPI_APPS=(immich Channels Emby ErsatzTV Frigate Jellyfin Plex Scrypted Tdarr Unmanic Ollama FileFlows "Open WebUI" Tunarr Debian)
  local is_vaapi_app=false a
  for a in "${VAAPI_APPS[@]}"; do [[ "$APP" == "$a" ]] && is_vaapi_app=true && break; done

  if [[ "$CTTYPE" == "0" || "$is_vaapi_app" == "true" ]]; then
    # 1. write config
    vaapi_select_and_apply "$CTID" "$CTTYPE"

    # 2. ensure CT is running
    if ! pct status "$CTID" | grep -q "status: running"; then
      msg_info "Starting CT $CTID for VAAPI setup"
      pct start "$CTID"
      for i in {1..10}; do
        if pct status "$CTID" | grep -q "status: running"; then
          msg_ok "CT $CTID is running"
          break
        fi
        sleep 1
      done
    fi

    # 3. run setup inside CT
    pct exec "$CTID" -- bash -lc "_hwaccel_setup_in_ct '$CTTYPE'"
  fi
}

# ----------------------------- NVIDIA -----------------------------------------

nvidia_passthrough_to_lxc() {
  local CTID="$1" CTTYPE="${2:-0}"
  local LXC_CONFIG="/etc/pve/lxc/${CTID}.conf"
  local found=0 dev mm

  for dev in /dev/nvidia0 /dev/nvidia1 /dev/nvidiactl /dev/nvidia-uvm /dev/nvidia-uvm-tools /dev/nvidia-modeset; do
    [[ -e "$dev" ]] || continue
    found=1
    if mm="$(stat -c '%t:%T' "$dev" 2>/dev/null | awk -F: '{printf "%d:%d","0x"$1,"0x"$2}')"; then
      echo "lxc.cgroup2.devices.allow: c $mm rwm" >>"$LXC_CONFIG"
      echo "lxc.mount.entry: $dev $dev none bind,optional,create=file" >>"$LXC_CONFIG"
    fi
  done

  if [[ "$found" -eq 0 ]]; then
    msg_warn "No /dev/nvidia* devices found on host; skipping NVIDIA passthrough."
    return 0
  fi

  # optional: expose /dev/dri for apps probing VAAPI; harmless with NVIDIA
  if [[ -d /dev/dri && "$CTTYPE" == "0" ]]; then
    echo "lxc.mount.entry: /dev/dri /dev/dri none bind,optional,create=dir" >>"$LXC_CONFIG"
  fi

  msg_ok "NVIDIA devices mapped to CT ${CTID}"
}

nvidia_handle_passthrough() {
  local CTID="$1" CTTYPE="$2" APP="$3"

  # Only offer if NVIDIA devices exist on host
  compgen -G "/dev/nvidia*" >/dev/null || return 0

  if whiptail --title "NVIDIA passthrough" \
    --yesno "NVIDIA GPU detected on host.\n\nMap /dev/nvidia* into CT ${CTID} and install NVIDIA userland inside the container?" 12 70; then

    # 1. Host: map devices into LXC config
    nvidia_passthrough_to_lxc "$CTID" "$CTTYPE"

    # 2. CT must be running before pct exec
    if ! pct status "$CTID" | grep -q "status: running"; then
      msg_info "Starting CT $CTID for NVIDIA setup"
      pct start "$CTID"
      for i in {1..10}; do
        if pct status "$CTID" | grep -q "status: running"; then
          break
        fi
        sleep 1
      done
    fi

    # 3. Run NVIDIA setup inside CT
    pct exec "$CTID" -- bash -lc "_nvidia_setup_in_ct '$CTTYPE'"

  else
    msg_warn "Skipped NVIDIA passthrough by user choice."
  fi
}

# ------------------------------------------------------------------------------
# Hardware acceleration setup inside container (Intel/AMD via VAAPI)
# Debian 12 (bookworm), Debian 13 (trixie), Ubuntu 24.04 (noble)
# Usage: hwaccel_setup_in_ct <CTTYPE 0|1> [--nonfree-intel]
# ------------------------------------------------------------------------------
hwaccel_setup_in_ct() {
  local CTTYPE="$1" NONFREE=0
  [[ "$2" == "--nonfree-intel" ]] && NONFREE=1

  local ID VERSION_CODENAME
  if [[ -r /etc/os-release ]]; then . /etc/os-release; fi
  ID="${ID:-debian}"
  VERSION_CODENAME="${VERSION_CODENAME:-bookworm}"

  msg_info "Setting up VAAPI userland for ${ID^} ($VERSION_CODENAME)"

  case "$ID" in
  debian | ubuntu)
    if ((NONFREE)) && [[ "$VERSION_CODENAME" =~ (trixie|noble) ]]; then
      # Debian 13 / Ubuntu 24.04 — enable non-free Intel media (Debian deb822)
      if [[ "$VERSION_CODENAME" == "trixie" ]]; then
        cat >/etc/apt/sources.list.d/non-free.sources <<'SRC'
Types: deb deb-src
URIs: http://deb.debian.org/debian
Suites: trixie
Components: non-free non-free-firmware

Types: deb deb-src
URIs: http://deb.debian.org/debian-security
Suites: trixie-security
Components: non-free non-free-firmware

Types: deb deb-src
URIs: http://deb.debian.org/debian
Suites: trixie-updates
Components: non-free non-free-firmware
SRC
      fi

      $STD apt-get update
      $STD apt-get install -y \
        intel-media-va-driver-non-free \
        ocl-icd-libopencl1 \
        mesa-opencl-icd \
        mesa-va-drivers \
        libvpl2 \
        vainfo \
        intel-gpu-tools
    else
      $STD apt-get update
      $STD apt-get install -y \
        va-driver-all \
        ocl-icd-libopencl1 \
        mesa-opencl-icd \
        mesa-va-drivers \
        vainfo \
        intel-gpu-tools
    fi
    ;;
  *)
    msg_warn "Unsupported distro ($ID $VERSION_CODENAME) – skipping VAAPI setup."
    return 0
    ;;
  esac

  if [[ "$CTTYPE" == "0" ]]; then
    $STD adduser "$(id -un)" video || true
    $STD adduser "$(id -un)" render || true
  fi

  msg_ok "VAAPI userland ready"
}

# ------------------------------------------------------------------------------
# NVIDIA userland inside container
# Debian 12/13, Ubuntu 24.04
# Usage: nvidia_setup_in_ct <CTTYPE 0|1>
# ------------------------------------------------------------------------------
nvidia_setup_in_ct() {
  local CTTYPE="$1"

  local ID VERSION_CODENAME
  if [[ -r /etc/os-release ]]; then . /etc/os-release; fi
  ID="${ID:-debian}"
  VERSION_CODENAME="${VERSION_CODENAME:-bookworm}"

  msg_info "Installing NVIDIA userland on ${ID^} ($VERSION_CODENAME)"

  case "$ID" in
  debian | ubuntu)
    $STD apt-get update
    $STD apt-get install -y \
      nvidia-driver \
      nvidia-utils \
      libnvidia-encode1 \
      libcuda1 || {
      msg_error "Failed to install NVIDIA packages"
      return 1
    }
    ;;
  *)
    msg_warn "Unsupported distro ($ID $VERSION_CODENAME) – skipping NVIDIA setup."
    return 0
    ;;
  esac

  if [[ "$CTTYPE" == "0" ]]; then
    $STD adduser "$(id -un)" video || true
  fi

  msg_ok "NVIDIA userland ready"
}