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

CTTYPE="${CTTYPE:-${CT_TYPE:-1}}"

# --------------------------- 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"
}

select_hw_passthrough() {
  local CT_ID="$1" CT_TYPE="$2" APP="$3"
  local LXC_CONFIG="/etc/pve/lxc/${CT_ID}.conf"

  if ! _is_gpu_app "$APP" && [[ "$CT_TYPE" != "0" ]]; then
    return
  fi

  local choices=()
  [[ -d /dev/dri ]] && choices+=("VAAPI" "Intel/AMD GPU via VAAPI" OFF)
  compgen -G "/dev/nvidia*" >/dev/null && choices+=("NVIDIA" "NVIDIA GPU passthrough" OFF)

  # no GPUs found
  [[ ${#choices[@]} -eq 0 ]] && {
    msg_info "No GPU devices detected"
    return
  }

  local SELECTED
  if [[ ${#choices[@]} -eq 2 ]]; then
    # both available → show whiptail
    SELECTED=$(whiptail --title "GPU Passthrough" \
      --checklist "Select GPU passthrough for CT $CT_ID:" 12 70 2 \
      "${choices[@]}" 3>&1 1>&2 2>&3) || return
  else
    # only one option → auto-select
    SELECTED="\"${choices[0]}\""
    msg_info "Auto-selecting GPU passthrough: ${choices[0]}"
  fi

  for sel in $SELECTED; do
    case "$sel" in
    "\"VAAPI\"")
      export ENABLE_VAAPI=1
      vaapi_select_and_apply "$CT_ID" "$CT_TYPE"
      ;;
    "\"NVIDIA\"")
      export ENABLE_NVIDIA=1
      nvidia_passthrough_to_lxc "$CT_ID" "$CT_TYPE"
      ;;
    esac
  done
}

# Apps that benefit from GPU passthrough (VAAPI + NVIDIA)
_GPU_APPS=(
  immich
  channels
  emby
  ersatztv
  frigate
  jellyfin
  plex
  scrypted
  tdarr
  unmanic
  ollama
  fileflows
  "open webui"
  tunarr
  debian
)

_is_gpu_app() {
  local app="$1"
  local a
  shopt -s nocasematch
  for a in "${_GPU_APPS[@]}"; do
    [[ "$app" == "$a" ]] && shopt -u nocasematch && return 0
  done
  shopt -u nocasematch
  return 1
}

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

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

  if [[ "$CT_TYPE" != "0" ]]; then
    return 0 # USB passthrough only for privileged CTs
  fi
  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)"
  case "$g" in
  video) echo "${gid:-44}" ;;
  render) echo "${gid:-104}" ;;
  *) echo "${gid:-44}" ;;
  esac
}

_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}"
    [[ "$idx" =~ ^[0-9]+$ ]] && idx=$((idx - 128)) || idx=0
    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 CT_ID="$1" CT_TYPE="$2"
  local LXC_CONFIG="/etc/pve/lxc/${CT_ID}.conf"

  mapfile -t pairs < <(_vaapi_pairs)
  ((${#pairs[@]} == 0)) && {
    msg_warn "No VAAPI devices detected – skipping."
    return
  }

  # only one device -> auto-select
  local SELECTED
  if [[ ${#pairs[@]} -eq 1 ]]; then
    SELECTED="${pairs[0]%%$'\t'*}"
    msg_info "Auto-selecting VAAPI device: ${pairs[0]#*$'\t'}"
  else
    # more than one device -> show whiptail
    local items=() maxlen=0
    for p in "${pairs[@]}"; do
      local devs="${p%%$'\t'*}" label="${p#*$'\t'}"
      items+=("$devs" "$label" "OFF")
      ((${#label} > maxlen)) && maxlen=${#label}
    done
    read -r h w < <(_whiptail_dims $((${#items[@]} / 3)) "$maxlen")

    SELECTED="$(whiptail --title "VAAPI Device Selection" \
      --checklist "Select VAAPI devices for CT $CT_ID:" "$h" "$w" 6 \
      "${items[@]}" 3>&1 1>&2 2>&3)" || {
      msg_warn "VAAPI selection cancelled."
      return
    }
    [[ -z "$SELECTED" ]] && {
      msg_warn "No VAAPI devices selected."
      return
    }
  fi

  # Apply selection to LXC config
  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 [[ "$CT_TYPE" == "0" ]]; then
        [[ $DID_MOUNT_DRI -eq 0 && -d /dev/dri ]] && {
          echo "lxc.mount.entry: /dev/dri /dev/dri none bind,optional,create=dir" >>"$LXC_CONFIG"
          DID_MOUNT_DRI=1
        }
        if mm=$(stat -c '%t:%T' "$d" | 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"
        fi
      else
        gid=$([[ "$d" =~ renderD ]] && _vaapi_gid render || _vaapi_gid video)
        echo "dev${idx}: $d,gid=${gid}" >>"$LXC_CONFIG"
        idx=$((idx + 1))
      fi
    done
  done

  # Fallback only for privileged CTs
  if [[ "$CT_TYPE" == "0" ]]; then
    cat <<'EOF' >>"$LXC_CONFIG"
# VAAPI fallback
lxc.mount.entry: /dev/dri /dev/dri none bind,optional,create=dir
lxc.cgroup2.devices.allow: c 226:* rwm
EOF
  fi
}

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

nvidia_passthrough_to_lxc() {
  local CT_ID="$1" CT_TYPE="$2"
  local LXC_CONFIG="/etc/pve/lxc/${CT_ID}.conf"
  local found=0

  for dev in /dev/nvidia*; do
    [[ -e "$dev" ]] || continue
    found=1
    if mm="$(stat -c '%t:%T' "$dev" | 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
  ((found == 0)) && {
    msg_warn "No NVIDIA devices found."
    return
  }

  if [[ -d /dev/dri && "$CT_TYPE" == "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 ${CT_ID}"
}

install_vaapi_userland_interactive() {
  . /etc/os-release
  if [[ "$VERSION_CODENAME" == "trixie" ]]; then
    read -r -p "${TAB3}Do you need the intel-media-va-driver-non-free driver for HW encoding (Debian 13 only)? <y/N> " prompt
    if [[ ${prompt,,} =~ ^(y|yes)$ ]]; then
      msg_info "Installing Intel Hardware Acceleration (non-free)"
      cat <<'EOF' >/etc/apt/sources.list.d/non-free.sources
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
EOF
      $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
      msg_ok "Installed Intel Hardware Acceleration (non-free)"
      return
    fi
  fi

  msg_info "Installing Intel Hardware Acceleration (open packages)"
  $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
  msg_ok "Installed Intel Hardware Acceleration (open packages)"
}

install_nvidia_userland_interactive() {
  msg_info "Installing NVIDIA Userland"
  $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_ok "Installed NVIDIA Userland"
}
