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 quiet = [
  AR
  AR_FOR_TARGET
  AS
  AS_FOR_TARGET
  CC
  CC_FOR_TARGET
  CXX
  CXX_FOR_TARGET
  LD
  LD_FOR_TARGET
  NM
  NM_FOR_TARGET
  RANLIB
  RANLIB_FOR_TARGET
  SIZE
  SIZE_FOR_TARGET
  STRINGS
  STRINGS_FOR_TARGET
  STRIP
  STRIP_FOR_TARGET

  READELF
  OBJCOPY
  OBJDUMP

  MACOSX_DEPLOYMENT_TARGET

  CONFIG_SHELL
  HOST_PATH
]
var ignore = [
  # keep user variables
  HOME
  OLDPWD
  SHELL
  TEMP
  TEMPDIR
  TERM
  TMP
  TMPDIR
  TZ

  # purity-related envs
  NIX_ENFORCE_PURITY
  SOURCE_DATE_EPOCH
  ZERO_AR_DATE

  # nix build variables
  buildInputs
  buildPhase
  builder
  cmakeFlags
  configureFlags
  depsBuildBuild
  depsBuildBuildPropagated
  depsBuildTarget
  depsBuildTargetPropagated
  depsHostHost
  depsHostHostPropagated
  depsTargetTarget
  depsTargetTargetPropagated
  doCheck
  doInstallCheck
  dontAddDisableDepTrack
  mesonFlags
  name
  nativeBuildInputs
  out
  outputs
  patches
  phases
  preferLocalBuild
  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 print-change {|key new|
  if (has-value $quiet $key) {
    return
  }

  if (str:has-prefix $key "NIX_") {
    return
  }

  if $new {
    put +$key
  } else {
    put -$key
  }
}

fn apply-env {|env|
  each {|row|
    var key value = (all $row)
    if (eq $value $nil) {
      print-change $key $false
      unset-env $key
    } else {
      print-change $key $true
      set-env $key $value
    }
  } $env | compact
}

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 "__") {
        continue
      } elif (has-value $ignore $key) {
        continue
      } elif (has-value $extend $key) {
        var extended = $v[value]
        if (eq $extended "") {
          continue
        }
        var existing = (get-env $key)
        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]