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 {|@a|
  var dir = $pwd
  if (> (count $a) 0) {
    set dir = $a[0]
  }

  if (path:is-regular $dir"/flake.nix") {
    put "nix"
  } elif (path:is-regular $dir"/shell.nix") {
    put "nix"
  } elif (path:is-regular $dir"/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
  LD_DYLD_PATH
  MACOSX_DEPLOYMENT_TARGET
  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 (str:has-prefix $key "__") {
        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 path|
  var env-diff = $nil

  if (not-eq $state-file $nil) {
    var env = []
    if (eq $type "nix") {
      set @env = (env-from-nix-flake $state-file)
    }

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

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

  var entry = [&dir=$path &env-diff=$env-diff]
  set dir-stack = [ $entry $@dir-stack ]
}

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

  if (not-eq $last[env-diff] $nil) {
    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 ]
      }

      # remove any extra links
      find $state-dir -name $hash'-*' -delete
      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 $pwd
  } 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 last = "/"
  if (> (count $dir-stack) 0) {
    set last = $dir-stack[0][dir]
  }

  var current = $pwd

  var candidates = []
  while (and (str:has-prefix $current $last) (not-eq $current $last)) {
    set candidates = [ $current $@candidates ]

    if (eq $current "/") {
      # put this at the end so it can run against `/` once
      break
    }
    set current = (path:dir $current)
  }

  for candidate $candidates {
    var type = (state-type $candidate)
    if (not-eq $type $nil) {
      var hash = (get-hash $candidate)
      var state-file = $state-dir/$hash'-'$type

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

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