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

# ============================================================
# function_rootfs.bash
#
# 说明：
#   存放符合 Bash Context Transaction Skill 第 5 条的函数。
#
# 对外 API 只使用以下模式：
#   init_rootfs_env CTX ...
#   step_${功能名称} CTX ...
#   clear_rootfs_env CTX
#
# 约束：
#   - 不提供 flow_xxx 之类的额外模式。
#   - 所有运行状态固定放入 CTX 关联数组。
#   - 数组参数传数组名，函数内部用 local -n 引用。
#   - 业务返回值使用 ret_ 变量名，函数内部用 local -n 写入。
#   - Bash return 只表示状态码。
# ============================================================

source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/base_rootfs.bash"

# ============================================================
# 1. init_${整体功能名称}_env / clear_${整体功能名称}_env
# ============================================================

init_rootfs_env() {
  if [[ "$#" -ne 5 ]]; then
    die "init_rootfs_env requires exactly 5 args: ctx_name iso_dev iso_mount rootfs apt_sources"
  fi

  local ctx_name="$1"
  local iso_dev="$2"
  local iso_mount="$3"
  local rootfs="$4"
  local apt_sources="$5"

  assert_assoc_name "$ctx_name"

  eval "declare -g -A ${ctx_name}=()"
  local -n C="$ctx_name"

  C["iso_dev"]="$iso_dev"
  C["iso_mount"]="$iso_mount"
  C["rootfs"]="$rootfs"
  C["apt_sources"]="$apt_sources"

  C["rootfs_marker"]="$rootfs/.wyc_rootfs_initialized"

  C["clean_done"]="0"
  C["clean_apt_after_install"]="0"

  if command -v mountpoint >/dev/null 2>&1 && mountpoint -q "$iso_mount"; then
    C["iso_mounted_before_init"]="1"
  else
    C["iso_mounted_before_init"]="0"
  fi
}

_rootfs_ensure_env() {
  local ctx_name="$1"
  assert_assoc_name "$ctx_name"

  local -n C="$ctx_name"

  if [[ -z "${C["rootfs"]:-}" ]]; then
    die "init_rootfs_env has not been called for ctx: $ctx_name"
  fi
}

_rootfs_ensure_exists() {
  local ctx_name="$1"
  _rootfs_ensure_env "$ctx_name"

  local -n C="$ctx_name"
  local rootfs="${C["rootfs"]}"

  if [[ ! -x "$rootfs/bin/bash" ]]; then
    die "rootfs not found or invalid: $rootfs"
  fi
}

_rootfs_require_host_cmds() {
  require_cmd sudo
  require_cmd mount
  require_cmd mountpoint
  require_cmd chroot
  require_cmd getent
  require_cmd awk
  require_cmd sed
  require_cmd sort
}

_rootfs_require_iso_cmds() {
  require_cmd unsquashfs
}

_rootfs_require_copy_cmds() {
  require_cmd sudo
  require_cmd cp
  require_cmd find
  require_cmd tar
}

_rootfs_clear_chroot_mounts() {
  local ctx_name="$1"
  _rootfs_ensure_env "$ctx_name"

  local -n C="$ctx_name"
  local rootfs="${C["rootfs"]}"

  log "cleanup chroot mounts..."

  sudo umount "$rootfs/run" 2>/dev/null || sudo umount -l "$rootfs/run" 2>/dev/null || true
  sudo umount "$rootfs/sys" 2>/dev/null || sudo umount -l "$rootfs/sys" 2>/dev/null || true
  sudo umount "$rootfs/proc" 2>/dev/null || sudo umount -l "$rootfs/proc" 2>/dev/null || true
  sudo umount "$rootfs/dev/pts" 2>/dev/null || sudo umount -l "$rootfs/dev/pts" 2>/dev/null || true
  sudo umount "$rootfs/dev" 2>/dev/null || sudo umount -l "$rootfs/dev" 2>/dev/null || true
}

