#!/bin/bash

set -uo pipefail

if [[ "${BASH_SOURCE[0]}" != "$0" ]]; then
   >&2 echo "FATAL: run this script, don't source it."
   return 1 2>/dev/null || exit 1
fi

rootpath() {
   local dir=$(dirname "$(readlink -e "${BASH_SOURCE[0]}")")
   if [[ -f "${dir}/LICENSE" ]] ; then
      echo "$dir"
   elif [[ -d /usr/share/rpi-image-gen ]] ; then
      echo /usr/share/rpi-image-gen
   else
      >&2 echo "FATAL: cannot locate project root" ; exit 1
   fi
}
readonly IGTOP=$(rootpath)

# Load helpers
source "$IGTOP/lib/cli.sh"
source "$IGTOP/lib/common.sh"
source "$IGTOP/lib/dependencies.sh"
source "$IGTOP/lib/tools.sh"


# Namespace setup (see bin/ns)
_ns_setup() {
   # unmount cache (defensive - cleanup hook should handle)
   if [ -n "${_NS_APT_ARCHIVES:-}" ]; then
      trap "umount -q \"$_NS_APT_ARCHIVES\" 2>/dev/null || true" EXIT
   fi
}
export NS_SETUP="$(declare -f _ns_setup)"


# cli
cmd=${1:-}
shift
case $cmd in
   build|clean|layer|metadata|config) ;;
   help) cli_help ; exit 0 ;;
   doc*)
      port=${1:-3142}
      [[ $port =~ ^[0-9]+$ ]] && (( port > 0 && port < 65536 )) || die "docs: invalid port '$port'"
      bind=$(hostname -I | cut -d' ' -f1)
      echo "Point your browser at http://${bind}:${port}"
      python3 -m http.server "$port" --directory "${IGTOP}/docs/" --bind "$bind"
      exit 0
      ;;
   *) cli_help; die "Unknown command" ;;
esac
argv=("$@")

if [[ ${argv[0]:-} == -h || ${argv[0]:-} == --help ]]; then
  case $cmd in
    layer|metadata|config) exec "$IGTOP/bin/ig" "$cmd" --help ;;
    build|clean) cli_help_"$cmd"; exit 0 ;;
  esac
fi


# Host context init
dependencies_check --category bootstrap "${IGTOP}/depends" || exit 1
TMPDIR=$(mktemp -d || exit 1 )
trap 'rm -rf "$TMPDIR"' EXIT
$IGTOP/bin/ig env "${argv[@]}" > "${TMPDIR}/host.json" || die
source <("${IGTOP}/builtin/scripts/host2sh.py" "${TMPDIR}/host.json")

: "${HOST_CONFIG_PATH?missing}"
: "${HOST_LAYER_PATH?missing}"
: "${HOST_EXEC_PATH?missing}"
: "${INTERACTIVE?missing}"
: "${ONLY_FS?missing}"
: "${ONLY_IMAGE?missing}"

# For dynamic layers
DYNROOT="${TMPDIR}/dynamic"
mkdir -p "${DYNROOT}/layer"
HOST_LAYER_PATH="DYNlayer=${DYNROOT}/layer:${HOST_LAYER_PATH}"

PATH="$HOST_EXEC_PATH:$PATH"
export PATH

# shortcut to loaders
case $cmd in
   build|clean)
      [[ -z ${HOST_CONFIG_FILE:-} ]] && { cli_help_"$cmd" ; die ;}
      ;;
   layer)
      exec "$IGTOP/bin/ig" layer --path "$HOST_LAYER_PATH" "${argv[@]}"
      ;;
   metadata|config)
    exec "$IGTOP/bin/ig" "$cmd" "${argv[@]}"
    ;;
esac


