#!/usr/bin/env bash
#
# A direnv hook for loading Nix flake-based development environments
# with built-in caching and garbage-collection root management.
#
# Based on
# https://github.com/nix-community/nix-direnv
# and
# https://github.com/paisano-nix/direnv
#
set -euo pipefail

readonly RENSA_DIRENV_VERSION="0.1.0"
readonly DIRENV_MIN_VERSION="2.21.3"
readonly BASH_MIN_VERSION="4.4"

#
# CONFIGURATION ENV VARS
#
REN_FLAKE_ATTR=${REN_FLAKE_ATTR-"__ren"}
# whether to enable watching the files for changes
REN_DO_WATCH=${REN_DO_WATCH-true}
# whether to add the shell profile or flake inputs to gcroots
REN_DO_GCROOTS=${REN_DO_GCROOTS-true}
# whether to archive the flake inputs and gcroot them (only works if REN_DO_GCROOT is true)
REN_DO_ARCHIVE=${REN_DO_ARCHIVE-true}
# whether to enable compat mode which sets PRJ env variables
REN_DO_PRJ_COMPAT=${REN_DO_PRJ_COMAT-true}

#
# LOGGING
#

# Usage: __ren_log <LEVEL> "message"
# LEVEL can be: TRACE, INFO, WARN, ERROR
__ren_log() {
  local level="$1" msg="$2"
  local format=$'\E[mren: \E[38;5;8m%s\E[m'
  case "$level" in
  TRACE)
    if [[ -n "${REN_TRACE:-}" ]]; then
      # shellcheck disable=SC2059
      log_status "$(printf "$format" "[TRACE] ${msg}")"
    fi
    ;;
  INFO)
    # shellcheck disable=SC2059
    log_status "$(printf "$format" "${msg}")"
    ;;
  WARN)
    # shellcheck disable=SC2059
    log_error "$(printf "$format" "[WARN] ${msg}")"
    ;;
  ERROR)
    # shellcheck disable=SC2059
    log_error "$(printf "$format" "[ERROR] ${msg}")"
    ;;
  esac
}

#
# UTILS
#

__ren_require_version() {
  local name="$1" version="$2" required="$3"
  if ! printf "%s\n%s\n" "$required" "$version" | LC_ALL=C sort --check --version-sort &>/dev/null; then
    __ren_log ERROR "Minimum required ${name} version is ${required} (found: ${version})."
    return 1
  fi
}

__ren_require_cmd_version() {
  local cmd="$1" required="$2" version
  if ! has "$cmd"; then
    __ren_log ERROR "Command not found: $cmd"
    return 1
  fi
  version=$($cmd --version)
  if [[ $version =~ ([0-9]+\.[0-9]+\.[0-9]+) ]]; then
    __ren_require_version "$cmd" "${BASH_REMATCH[1]}" "$required"
  else
    __ren_log WARN "Could not parse version for '$cmd'. Skipping version check."
  fi
}

__cell_path_cache=""
__ren_get_cell_path() {
  if [[ -z "$__cell_path_cache" ]]; then
    local cell="$1"
    local cellsFrom
    cellsFrom=$(__ren_nix eval --raw ".#${REN_FLAKE_ATTR}.cellsFrom" 2>/dev/null || true)
    if [[ -z "$cellsFrom" ]]; then
      __ren_log ERROR "Could not determine 'cellsFrom' path from flake."
      return 1
    fi
    __cell_path_cache=$cellsFrom
  fi
  echo "${REN_ROOT}/${__cell_path_cache}/${cell}"
}

# set in the entrypoint
__nix_cmd=""
__ren_nix() {
  "$__nix_cmd" --no-warn-dirty --extra-experimental-features "nix-command flakes" "$@"
}

#
# INIT
#