clear_rootfs_env() {
  local ctx_name="$1"
  _rootfs_ensure_env "$ctx_name"

  local -n C="$ctx_name"

  if [[ "${C["clean_done"]}" == "1" ]]; then
    return 0
  fi

  C["clean_done"]="1"

  log "clear rootfs env..."

  _rootfs_clear_chroot_mounts "$ctx_name"

  local iso_mount="${C["iso_mount"]}"
  local iso_before="${C["iso_mounted_before_init"]:-0}"

  if [[ "$iso_before" == "0" ]]; then
    if mountpoint -q "$iso_mount"; then
      log "unmount ISO mounted after init_rootfs_env: $iso_mount"
      sudo umount "$iso_mount" 2>/dev/null || sudo umount -l "$iso_mount" 2>/dev/null || true
    fi
  else
    log "skip ISO unmount, ISO was mounted before init_rootfs_env: $iso_mount"
  fi
}

# ============================================================
# 2. step_${功能名称}: 返回值示例 ret_ + local -n
# ============================================================

step_get_apt_source_hosts() {
  if [[ "$#" -ne 2 ]]; then
    die "step_get_apt_source_hosts requires: ctx_name ret_hosts"
  fi

  local ctx_name="$1"
  local ret_hosts_name="$2"

  _rootfs_ensure_env "$ctx_name"
  assert_var_name "$ret_hosts_name"

  local -n C="$ctx_name"
  local -n ret_hosts="$ret_hosts_name"

  ret_hosts=()

  local host
  while IFS= read -r host; do
    [[ -z "$host" ]] && continue
    [[ "$host" == file:* ]] && continue
    ret_hosts+=("$host")
  done < <(
    printf '%s\n' "${C["apt_sources"]}" \
      | awk '/^[[:space:]]*deb[[:space:]]+/ { print $2 }' \
      | sed -E 's#^[a-zA-Z]+://([^/:]+).*#\1#' \
      | sort -u
  )
}

# ============================================================
# 3. step_${功能名称}: ISO / rootfs / chroot
# ============================================================

step_check_host_network() {
  local ctx_name="$1"
  _rootfs_ensure_env "$ctx_name"

  local ret_hosts=()
  step_get_apt_source_hosts "$ctx_name" ret_hosts

  log "check host DNS from apt sources..."

  if [[ "${#ret_hosts[@]}" -eq 0 ]]; then
    warn "no host parsed from apt_sources"
    return 0
  fi

  local host
  for host in "${ret_hosts[@]}"; do
    if ! getent hosts "$host" >/dev/null 2>&1; then
      warn "host cannot resolve: $host"
    fi
  done
}

step_mount_iso() {
  local ctx_name="$1"
  _rootfs_ensure_env "$ctx_name"

  local -n C="$ctx_name"

  _rootfs_require_host_cmds

  log "prepare ISO mount..."

  sudo mkdir -p "${C["iso_mount"]}"

  if mountpoint -q "${C["iso_mount"]}"; then
    log "${C["iso_mount"]} already mounted"
  else
    if [[ ! -b "${C["iso_dev"]}" ]]; then
      err "ISO device not found: ${C["iso_dev"]}"
      err "current block devices:"
      lsblk
      return 1
    fi

    log "mount ${C["iso_dev"]} -> ${C["iso_mount"]}"
    sudo mount "${C["iso_dev"]}" "${C["iso_mount"]}"
  fi

  if [[ ! -f "${C["iso_mount"]}/casper/filesystem.squashfs" ]]; then
    err "not found: ${C["iso_mount"]}/casper/filesystem.squashfs"
    err "search filesystem.squashfs:"
    sudo find "${C["iso_mount"]}" -name 'filesystem.squashfs' -print || true
    return 1
  fi

  log "found squashfs: ${C["iso_mount"]}/casper/filesystem.squashfs"
}