# cmdline
OVRF="${TMPDIR}/overrides"
: > "$OVRF"
[[ ${#OVERRIDES[@]} -gt 0 ]] && printf '%s\n' "${OVERRIDES[@]}" >> "$OVRF"
[[ -n ${HOST_BUILD_DIR:-} ]] && printf 'IGconf_sys_workroot="%s"\n' "$HOST_BUILD_DIR" >> "$OVRF"
# Must hold sane key=value pairs
while IFS= read -r line || [ -n "$line" ]; do
   case "$line" in
      [A-Za-z_][A-Za-z0-9_]*=*) ;;
      *)
         nl -ba "$OVRF" >&2
         die "Overrides must be provided as key=value pairs."
         ;;
    esac
done < "$OVRF"


# Ready for pipeline flow
dependencies_check --category all "${IGTOP}/depends" || exit 1
declare -A ctx=()
ctx[INTERACTIVE]="${INTERACTIVE:?missing}"
ctx[ONLY_FS]="${ONLY_FS:?missing}"
ctx[ONLY_IMAGE]="${ONLY_IMAGE:?missing}"
ctx[SRCROOT]="${SRCROOT:-$(realpath -e .)}"


###############################################################################
# Stage 1: Parameter assembly
#   Read registry
#   Parse config
#   Seed env for subsequent stages
###############################################################################
parameter_assembly()
{
   msg "\nPARAM"

   # Seed with registry defaults
   ig metadata --emit "$IGTOP/registry.defs" > "${TMPDIR}/registry.env" \
      || die "Failed to parse registry"

   # Read config, writing all settings to file. Deferred variable resolution in pipeline
   # means that variable expansion does not happen here, so a config file can use any
   # variable it wants (from layers, env, anchors etc). Nothing is resolved until pipeline runs.
   env \
      ig config \
      --path "$HOST_CONFIG_PATH" \
      "$HOST_CONFIG_FILE"     \
      --overrides "$OVRF" \
      --write-to "${TMPDIR}/userconfig.env" \
      || die "Config parse failed"

   # Prepend registry defaults so config overrides (last wins)
   cat "${TMPDIR}/registry.env" "${TMPDIR}/userconfig.env" > "${TMPDIR}/config.env"

   # Toolchain variables. Native unless ARCH forces cross.
   if [ -z "${ARCH:-}" ]; then
      ARCH=$(get_var ARCH "${TMPDIR}/config.env") || true
   fi

   DEB_BUILD_ARCH_VAL="${DEB_BUILD_ARCH:-$(dpkg-architecture -qDEB_BUILD_ARCH)}"
   DEB_BUILD_GNU_TYPE_VAL="${DEB_BUILD_GNU_TYPE:-$(dpkg-architecture -a"$DEB_BUILD_ARCH_VAL" -qDEB_BUILD_GNU_TYPE)}"

   if [ -n "${ARCH:-}" ]; then
      DEB_HOST_ARCH_VAL=${DEB_HOST_ARCH:-$(dpkg-architecture -a"$ARCH" -qDEB_HOST_ARCH 2>/dev/null)}
   else
      DEB_HOST_ARCH_VAL=${DEB_HOST_ARCH:-$DEB_BUILD_ARCH_VAL}
   fi
   DEB_HOST_GNU_TYPE_VAL="${DEB_HOST_GNU_TYPE:-$(dpkg-architecture -a"$DEB_HOST_ARCH_VAL" -qDEB_HOST_GNU_TYPE 2>/dev/null)}"

   TOOLCHAIN_MODE=$([ "$DEB_BUILD_ARCH_VAL" = "$DEB_HOST_ARCH_VAL" ] && echo native || echo cross)

   # These canonical variables are always provided to pipeline
   cat >> "${TMPDIR}/config.env" <<EOF
DEB_BUILD_ARCH="$DEB_BUILD_ARCH_VAL"
DEB_BUILD_GNU_TYPE="$DEB_BUILD_GNU_TYPE_VAL"
DEB_HOST_ARCH="$DEB_HOST_ARCH_VAL"
DEB_HOST_GNU_TYPE="$DEB_HOST_GNU_TYPE_VAL"
TOOLCHAIN_MODE="$TOOLCHAIN_MODE"
IG_BUILD_ID="$(git -C "${ctx[SRCROOT]}" describe --tags --always --dirty 2> /dev/null || echo "999-unknown")"
IGTOP="$IGTOP"
IGROOT="$IGTOP"
LAYER_HOOKS="${IGTOP}/layer-hooks"
RPI_TEMPLATES="${IGTOP}/templates/rpi"
DYNROOT="$DYNROOT"
EOF

   # Conditional
   [[ -n "${SRCROOT:-}" ]] && \
      printf 'SRCROOT="%s"\n' "$SRCROOT" >> "${TMPDIR}/config.env"

  msg "\nENV" 
  cat ${TMPDIR}/config.env
}


###############################################################################
# Stage 2: Layer processing
#   Collect all specified layers
#   Ingest config environment and run pipeline to output the build config
#   Expand all variables and generate layer order
#   Write bootstrap information to file
##############################################################################
collect_layers()
{
   msg "\nLAYER"

   local -a layers=(essential)
   _collect() {
      local k=$1 v=$2
      case $k in
         IGconf_device_layer|IGconf_image_layer|IGconf_layer_*) layers+=("$v");;
      esac
   }
   mapfile_kv "${TMPDIR}/config.env" _collect

   msg "PIPELINE: ${layers[*]}"
   msg "SEARCH: $HOST_LAYER_PATH"

   # Generate the bootstrap information, resolving all variables. Configuration
   # input is only from the input env file, so run it in a clean room.
   env -i PATH="$PATH" \
     ig pipeline \
      --env-in "${TMPDIR}/config.env" \
      --layers "${layers[@]}" \
      --path "$HOST_LAYER_PATH" \
      --env-out "${TMPDIR}/env.out" \
      --plan-out "${TMPDIR}/layer.plan" \
      || die "pipeline failed"

   msg "PIPELINE: OK"

   # Normalise
   safe_kv "${TMPDIR}/env.out" > "${TMPDIR}/env.safe"

   # Expand all variables with a strict policy, posix only
   local -a vars
   mapfile -t vars < <(grep -oE '^[A-Za-z_][A-Za-z0-9_]*' "${TMPDIR}/env.safe")

   if ! env -i POSIXLY_CORRECT=1 \
      /bin/sh -c '
         set -aeu
         . "$1"
         shift
         for v in "$@"; do
            eval "val=\${$v-}"
            printf "%s=\"%s\"\n" "$v" "$val"
         done
         ' _ "${TMPDIR}/env.safe" "${vars[@]}" > "${TMPDIR}/final.env"
      then
         nl -ba "${TMPDIR}/env.safe" >&2
         die "Layer env expansion failed"
   fi

   # Write bootstrap information
   local bdir=$(get_var IGconf_sys_bootstrapdir "${TMPDIR}/final.env") || \
      die "No bootstrap dir"

   mkdir -p "$bdir" || die
   for f in final.env layer.plan host.json ; do
      src="${TMPDIR}/$f"
      [[ -f "$src" ]] || die "Missing bootstrap file: $src"
      cp "$src" "$bdir" || die
   done
   if [ -d "${DYNROOT}" ] ; then
      mkdir -p "${bdir}/dynamic" && rsync -a ${DYNROOT}/ "${bdir}/dynamic/"
   fi

   ctx[FINALENV]="${bdir}/final.env"
   ctx[LAYER_PLAN]="${bdir}/layer.plan"
}