# discovers the project root (git repo) and sets up state directories and variables
__ren_init_project() {
  [[ -n "${REN_ROOT:-}" ]] && return 0
  local git_repo
  git_repo=$(git rev-parse --show-toplevel 2>/dev/null || true)
  if [[ -z "$git_repo" ]]; then
    __ren_log ERROR "Not inside a git repository. Cannot determine project root."
    return 1
  fi

  export REN_ROOT="$git_repo"
  export REN_STATE="${REN_ROOT}/.ren"
  mkdir -p "${REN_STATE}"

  # shellcheck disable=SC2034
  direnv_layout_dir="${REN_STATE}/direnv"
  mkdir -p "${REN_STATE}/direnv"

  # add nested gitignore which ignores the whole state dir
  echo "**/*" >"$REN_STATE/.gitignore"

  if [ "$REN_DO_PRJ_COMPAT" = true ]; then
    __ren_log INFO "Setting PRJ env variables for compat"
    __ren_init_prj_spec
  else
    __ren_log TRACE "Not setting PRJ env variables"
  fi
}

# compat with other stuff that uses PRJ_...
# we explicitly don't want any dirs to be in the project dir, put them all
#  in .ren, so they don't clutter everything up
__ren_init_prj_spec() {
  export PRJ_ROOT=${PRJ_ROOT:="${REN_ROOT}"}

  export PRJ_CONFIG_HOME=${PRJ_CONFIG_HOME:="${REN_STATE}/config"}
  mkdir -p "${PRJ_CONFIG_HOME}"

  export PRJ_RUNTIME_DIR=${PRJ_RUNTIME_DIR:="${REN_STATE}/runtime"}
  mkdir -p "${PRJ_RUNTIME_DIR}"

  # load project id if exists
  if [[ -z "${PRJ_ID:-}" && -f "${PRJ_CONFIG_HOME}/prj_id" ]]; then
    export PRJ_ID=$(<"${PRJ_CONFIG_HOME}/prj_id")
  fi

  # PRJ_CACHE_HOME - shared if PRJ_ID is set
  if [[ -z "${PRJ_CACHE_HOME:-}" ]]; then
    if [[ -n "${PRJ_ID:-}" ]]; then
      export PRJ_CACHE_HOME="${XDG_CACHE_HOME}/prj/${PRJ_ID}"
    else
      export PRJ_CACHE_HOME="${REN_STATE}/cache"
    fi
  fi
  mkdir -p "${PRJ_CACHE_HOME}"

  # PRJ_DATA_HOME - shared if PRJ_ID is set
  if [[ -z "${PRJ_DATA_HOME:-}" ]]; then
    if [[ -n "${PRJ_ID:-}" ]]; then
      export PRJ_DATA_HOME="${XDG_DATA_HOME}/prj/${PRJ_ID}"
    else
      export PRJ_DATA_HOME="${REN_STATE}/data"
    fi
  fi
  # NOTE: stuff like devshell uses PRJ_DATA_DIR instead, even though numtide
  # wrote the spec? :D
  # technically not supported so we just overwrite it, use PRJ_DATA_HOME instead
  export PRJ_DATA_DIR=$PRJ_DATA_HOME
  mkdir -p "${PRJ_DATA_HOME}"

  # PRJ_PATH - shared if PRJ_ID is set
  if [[ -z "${PRJ_PATH:-}" ]]; then
    if [[ -n "${PRJ_ID:-}" ]]; then
      export PRJ_PATH="${HOME}/.local/bin/prj/${PRJ_ID}"
    else
      export PRJ_PATH="${REN_STATE}/bin"
    fi
  fi
  mkdir -p "${PRJ_PATH}"
}

#
# NIX GCROOTS & ENV STUFF
#