step_persistent_init_from_iso() {
  local ctx_name="$1"
  _rootfs_ensure_env "$ctx_name"

  local -n C="$ctx_name"

  _rootfs_require_iso_cmds
  _rootfs_require_host_cmds

  step_mount_iso "$ctx_name"
  _rootfs_clear_chroot_mounts "$ctx_name"

  if [[ -f "${C["rootfs_marker"]}" && -x "${C["rootfs"]}/bin/bash" ]]; then
    log "rootfs already initialized: ${C["rootfs"]}"
    log "skip extract"
    return 0
  fi

  log "remove old rootfs: ${C["rootfs"]}"
  sudo rm -rf "${C["rootfs"]}"

  log "create rootfs directory..."
  sudo mkdir -p "${C["rootfs"]}"

  log "extract filesystem.squashfs to rootfs..."
  sudo unsquashfs -f -d "${C["rootfs"]}" "${C["iso_mount"]}/casper/filesystem.squashfs"

  if [[ ! -x "${C["rootfs"]}/bin/bash" ]]; then
    die "rootfs extract failed: ${C["rootfs"]}/bin/bash not found"
  fi

  log "write rootfs initialized marker..."

  {
    echo "initialized_at=$(date -Iseconds)"
    echo "source_iso=${C["iso_dev"]}"
    echo "rootfs=${C["rootfs"]}"
  } | sudo tee "${C["rootfs_marker"]}" >/dev/null

  log "persistent init done: ${C["rootfs"]}"
}

step_force_rebuild_from_iso() {
  local ctx_name="$1"
  _rootfs_ensure_env "$ctx_name"

  local -n C="$ctx_name"

  _rootfs_require_iso_cmds
  _rootfs_require_host_cmds

  step_mount_iso "$ctx_name"

  log "force rebuild rootfs..."
  _rootfs_clear_chroot_mounts "$ctx_name"

  sudo rm -rf "${C["rootfs"]}"
  sudo mkdir -p "${C["rootfs"]}"

  sudo unsquashfs -f -d "${C["rootfs"]}" "${C["iso_mount"]}/casper/filesystem.squashfs"

  if [[ ! -x "${C["rootfs"]}/bin/bash" ]]; then
    die "rootfs extract failed: ${C["rootfs"]}/bin/bash not found"
  fi

  {
    echo "initialized_at=$(date -Iseconds)"
    echo "source_iso=${C["iso_dev"]}"
    echo "rootfs=${C["rootfs"]}"
  } | sudo tee "${C["rootfs_marker"]}" >/dev/null

  log "force rebuild done: ${C["rootfs"]}"
}

step_mount_chroot_fs() {
  local ctx_name="$1"
  _rootfs_ensure_exists "$ctx_name"

  local -n C="$ctx_name"
  local rootfs="${C["rootfs"]}"

  _rootfs_require_host_cmds

  log "mount chroot filesystems..."

  sudo mkdir -p \
    "$rootfs/dev" \
    "$rootfs/dev/pts" \
    "$rootfs/proc" \
    "$rootfs/sys" \
    "$rootfs/run"

  mountpoint -q "$rootfs/dev" || sudo mount --bind /dev "$rootfs/dev"
  mountpoint -q "$rootfs/dev/pts" || sudo mount --bind /dev/pts "$rootfs/dev/pts"
  mountpoint -q "$rootfs/proc" || sudo mount -t proc proc "$rootfs/proc"
  mountpoint -q "$rootfs/sys" || sudo mount -t sysfs sysfs "$rootfs/sys"
  mountpoint -q "$rootfs/run" || sudo mount -t tmpfs tmpfs "$rootfs/run"
}

# ============================================================
# 4. step_${功能名称}: apt
# ============================================================

step_write_apt_sources() {
  local ctx_name="$1"
  _rootfs_ensure_exists "$ctx_name"

  local -n C="$ctx_name"
  local rootfs="${C["rootfs"]}"

  if [[ -z "${C["apt_sources"]:-}" ]]; then
    die "apt_sources is empty"
  fi

  log "write apt sources.list..."

  sudo mkdir -p "$rootfs/etc/apt"
  printf '%s\n' "${C["apt_sources"]}" | sudo tee "$rootfs/etc/apt/sources.list" >/dev/null

  log "remove possible cdrom source files..."
  sudo rm -f "$rootfs/etc/apt/sources.list.d/"*.list 2>/dev/null || true
  sudo rm -f "$rootfs/etc/apt/sources.list.d/"*.sources 2>/dev/null || true
}