###############################################################################
# Stage 3: Filesystem build preparation
#   Create output directories
#   Resolve host paths
#   Build mission critical tools
#   Initialise bdebstrap environment - apt, proxy, cache, etc
#   Assemble compatible layers
###############################################################################
prepare_build_config()
{
   msg "\nPREPARE"

   set_kv() {
      local key="$1" val="$2"
      case $key in
         IGconf_device_hostname|\
         IGconf_image_name|\
         IGconf_image_outputdir|\
         IGconf_sys_workroot|\
         IGconf_target_dir|\
         IGconf_target_path|\
         IGconf_sys_workroot|\
         IGconf_sys_cachedir|\
         DEB_BUILD_GNU_TYPE)
            declare -g "$key"="$val"
            ;;
      esac
   }
   # Set these variables in the shell to simply further processing
   mapfile_kv "${ctx[FINALENV]}" set_kv

   # Create output dirs
   install -d -m 0755 "$IGconf_target_dir" "$IGconf_sys_workroot" "$IGconf_sys_cachedir"
   [[ -n "${IGconf_image_outputdir:-}" ]] && install -d -m 0755 "$IGconf_image_outputdir"

   # Build required host tools
   bootstrap_build_tools

   # Layers can ship host executables
   PLAN="${ctx[LAYER_PLAN]}" foreach_layer_in_plan prepend_layer_paths
   export PATH

   # Seed bdebstrap env with sane defaults
   _bdebstrap=()
   _bdebstrap+=( --force --env PATH="$PATH" )
   _bdebstrap+=( --name     "${IGconf_image_name:-rpi-image-gen}" )
   _bdebstrap+=( --hostname "${IGconf_device_hostname:-rpi-image-gen}" )
   _bdebstrap+=( --output   "${IGconf_target_dir:?not set}" )
   _bdebstrap+=( --target   "${IGconf_target_path:?not set}" )

   # Process all options from the configuration and selectively set apt/dpkg opts
   # or other dynamically driven settings.
   process_conf_opt() {
      local key=${1:-} value=${2:-}
      msg "$key : $value"
      case $key in
         IGconf_sys_apt_keydir)
            xmkdir "$value"
            [[ -d /usr/share/keyrings ]] && rsync -a /usr/share/keyrings/ "$value"
            if [[ -n ${HOME:-} && -d ${HOME}/.local/share/keyrings ]]; then
               rsync -a "${HOME}/.local/share/keyrings/" "$value"
            fi
            rsync -a "${IGTOP}/keydir/" "$value"
            _bdebstrap+=( --aptopt "Dir::Etc::TrustedParts $value" ) ;;

         IGconf_sys_apt_cachedir)
            if [[ -n "$value" ]]; then
              local cache
              cache=$(realpath -e "$value" 2>/dev/null) || die "$value does not exist"
              local archives=/var/cache/apt/archives
              export _NS_APT_ARCHIVES="${IGconf_target_path}${archives}"

              local num=$(find "$cache" -maxdepth 1 -name '*.deb' 2>/dev/null | wc -l)
              msg "Apt cache: $num pkgs @ $cache"

              _bdebstrap+=( --skip=cleanup/apt/cache )
              _bdebstrap+=( --skip=cleanup/apt/lists )
              _bdebstrap+=( --skip=essential/unlink )
              _bdebstrap+=( --aptopt 'APT::Keep-Downloaded-Packages "true"' )
              _bdebstrap+=( --setup-hook   "mkdir -p \"\$1${archives}/partial\" && mount --bind '$cache' \"\$1${archives}\"" )
              _bdebstrap+=( --cleanup-hook "umount \"\$1${archives}\"" )
              value="$cache"
            fi
            ;;

         IGconf_sys_apt_proxy_http)
            curl --silent --head --max-time 2 "$value" >/dev/null \
               || die "$key proxy unreachable ($value)"
               _bdebstrap+=( --aptopt "Acquire::http { Proxy \"$value\"; }" ) ;;

         IGconf_sys_apt_get_purge)
            [[ ${value,,} =~ ^y(es)?$ ]] \
               && _bdebstrap+=( --aptopt 'APT::Get::Purge true' ) ;;
      esac

      # pass through
      _bdebstrap+=( --env "$key=$value" )
   }
   msg "\nFINAL ENV"
   mapfile_kv "${ctx[FINALENV]}" process_conf_opt

   msg "\nMMDEBSTRAP"
   local total=0 added=0
   local PLAN="${ctx[LAYER_PLAN]}"

   # Sanity check resolved layers from the plan
   _sanity_check_layer() {
      local layer=$1 resolved=$4
      [[ -n $resolved ]] || return 0
      ((total++))
      [[ -f $resolved ]] || die "Layer $layer ($resolved) not found"
   }
   foreach_layer_in_plan _sanity_check_layer

   # Layers declaring an mmdebstrap mapping contribute their own config
   local mmpaths
   mmpaths=$(mmdebstrap_layer_paths "${ctx[LAYER_PLAN]}") \
      || die "Failed to parse one or more YAML layers"

   local -A mmset=()
   local mmpath
   while IFS= read -r mmpath; do
      [[ -n $mmpath ]] && mmset[$mmpath]=1
   done <<< "$mmpaths"

   # Construct the bdebstrap ordered chain, including synthesised layers.
   # Every layer is offered pre/post synth, whether or not it has an mmdebstrap
   # mapping of its own, so that a config-only layer can still ship an overlay.
   # Use the per-layer workdir for pre/post synth, created but never wiped, so
   # build artefacts survive across builds. The exception is the registry's
   # auto-pruned sub-directories.
   _build_layer_config() {
      local layer=$1 version=$2 stempath=$3 resolved=$4 workdir=$5
      local cfg=()

      mkdir -p "$workdir"

      local prefile="${workdir}/pre.yaml" postfile="${workdir}/post.yaml"

      synth_layer_pre "$layer" "$version" "$stempath" "$prefile" \
         || die "Layer $layer ($stempath): pre-synth failed"
      [[ -f $prefile ]] && cfg+=( --config "$prefile" )

      [[ -n ${mmset[$resolved]:-} ]] && cfg+=( --config "$resolved" )

      synth_layer_post "$layer" "$version" "$stempath" "$postfile" \
         || die "Layer $layer ($stempath): post-synth failed"
      [[ -f $postfile ]] && cfg+=( --config "$postfile" )

      [[ ${#cfg[@]} -gt 0 ]] || return 0
      ((added++))

      _bdebstrap+=( "${cfg[@]}" )

      msg "Loaded $layer${version:+ ${version}}"
   }
   foreach_layer_in_plan _build_layer_config

   local skipped=$((total - added))
   msg "Loaded ${added}/${total}, skipped $skipped"

   # Store variable name so we can get it by nameref. ctx is an associative
   # array and bash can't store another (indexed) array inside it.
   ctx[ENV_BDEBSTRAP]='_bdebstrap'

   msg "\nREADY"
}


###############################################################################
# Stage 4: Filesystem generation
#   run bdebstrap
#   SBOM
###############################################################################
generate_filesystem()
{
   [[ ${ctx[INTERACTIVE]}  == y ]] && { ask "Generate filesystem?" y || exit 0; }

   msg "\nFILESYSTEM"

   # Auto-prune layer workdir sub-directories before building the filesystem
   local -a prune
   IFS=',' read -ra prune <<< "$(get_var IG_WORKDIR_PRUNE "${ctx[FINALENV]}")"

   _prune_layer_workdir() {
      local workdir=$5 sub
      for sub in "${prune[@]}"; do
         sub=${sub// /}
         rm -rf -- "${workdir:?}/${sub:?}"
      done
   }
   local PLAN="${ctx[LAYER_PLAN]}"
   foreach_layer_in_plan _prune_layer_workdir

   local -n _benv="${ctx[ENV_BDEBSTRAP]}" # via nameref

   runenv "${ctx[FINALENV]}" -s ns runner prebuild "${ctx[LAYER_PLAN]}"

   rund "${ctx[SRCROOT]}" -s ns bdebstrap \
      "${_benv[@]}" \
      --setup-hook     "runner setup     '${ctx[LAYER_PLAN]}' \"\$@\"" \
      --extract-hook   "runner extract   '${ctx[LAYER_PLAN]}' \"\$@\"" \
      --essential-hook "runner essential '${ctx[LAYER_PLAN]}' \"\$@\"" \
      --customize-hook "runner customize '${ctx[LAYER_PLAN]}' \"\$@\"" \
      --cleanup-hook   "runner cleanup   '${ctx[LAYER_PLAN]}' \"\$@\""

   runenv "${ctx[FINALENV]}" -s ns runner postbuild "${ctx[LAYER_PLAN]}"

   [[ ${ctx[INTERACTIVE]} == y ]] && \
      { ask "Filesystem complete. Proceed to SBOM?" y || exit 0 ; }

   msg "\nSBOM"
   runenv "${ctx[FINALENV]}" -s ns runner sbom "${ctx[LAYER_PLAN]}"
}



###############################################################################
# Stage 5: Image generation
#   preimage hooks
#   Invoke image provider
#   postimage hooks
###############################################################################
generate_images() {
   local provider=$(get_var IGconf_image_provider "${ctx[FINALENV]}")
   [[ -z "$provider" ]] && return 0

   [[ ${ctx[INTERACTIVE]} == y ]] && { ask "Generate image(s)?" y || exit 0 ; }

   local filesystem=$(get_var IGconf_target_path "${ctx[FINALENV]}") || die "no filesystem"
   local output_path=$(get_var IGconf_image_outputdir "${ctx[FINALENV]}") || die "no out path"

   msg "\nIMAGE"

   runenv "${ctx[FINALENV]}" -s ns runner preimage "${ctx[LAYER_PLAN]}"

   if [[ "$provider" == genimage && -d "$filesystem" ]] ; then
      mkdir -p "${TMPDIR}/genimage"
      local cfg
      for cfg in "${output_path}"/genimage*.cfg; do
         [[ -f $cfg ]] || continue
         runenv "${ctx[FINALENV]}" -s ns env genimage \
            --rootpath   "$filesystem" \
            --tmppath    "${TMPDIR}/genimage" \
            --inputpath  "$output_path" \
            --outputpath "$output_path" \
            --loglevel=1 \
            --config     "$cfg"  \
            | pv -t -F 'Generating image...%t' \
            || die "genimage error ($cfg)"
      done
   fi

   runenv "${ctx[FINALENV]}" -s ns runner postimage "${ctx[LAYER_PLAN]}"
}



###############################################################################
# Stage 6: Deploy
#   Install build assets for distribution
###############################################################################
deploy() {
   runenv "${ctx[FINALENV]}" -s ns runner finalize "${ctx[LAYER_PLAN]}"

   [[ ${ctx[INTERACTIVE]} == y ]] && { ask "Deploy assets?" y || exit 0 ; }

   msg "\nDEPLOY"
   runenv "${ctx[FINALENV]}" -s ns runner deploy "${ctx[LAYER_PLAN]}"
}



###############################################################################
# Clean
#   Output directories
#   Host build artefacts
###############################################################################
clean_worktree() {
   msg "\nCLEAN"
   : "${ctx[FINALENV]?FINALENV is not set}"

   local -a keys=(
      IGconf_target_path
      IGconf_image_outputdir
      IGconf_image_deploydir
      IGconf_sys_buildroot
      IGconf_sys_workroot
   )
   local key val
   for key in "${keys[@]}"; do
      val=$(get_var "$key" "${ctx[FINALENV]}") || continue
      case $key in
         IGconf_target_path)
            [[ -e $val ]] || continue
            ask "Remove $key [$val]?" y && run ns rm -rf -- "$val"
            ;;
         IGconf_sys_workroot)
            local gnutype
            gnutype=$(get_var DEB_BUILD_GNU_TYPE "${ctx[FINALENV]}") || continue
            val="${val}/${gnutype}"
            [[ -e $val ]] || continue
            ask "Remove host support packages [$val]?" y && rm -rf -- "$val"
            ;;
         *)
            [[ -e $val ]] || continue
            ask "Remove $key [$val]?" y && rm -rf -- "$val"
            ;;
      esac
   done
}



main() {
   run_stage() {
      local fn=$1; shift
      printf '\n\033[1m==> %s\033[0m\n' "$fn"
      "$fn" "$@" || die "Stage '$fn' failed"
   }

   case $cmd in
      build)
         run_stage parameter_assembly
         run_stage collect_layers
         run_stage prepare_build_config
         [[ "${ctx[ONLY_IMAGE]}" == y ]] || run_stage generate_filesystem
         [[ "${ctx[ONLY_FS]}" == y ]] || run_stage generate_images
         run_stage deploy
         ;;

      clean)
         run_stage parameter_assembly
         run_stage collect_layers
         run_stage clean_worktree
         ;;
   esac
}


main "$@"
