use str
use re
use platform

var in-ssh = (has-env SSH_CONNECTION)
var hostname = (platform:hostname)

fn nix-shell-prompt-string {
  put '['
  styled nix yellow
  if (has-env NIX_SHELL_NAME) {
    put ':'
    styled $E:NIX_SHELL_NAME green
  }
  put ']'
}

fn k8s-prompt-string {
  var k8s-config = (kubectl config view -o json | from-json)
  var k8s-context = $k8s-config[current-context]
  var k8s-namespace = (each { |ctx| if (eq $ctx[name] $k8s-context) { put $ctx[context][namespace] } } $k8s-config[contexts])
  put '['
  styled $k8s-context magenta
  put '|'
  styled $k8s-namespace blue
  put ']'
}

fn aws-prompt-string {
  put '['
  styled $E:AWS_PROFILE yellow
  put ']'
}

fn parse-git-branch {
  try {
    git symbolic-ref -q HEAD 2> /dev/null
  } catch e {
    try {
      git rev-parse --short HEAD 2> /dev/null
    } catch e {
      echo ""
    }
  }
}

fn resolve-symbolic {|ref|
  try {
    put (git rev-parse --symbolic-full-name $ref 2> /dev/null)
  } catch {
    put $nil
  }
}

fn parse-git-state {
  var git-revision = ''
  var git-branch = ''
  var git-changed = (num 0)
  var git-staged = (num 0)
  var git-conflicts = (num 0)
  var git-untracked = (num 0)
  var git-ahead-upstream = (num 0)
  var git-behind-upstream = (num 0)
  var git-ahead-push = (num 0)
  var git-behind-push = (num 0)

  git status --porcelain=v2 -b 2>/dev/null | each {|line|
    if (eq '#' $line[0]) {
      try {
        var match = (re:find '# branch.ab \+(\d+) -(\d+)' $line)
        set git-ahead-upstream = (num $match[groups][1][text])
        set git-behind-upstream = (num $match[groups][2][text])
      } catch e {
        # noop
      }

      if (str:has-prefix $line '# branch.head') {
        set git-branch = $line[14..]
      }

      if (str:has-prefix $line '# branch.oid') {
        set git-revision = $line[13..20]
      }

    } elif (eq u $line[0]) {
      set git-conflicts = (+ $git-conflicts 1)
    } elif (eq '?' $line[0]) {
      set git-untracked = (+ $git-untracked 1)
    } else {
      if (not-eq '.' $line[2]) {
        set git-staged = (+ $git-staged 1)
      }
      if (not-eq '.' $line[3]) {
        set git-changed = (+ $git-changed 1)
      }
    }
  }

  var git-ref = $git-branch
  var git-ref-prefix = ''
  if (eq $git-branch '(detached)') {
    try {
      var tag @rest = (git tag --points-at HEAD)
      set git-ref-prefix = '#'
      set git-ref = $tag
    } catch e {
      set git-ref-prefix = '@'
      set git-ref = $git-revision
    }
  } else {
    var git-push = (resolve-symbolic '@{push}')
    var git-upstream = (resolve-symbolic '@{upstream}')
    var git-origin-head = (resolve-symbolic origin/HEAD)

    if (not-eq $nil $git-push) {
      if (eq $git-push $git-upstream) {
        # if push is the same as upstream, use the values from git status
        set git-behind-push = $git-behind-upstream
        set git-ahead-push = $git-ahead-upstream
      } else {
        # else calculate it ourselves
        var behind ahead = (str:split "\t" (git rev-list --count --left-right '@{push}...HEAD'))
        set git-behind-push = (num $behind)
        set git-ahead-push = (num $ahead)
      }
    }

    # if upstream is nil or the same as push, use origin/HEAD instead if
    # present
    if (and ^
      (or (eq $git-push $git-upstream) (eq $nil $git-upstream)) ^
      (not-eq $nil $git-origin-head) ^
    ) {
      var behind ahead = (str:split "\t" (git rev-list --count --left-right origin/HEAD...HEAD))
      set git-behind-upstream = (num $behind)
      set git-ahead-upstream = (num $ahead)
    }
  }

  put '('
  put $git-ref-prefix
  styled $git-ref magenta bold
  if (> $git-ahead-push 0) {
    put '^'$git-ahead-push
  }
  if (> $git-behind-push 0) {
    put 'v'$git-behind-push
  }
  if (> $git-behind-upstream 0) {
    put '!'
  }
  put '|'
  if (> $git-staged 0) {
    styled +$git-staged green
  }
  if (> $git-conflicts 0) {
    styled x$git-conflicts red
  }
  if (> $git-changed 0) {
    styled '*'$git-changed red
  }
  if (> $git-untracked 0) {
    put '..'
  }
  if (== 0 $git-changed $git-staged $git-conflicts $git-untracked) {
    styled ok green bold
  }
  put ')'
}

fn git-prompt-string {
  try {
    parse-git-state
  } catch e {
  }
}

set edit:prompt = {
  put "\n"
  styled '['(date +%H:%M)']' bold bright-cyan
  put ' '
  if $in-ssh {
    put '('
    styled $hostname bright-blue
    put ')'
    put ' '
  }
  styled (tilde-abbr $pwd) yellow
  put ' '
  git-prompt-string
  put "\n% "
}

set edit:rprompt = {
  put '['
  styled elv green
  put ']'
  if (has-env IN_NIX_SHELL) {
    nix-shell-prompt-string
  }
  if (has-env IN_K8S) {
    k8s-prompt-string
  }
  if (has-env AWS_PROFILE) {
    aws-prompt-string
  }
}

# when stale replace clock with --:--
set edit:prompt-stale-transform = {|x|
  put $x[0] # newline
  styled '[--:--]' bright-cyan
  put $x[2..]
}