step_write_dns() {
  local ctx_name="$1"
  _rootfs_ensure_exists "$ctx_name"

  local -n C="$ctx_name"
  local rootfs="${C["rootfs"]}"

  log "write DNS config..."

  sudo rm -f "$rootfs/etc/resolv.conf"

  if [[ -f /run/systemd/resolve/resolv.conf ]] && grep -q '^nameserver ' /run/systemd/resolve/resolv.conf; then
    log "copy host DNS as regular file from /run/systemd/resolve/resolv.conf"

    sudo awk '
      /^nameserver / { print }
      /^search / { print }
      /^options / { print }
    ' /run/systemd/resolve/resolv.conf | sudo tee "$rootfs/etc/resolv.conf" >/dev/null

  elif [[ -f /etc/resolv.conf ]] && grep -q '^nameserver ' /etc/resolv.conf; then
    log "copy host DNS as regular file from /etc/resolv.conf"

    sudo awk '
      /^nameserver / { print }
      /^search / { print }
      /^options / { print }
    ' /etc/resolv.conf | sudo tee "$rootfs/etc/resolv.conf" >/dev/null

  else
    warn "no usable host DNS found, use fallback DNS"

    sudo tee "$rootfs/etc/resolv.conf" >/dev/null <<'EOF'
nameserver 223.5.5.5
nameserver 223.6.6.6
nameserver 8.8.8.8
nameserver 1.1.1.1
EOF
  fi

  log "rootfs resolv.conf:"
  sudo cat "$rootfs/etc/resolv.conf"
}

step_disable_chroot_apt_hooks() {
  local ctx_name="$1"
  _rootfs_ensure_exists "$ctx_name"

  local -n C="$ctx_name"
  local rootfs="${C["rootfs"]}"

  log "disable chroot-unfriendly apt hooks..."

  sudo mkdir -p "$rootfs/etc/apt/apt.conf.d.disabled"

  sudo find "$rootfs/etc/apt/apt.conf.d" -maxdepth 1 -type f -iname '*zsys*' -exec \
    sudo mv -t "$rootfs/etc/apt/apt.conf.d.disabled" {} + 2>/dev/null || true
}

step_prepare_apt() {
  local ctx_name="$1"

  _rootfs_require_host_cmds
  _rootfs_ensure_exists "$ctx_name"

  step_check_host_network "$ctx_name"
  step_write_apt_sources "$ctx_name"
  step_write_dns "$ctx_name"
  step_disable_chroot_apt_hooks "$ctx_name"
  step_mount_chroot_fs "$ctx_name"
}

