use path
use builtin
use str

fn get-env {|name &default=$nil|
  if (has-env $name) {
    builtin:get-env $name
  } else {
    put $default
  }
}

fn get-hash {|@a|
  var dir = $pwd
  if (> (count $a) 0) {
    set dir = $a[0]
  }
  str:trim-right (echo $dir | sha256sum) "- "
}

fn state-type {
  if (path:is-regular "flake.nix") {
    put "nix"
  } elif (path:is-regular "shell.nix") {
    put "nix"
  } elif (path:is-regular "default.nix") {
    put "nix"
  } else {
    put $nil
  }
}

var state-dir = (get-env XDG_STATE_HOME &default=~/.local/state)/direlv
var dir-stack = []
var ignore = [
  AR
  AS
  CC
  CONFIG_SHELL
  CXX
  HOME
  HOST_PATH
  LD
  NM
  OBJCOPY
  OBJDUMP
  OLDPWD
  RANLIB
  READELF
  SHELL
  SIZE
  SOURCE_DATE_EPOCH
  STRINGS
  STRIP
  TEMP
  TEMPDIR
  TERM
  TMP
  TMPDIR
  TZ
  buildInputs
  buildPhase
  builder
  cmakeFlags
  configureFlags
  depsBuildBuild
  depsBuildBuildPropagated
  depsBuildTarget
  depsBuildTargetPropagated
  depsHostHost
  depsHostHostPropagated
  depsTargetTarget
  depsTargetTargetPropagated
  doCheck
  doInstallCheck
  dontAddDisableDepTrack
  mesonFlags
  name
  nativeBuildInputs
  out
  outputs
  patches
  phases
  propagatedBuildInputs
  propagatedNativeBuildInputs
  shell
  shellHook
  src
  stdenv
  strictDeps
  system
]
var extend = [
  PATH
  XDG_DATA_DIRS
]

fn diff-env {|env|
  each {|row|
    var key value = (all $row)
    put [$key (get-env $key)]
  } $env
}

fn apply-env {|env|
  each {|row|
    var key value = (all $row)
    if (eq $value $nil) {
      put -$key
      unset-env $key
    } else {
      put +$key
      set-env $key $value
    }
  } $env
}

fn env-from-nix-flake {|state-file|
  var data = (from-json < $state-file)
  keys $data[variables] | each {|key|
    var v = $data[variables][$key]
    if (eq $v[type] "exported") {
      if (str:has-prefix $key "NIX_") {
        continue
      } elif (has-value $ignore $key) {
        continue
      } elif (has-value $extend $key) {
        var existing = (get-env $key)
        var extended = $v[value]
        if (not-eq $existing $nil) {
          set extended = $extended":"$existing
        }
        put [$key $extended]
      } else {
        put [$key $v[value]]
      }
    }
  }
}

fn push-dir {|type state-file|
  var env = []
  if (eq $type "nix") {
    set @env = (env-from-nix-flake $state-file)
  }

  var @env-diff = (diff-env $env)

  var entry = [&dir=$pwd &env-diff=$env-diff]

  edit:notify "Entered direlv "$pwd
  edit:notify (print (apply-env $env))

  set dir-stack = [ $entry $@dir-stack ]
}

fn pop-dir {
  var last = $dir-stack[0]

  edit:notify "Exited direlv "$last[dir]
  edit:notify (print (apply-env $last[env-diff]))

  set dir-stack = $dir-stack[1..]
}

fn current-dir {
  if (> (count $dir-stack) 0) {
    put $dir-stack[0][dir]
  } else {
    put $nil
  }
}

edit:add-var direlv~ {|@a|
  if (eq (count $a) 0) {
    echo "direlv <cache|clear>"
    return
  } elif (eq $a[0] "cache") {
    var type = (state-type)
    if (eq $type $nil) {
      fail "No flake.nix or shell.nix found"
    }

    mkdir -p $state-dir
    var hash = (get-hash)
    var state-file = $state-dir/$hash'-'$type

    if (eq $type "nix") {
      var args = [
        "--print-build-logs"
        "--profile"
        $state-file
        "--command"
        "true"
      ]
      if (path:is-regular "flake.nix") {
        # noop
      } elif (path:is-regular "shell.nix") {
        set args = [ "--file" "shell.nix" $@args ]
      } elif (path:is-regular "default.nix") {
        set args = [ "--file" "default.nix" $@args ]
      }

      e:nix develop $@args
    }

    # if we're already in a direlv, unload first
    var dir = (current-dir)
    if (eq $dir $pwd) {
      pop-dir
    }
    push-dir $type $state-file
  } elif (eq $a[0] "clear") {
    var dir = (current-dir)
    if (not-eq $dir $nil) {
      var hash = (get-hash $dir)
      find $state-dir -name $hash'*' -delete
      pop-dir
    }
  }
}

fn on-chdir {|dir|
  while (> (count $dir-stack) 0) {
    if (eq $pwd $dir-stack[0][dir]) {
      return
    } elif (str:has-prefix $pwd $dir-stack[0][dir]) {
      break
    } else {
      pop-dir
    }
  }

  var type = (state-type)
  if (eq $type $nil) {
    return
  }

  var hash = (get-hash)
  var state-file = $state-dir/$hash'-'$type

  if (path:is-regular $state-file &follow-symlink=$true) {
    push-dir $type $state-file
  }
}

set after-chdir = [$on-chdir~]