__ren_import_env() {
  local profile_rc="$1"
  [[ -f "$profile_rc" ]] || return 0
  local -A to_restore
  to_restore=(
    ["NIX_BUILD_TOP"]="${NIX_BUILD_TOP:-__UNSET__}"
    ["TMP"]="${TMP:-__UNSET__}"
    ["TMPDIR"]="${TMPDIR:-__UNSET__}"
    ["TEMP"]="${TEMP:-__UNSET__}"
    ["TEMPDIR"]="${TEMPDIR:-__UNSET__}"
    ["terminfo"]="${terminfo:-__UNSET__}"
  )
  local old_xdg_data_dirs="${XDG_DATA_DIRS:-}"

  # shellcheck source=/dev/null
  eval "$(<"$profile_rc")"

  if [[ -n "${NIX_BUILD_TOP+x}" && $NIX_BUILD_TOP == */nix-shell.* && -d $NIX_BUILD_TOP ]]; then rm -rf "$NIX_BUILD_TOP"; fi
  for key in "${!to_restore[@]}"; do
    if [[ "${to_restore[$key]}" == "__UNSET__" ]]; then unset "$key"; else export "$key=${to_restore[$key]}"; fi
  done
  local new_xdg_data_dirs="${XDG_DATA_DIRS:-}"
  export XDG_DATA_DIRS=""
  local IFS=:
  for dir in $new_xdg_data_dirs${old_xdg_data_dirs:+:}$old_xdg_data_dirs; do
    dir="${dir%/}"
    [[ -z "$dir" || ":$XDG_DATA_DIRS:" == *:"$dir":* ]] && continue
    XDG_DATA_DIRS="${XDG_DATA_DIRS}${XDG_DATA_DIRS:+:}$dir"
  done
}

__ren_add_gcroot() {
  local storepath="$1" symlink="$2"
  __ren_log TRACE "Adding GC root: $symlink -> $storepath"
  __ren_nix build --out-link "$symlink" "$storepath" >/dev/null
}

__ren_clean_old_gcroots() {
  __ren_log TRACE "Cleaning old GC roots in $1"
  rm -rf "${1}/flake-inputs/" "${1}/profile-"*
}

__ren_argsum_suffix() {
  local out checksum
  if [[ -n "$1" ]]; then
    if has sha1sum; then
      out=$(sha1sum <<<"$1")
    elif has shasum; then
      out=$(shasum <<<"$1")
    else
      __ren_log WARN "sha1sum/shasum not found; cannot create stable profile name."
      return
    fi
    read -r checksum _ <<<"$out"
    echo "-$checksum"
  fi
}