step_apt_update_strict() {
  local ctx_name="$1"
  _rootfs_ensure_exists "$ctx_name"

  local -n C="$ctx_name"
  local rootfs="${C["rootfs"]}"

  log "run strict apt update inside chroot..."

  sudo chroot "$rootfs" /bin/bash -s <<'EOF'
set -e

echo "[CHROOT] sources.list:"
cat /etc/apt/sources.list

echo "[CHROOT] DNS config:"
cat /etc/resolv.conf

echo "[CHROOT] clean old apt lists"
apt-get clean
rm -rf /var/lib/apt/lists/*

echo "[CHROOT] strict apt update"
apt-get update \
  -o Acquire::Retries=3 \
  -o APT::Update::Error-Mode=any
EOF
}

step_install_packages() {
  local ctx_name="$1"
  local packages_name="$2"

  _rootfs_ensure_exists "$ctx_name"
  assert_var_name "$packages_name"

  local -n C="$ctx_name"
  local -n packages="$packages_name"
  local rootfs="${C["rootfs"]}"

  if [[ "${#packages[@]}" -eq 0 ]]; then
    die "step_install_packages requires non-empty package array"
  fi

  log "install packages incrementally:"
  printf '  %s\n' "${packages[@]}"

  sudo chroot "$rootfs" /usr/bin/env DEBIAN_FRONTEND=noninteractive \
    apt-get install -y "${packages[@]}" || return $?

  if [[ "${C["clean_apt_after_install"]}" == "1" ]]; then
    step_apt_clean "$ctx_name"
  fi
}

step_remove_packages() {
  local ctx_name="$1"
  local packages_name="$2"

  _rootfs_ensure_exists "$ctx_name"
  assert_var_name "$packages_name"

  local -n C="$ctx_name"
  local -n packages="$packages_name"
  local rootfs="${C["rootfs"]}"

  if [[ "${#packages[@]}" -eq 0 ]]; then
    die "step_remove_packages requires non-empty package array"
  fi

  log "remove packages:"
  printf '  %s\n' "${packages[@]}"

  sudo chroot "$rootfs" /usr/bin/env DEBIAN_FRONTEND=noninteractive \
    apt-get remove -y "${packages[@]}" || return $?
}

step_apt_clean() {
  local ctx_name="$1"
  _rootfs_ensure_exists "$ctx_name"

  local -n C="$ctx_name"
  local rootfs="${C["rootfs"]}"

  log "clean apt cache inside rootfs..."

  sudo chroot "$rootfs" /bin/bash -c '
set -e
apt-get clean
rm -rf /var/lib/apt/lists/*
rm -rf /tmp/*
'
}

# ============================================================
# 5. step_${功能名称}: export / pack
# ============================================================

step_export_dev_sysroot() {
  local ctx_name="$1"
  local dst_root="$2"

  _rootfs_ensure_exists "$ctx_name"
  _rootfs_require_copy_cmds

  if [[ -z "$dst_root" ]]; then
    die "step_export_dev_sysroot requires destination directory"
  fi

  local -n C="$ctx_name"
  local src_root="${C["rootfs"]}"

  _rootfs_clear_chroot_mounts "$ctx_name"

  log "export slim dev sysroot"
  log "source rootfs: $src_root"
  log "target sysroot: $dst_root"

  sudo rm -rf "$dst_root"
  sudo mkdir -p "$dst_root"

  copy_path_if_exists "$src_root" "$dst_root" "include"
  copy_path_if_exists "$src_root" "$dst_root" "usr/include"
  copy_path_if_exists "$src_root" "$dst_root" "usr/local/include"

  copy_path_if_exists "$src_root" "$dst_root" "lib"
  copy_path_if_exists "$src_root" "$dst_root" "lib64"
  copy_path_if_exists "$src_root" "$dst_root" "usr/lib"
  copy_path_if_exists "$src_root" "$dst_root" "usr/local/lib"

  copy_path_if_exists "$src_root" "$dst_root" "usr/lib/gcc"
  copy_path_if_exists "$src_root" "$dst_root" "usr/libexec/gcc"

  copy_path_if_exists "$src_root" "$dst_root" "usr/share/pkgconfig"
  copy_path_if_exists "$src_root" "$dst_root" "usr/local/share/pkgconfig"

  copy_path_if_exists "$src_root" "$dst_root" "etc/ld.so.conf"
  copy_path_if_exists "$src_root" "$dst_root" "etc/ld.so.conf.d"

  copy_path_if_exists "$src_root" "$dst_root" "etc/os-release"
  copy_path_if_exists "$src_root" "$dst_root" "usr/lib/os-release"

  log "remove unnecessary paths if any..."
  sudo rm -rf \
    "$dst_root/var" \
    "$dst_root/tmp" \
    "$dst_root/run" \
    "$dst_root/proc" \
    "$dst_root/sys" \
    "$dst_root/dev" \
    "$dst_root/lib/modules" \
    "$dst_root/usr/lib/modules" \
    2>/dev/null || true

  log "slim dev sysroot export done: $dst_root"
}

step_pack_rootfs() {
  local ctx_name="$1"
  local output="$2"

  _rootfs_ensure_exists "$ctx_name"

  if [[ -z "$output" ]]; then
    die "step_pack_rootfs requires output file"
  fi

  local -n C="$ctx_name"
  local rootfs="${C["rootfs"]}"

  _rootfs_clear_chroot_mounts "$ctx_name"

  log "pack rootfs to: $output"
  sudo mkdir -p "$(dirname "$output")"
  sudo tar --numeric-owner -C "$rootfs" -czf "$output" .
}

step_pack_dev_sysroot() {
  local ctx_name="$1"
  local dst_root="$2"
  local output="$3"

  if [[ -z "$dst_root" || -z "$output" ]]; then
    die "step_pack_dev_sysroot requires dst_root and output"
  fi

  step_export_dev_sysroot "$ctx_name" "$dst_root"

  log "pack slim dev sysroot to: $output"
  sudo mkdir -p "$(dirname "$output")"
  sudo tar --numeric-owner -C "$dst_root" -czf "$output" .
}