# archive flake inputs and gcroot them
__ren_add_flake_input_gcroots() {
  local flake_dir="$1" __layout_dir="$2"
  __ren_log TRACE "Archiving flake inputs from '$flake_dir'"

  local flake_inputs_dir="${__layout_dir}/flake-inputs/"
  mkdir -p "$flake_inputs_dir"

  local flake_inputs_json
  # run in subshell to not affect the current directory and handle errors gracefully
  flake_inputs_json=$( (cd "$flake_dir" && __ren_nix flake archive --json --no-write-lock-file -- ".") 2>/dev/null || true)

  if [[ -z "$flake_inputs_json" ]]; then
    __ren_log TRACE "No inputs found or error archiving flake at '$flake_dir'. Skipping."
    return
  fi

  while [[ $flake_inputs_json =~ /nix/store/[^\"]+ ]]; do
    local store_path="${BASH_REMATCH[0]}"
    [[ -z "$store_path" ]] && continue
    __ren_add_gcroot "$store_path" "${flake_inputs_dir}/${store_path##*/}"
    flake_inputs_json="${flake_inputs_json/${store_path}/}"
  done
}

# updates gcroots
__ren_update_gcroots() {
  local __layout_dir="$1" tmp_profile="$2" profile="$3" cell="$4"
  __ren_add_gcroot "$tmp_profile" "$profile"
  rm -f "$tmp_profile"*
  if [ "$REN_DO_ARCHIVE" = true ]; then
    # add gcroots for the main flake
    __ren_add_flake_input_gcroots "$REN_ROOT" "$__layout_dir"

    # if a cell is specified, also gcroot it's flake inputs, if it has a flake
    if [[ -n "$cell" ]]; then
      local cell_path
      if cell_path=$(__ren_get_cell_path "$cell"); then
        if [[ -f "${cell_path}/flake.nix" ]]; then
          __ren_log TRACE "Found cell flake. Adding GC roots for inputs from '${cell_path}'."
          __ren_add_flake_input_gcroots "$cell_path" "$__layout_dir"
        fi
      fi
    fi
  fi
}

#
# WATCHERS
#

__ren_get_direnv_watches() {
  local -n _watches_ref=$1
  [[ -z "${DIRENV_WATCHES-}" ]] && return 0
  # shellcheck disable=SC2154
  while IFS= read -r line; do
    local regex='"[Pp]ath": "(.+)"$'
    if [[ $line =~ $regex ]]; then
      local path
      # shellcheck disable=SC2059
      path=$(printf "${BASH_REMATCH[1]}")
      if [[ $path != "${XDG_DATA_HOME:-${HOME:-/var/empty}/.local/share}/direnv/allow/"* ]]; then
        _watches_ref+=("$path")
      fi
    fi
  done < <("$direnv" show_dump "${DIRENV_WATCHES}")
}

__ren_watch_cell() {
  local target_spec="$1"
  read -r cell block _ <<<"${target_spec//\// }"
  local cell_path
  if ! cell_path=$(__ren_get_cell_path "$cell"); then
    __ren_log WARN "Could not determine 'cellsFrom' path from flake. File watching for target files will be incomplete."
    return
  fi

  if [[ -f "${cell_path}/${block}.nix" ]]; then
    __ren_log TRACE "Watching: ${cell_path}/${block}.nix"
    watch_file "${cell_path}/${block}.nix"
  fi
  if [[ -f "${cell_path}/default.nix" ]]; then
    __ren_log TRACE "Watching: ${cell_path}/default.nix"
    watch_file "${cell_path}/default.nix"
  fi
  if [[ -d "${cell_path}/${block}" ]]; then
    __ren_log TRACE "Watching dir: ${cell_path}/${block}"
    watch_dir "${cell_path}/${block}"
  fi

  # watch for cell-level flake.nix and flake.lock
  if [[ -f "${cell_path}/flake.nix" ]]; then
    __ren_log TRACE "Watching cell flake: ${cell_path}/flake.nix"
    watch_file "${cell_path}/flake.nix"
  fi
  if [[ -f "${cell_path}/flake.lock" ]]; then
    __ren_log TRACE "Watching cell flake lock: ${cell_path}/flake.lock"
    watch_file "${cell_path}/flake.lock"
  fi
}

__ren_setup_watches() {
  watch_file "$REN_ROOT/flake.nix" "$REN_ROOT/flake.lock"
  for target in "$@"; do
    if [[ "$target" != *"#"* ]]; then
      __ren_watch_cell "$target"
    fi
  done
}

#
# BUILD & CACHE
#

# returns 1 if rebuild is needed
__ren_is_rebuild_needed() {
  local profile_rc="$1"
  if [[ ! -f "$profile_rc" ]]; then
    __ren_log TRACE "No cached profile found at '$profile_rc'. Rebuilding."
    echo 1
    return
  fi

  if [ "$REN_DO_WATCH" = true ]; then
    local watches=()
    __ren_get_direnv_watches watches
    for file in "${watches[@]}"; do
      if [[ "$file" -nt "$profile_rc" ]]; then
        __ren_log INFO "Cache invalidated by: $file"
        echo 1
        return
      fi
    done
  fi
  echo 0
}

# returns 0 on success, 1 on failure
__ren_build_and_cache() {
  local target_spec="$1" __layout_dir="$2" profile="$3" profile_rc="$4"
  __ren_log INFO "Building shell from '$target_spec'..."
  local flake_attr cell=""
  # allows specifying either flake attr like `./path#devShells.default` or through
  # ren style `//repo/devShells/default`
  if [[ "$target_spec" == *"#"* ]]; then
    flake_attr="$target_spec"
    __ren_log TRACE "Using direct flake attribute: $flake_attr"
  else
    read -r cell block target <<<"${target_spec//\// }"
    local system
    system=$(__ren_nix eval --raw --impure --expr builtins.currentSystem)
    flake_attr=".#${system}.${cell}.${block}.${target}"
  fi
  local tmp_profile="${__layout_dir}/tmp-profile.$$"
  if build_output=$(__ren_nix print-dev-env --profile "$tmp_profile" "$flake_attr"); then
    __ren_clean_old_gcroots "$__layout_dir"
    echo "$build_output" >"$profile_rc"

    # only add new gcroots if enabled, cleaning old is always fine
    if [ "$REN_DO_GCROOTS" = true ]; then
      __ren_update_gcroots "$__layout_dir" "$tmp_profile" "$profile" "$cell"
    fi

    __ren_log TRACE "Cache for '$target_spec' renewed successfully."
    return 0
  else
    rm -f "$tmp_profile"*
    return 1
  fi
}

#
# ENTRYPOINTS
#

# usage in .envrc: use ren <target>
# <target> can be `//cell/block/target` or a direct flake attr like `.#myShell`
use_ren() {
  if [[ -z ${REN_SKIP_VERSION_CHECK:-} ]]; then
    __ren_require_version "bash" "$BASH_VERSION" "$BASH_MIN_VERSION"
    __ren_require_cmd_version "direnv" "$DIRENV_MIN_VERSION"
  fi

  if command -v nix &>/dev/null; then
    __nix_cmd=$(command -v nix)
  elif [[ -n "${REN_FALLBACK_NIX:-}" ]]; then
    __nix_cmd="${REN_FALLBACK_NIX}"
  else
    __ren_log ERROR "Could not find Nix. Add 'nix' to PATH or set REN_FALLBACK_NIX."
    return 1
  fi

  local target_spec="${1-}"
  if [[ -z "$target_spec" ]]; then
    __ren_log ERROR "use ren requires a target argument, e.g., //repo/devShells/default or .#myShell"
    return 1
  fi

  if [ "$REN_DO_WATCH" = true ]; then
    __ren_setup_watches "$@"
  fi

  local __layout_dir profile profile_rc
  __layout_dir=$(direnv_layout_dir)
  profile="${__layout_dir}/profile$(__ren_argsum_suffix "$target_spec")"
  profile_rc="${profile}.rc"

  if [[ $(__ren_is_rebuild_needed "$profile_rc") -eq 1 ]]; then
    if ! __ren_build_and_cache "$target_spec" "$__layout_dir" "$profile" "$profile_rc"; then
      if [[ -f "$profile_rc" ]]; then
        __ren_log WARN "Nix evaluation failed. Falling back to the last known-good environment."
        export REN_DID_FALLBACK=1
      else
        __ren_log ERROR "Nix evaluation failed, and no cached environment is available."
        return 1
      fi
    fi
  else
    __ren_log INFO "Using cached environment for '$target_spec'."
  fi

  __ren_import_env "$profile_rc"
}

use_rensa() {
  use_ren "$@"
}

# initialize project environment variables immediately upon sourcing the script,
# allows the user to use $REN_STATE etc. in .envrc
__ren_init_project
__ren_log TRACE "Rensa v${RENSA_DIRENV_VERSION}"

#
# HELPERS
# to use in .envrc, eg. to switch between shells etc.
#

# read state from the project
read_state() {
  if [[ -f "$REN_STATE/$1" ]]; then
    cat "$REN_STATE/$1"
  fi
}

# set/write state for the project
write_state() {
  echo "$2" >"$REN_STATE/$1"
}