# Fast prompt. Git is read from files; jj uses --ignore-working-copy and never snapshots.

const ICON_NOTES = (char -u "03C0")
const ICON_WORKSPACE = (char -u "03BB")
const ICON_HOME = (char home)
const ICON_GIT = (char nf_git)
const ICON_BRANCH = (char nf_branch)
const ICON_CHECK = (char -u "2713")
const ICON_FAILED = (char failed)
const ICON_PIPE = (char pipe)
const ICON_PROMPT = (char prompt)
const ICON_JJ = (char -u "f329") # nf-linux-parrot
const ICON_JJ_AT = "@"
const ICON_UNKNOWN = "?"
const FLAG_EMPTY = "∅"
const FLAG_CONFLICT = "×"
const FLAG_DIVERGENT = "≈"

const ANSI_DARK_GRAY = (ansi dark_gray)
const ANSI_DARK_GRAY_DIMMED = (ansi dark_gray_dimmed)
const ANSI_GREEN = (ansi green)
const ANSI_MAGENTA = (ansi magenta)
const ANSI_RED = (ansi red)
const ANSI_RESET = (ansi reset)

const PATH_SEP = (char path_sep)
const TRUNCATE_DEPTH = 5
const TRUNCATE_SYMBOL = $"…($PATH_SEP)"
const NOTES_DIR = "Notes"
const WORKSPACE_DIR = "Workspace"

const VCS_GIT = "git"
const VCS_JJ = "jj"
const DOT_GIT = ".git"
const DOT_JJ = ".jj"
const HEAD_FILE = "HEAD"
const OBJECTS_DIR = "objects"
const REFS_DIR = "refs"
const GITDIR_PREFIX = "gitdir: "
const REF_HEADS_PREFIX = "ref: refs/heads/"
const DETACHED_LABEL = "HEAD"

const JJ_TEMPLATE = '"{\"change\":" ++ change_id.short().escape_json() ++ ",\"bookmarks\":" ++ stringify(local_bookmarks.join(",")).escape_json() ++ ",\"empty\":" ++ json(empty) ++ ",\"conflict\":" ++ json(conflict) ++ ",\"divergent\":" ++ json(divergent) ++ "}\n"'

const PROMPT_INDICATOR_TEXT = $"($ANSI_GREEN) ⚈  ($ANSI_RESET)"
const PROMPT_INSERT_TEXT = $"($ANSI_MAGENTA) ($ICON_PIPE) ($ANSI_RESET)"
const PROMPT_NORMAL_TEXT = $"($ANSI_MAGENTA) ($ICON_PROMPT) ($ANSI_RESET)"
const PROMPT_MULTILINE_TEXT = $"($ANSI_MAGENTA) ($ICON_PROMPT)($ICON_PROMPT) ($ANSI_RESET)"

# Expand `~`; leave absolute/relative paths to `path join`/`path expand`.
def expand-home [path: string] {
    if ($path | str starts-with "~") { $path | path expand } else { $path }
}

# Resolve gitdir paths from `.git` files and `$GIT_DIR`.
def resolve-git-path [path: string, base: path] {
    $base | path join (expand-home $path) | path expand
}

def git-dir-from-file [dotgit: path, base: path] {
    let line = (try { open --raw $dotgit | str trim } catch { return null })
    if ($line | str starts-with $GITDIR_PREFIX) {
        resolve-git-path ($line | str replace $GITDIR_PREFIX "") $base
    }
}

def safe-path-type [value: any] {
    if $value == null { null } else { try { $value | path type } catch { null } }
}

def is-conventional-bare-git-dir [dir: path] {
    ($dir | path basename) == $DOT_GIT and (safe-path-type ($dir | path join $HEAD_FILE)) == file and (($dir | path join $OBJECTS_DIR) | path exists) and (($dir | path join $REFS_DIR) | path exists)
}

def env-git-vcs [] {
    let git_dir = $env.GIT_DIR?
    if ($git_dir == null) or ($git_dir | is-empty) { return null }

    {kind: $VCS_GIT, git_dir: (resolve-git-path $git_dir ($env.PWD? | default "."))}
}

def find-git-or-colocated-jj [] {
    let pwd = $env.PWD?
    if $pwd == null { return null }

    mut dir = ($pwd | path expand)

    loop {
        let dotgit = ($dir | path join $DOT_GIT)
        let type = (safe-path-type $dotgit)

        if $type == dir {
            if (safe-path-type ($dir | path join $DOT_JJ)) == dir { return {kind: $VCS_JJ, root: $dir} }
            return {kind: $VCS_GIT, git_dir: $dotgit}
        }

        if $type == file {
            let git_dir = (git-dir-from-file $dotgit $dir)
            if $git_dir != null {
                if (safe-path-type ($dir | path join $DOT_JJ)) == dir { return {kind: $VCS_JJ, root: $dir} }
                return {kind: $VCS_GIT, git_dir: $git_dir}
            }
        }

        if (is-conventional-bare-git-dir $dir) { return {kind: $VCS_GIT, git_dir: $dir} }

        let parent = ($dir | path dirname)
        if $parent == $dir { return (env-git-vcs) }
        $dir = $parent
    }
}

def find-jj-only [] {
    let pwd = $env.PWD?
    if $pwd == null { return null }

    mut dir = ($pwd | path expand)

    loop {
        if (safe-path-type ($dir | path join $DOT_JJ)) == dir { return {kind: $VCS_JJ, root: $dir} }

        let parent = ($dir | path dirname)
        if $parent == $dir { return null }
        $dir = $parent
    }
}

def find-vcs [] {
    let vcs = (find-git-or-colocated-jj)
    if $vcs != null { return $vcs }

    find-jj-only
}

def branch-from-head [head: string] {
    if ($head | str starts-with $REF_HEADS_PREFIX) {
        $head | str replace $REF_HEADS_PREFIX ""
    } else {
        $DETACHED_LABEL
    }
}

def git-segment [git_dir: path] {
    let head = (try { open --raw ($git_dir | path join $HEAD_FILE) | str trim } catch { return null })
    let branch = (branch-from-head $head)
    $" ($ICON_GIT) ($ICON_BRANCH) ($branch)"
}

def jj-installed [] {
    which jj | is-not-empty
}

# Query jj in prompt-safe mode: no snapshot, no pager, no color, JSON template.
def jj-info [root: path] {
    if not (jj-installed) { return {available: false} }

    let res = (try { ^jj -R $root --ignore-working-copy --no-pager --color never log --no-graph -r @ -T $JJ_TEMPLATE | complete } catch { return {available: false} })
    if $res.exit_code != 0 { return null }

    try { $res.stdout | from json | upsert available true } catch { null }
}

def jj-flags [info: record] {
    [
        (if $info.empty { $FLAG_EMPTY })
        (if $info.conflict { $FLAG_CONFLICT })
        (if $info.divergent { $FLAG_DIVERGENT })
    ] | str join
}

def suffix [value: string] {
    if ($value | is-empty) { "" } else { $" ($value)" }
}

def jj-segment [root: path] {
    let info = (jj-info $root)
    if $info == null { return null }
    if not $info.available { return $" ($ICON_JJ) ($ICON_UNKNOWN)" }

    $" ($ICON_JJ) ($ICON_JJ_AT) ($info.change)(suffix $info.bookmarks)(suffix (jj-flags $info))"
}

def get_vcs_path [] {
    let vcs = (find-vcs)
    if $vcs == null { return null }

    if $vcs.kind == $VCS_JJ { jj-segment $vcs.root } else { git-segment $vcs.git_dir }
}

def path-segment [] {
    let home = if $env.HOME? != null { $env.HOME? } else { $env.USERPROFILE? }
    let notes = if $env.NOTES? != null { $env.NOTES? } else if $home != null { $home | path join $NOTES_DIR } else { null }
    let workspace = if $env.WORKSPACE? != null { $env.WORKSPACE? } else if $home != null { $home | path join $WORKSPACE_DIR } else { null }

    let pwd = $env.PWD?
    if $pwd == null { return "" }

    mut path = ($pwd | path expand)
    if $notes != null { $path = ($path | str replace $notes $ICON_NOTES) }
    if $workspace != null { $path = ($path | str replace $workspace $ICON_WORKSPACE) }
    if $home != null { $path = ($path | str replace $home $ICON_HOME) }

    if ($path | path split | length) >= $TRUNCATE_DEPTH {
        $path | path dirname --num-levels 2 --replace $TRUNCATE_SYMBOL
    } else {
        $path
    }
}

def create_left_prompt [] {
    [
        $ANSI_DARK_GRAY
        (path-segment)
        $ANSI_RESET
        $ANSI_GREEN
        (get_vcs_path)
        $ANSI_RESET
    ] | str join
}

def create_right_prompt [] {
    let status = if (($env.LAST_EXIT_CODE? | default 0) == 0) { $" ($ICON_CHECK) " } else { $"($ANSI_RED) ($ICON_FAILED) ($ANSI_RESET)" }

    [
        $ANSI_DARK_GRAY_DIMMED
        $status
        $ANSI_RESET
    ] | str join
}

export-env {
    $env.PROMPT_COMMAND = { || create_left_prompt }
    $env.PROMPT_COMMAND_RIGHT = { || create_right_prompt }
    $env.PROMPT_INDICATOR = $PROMPT_INDICATOR_TEXT
    $env.PROMPT_INDICATOR_VI_INSERT = $PROMPT_INSERT_TEXT
    $env.PROMPT_INDICATOR_VI_NORMAL = $PROMPT_NORMAL_TEXT
    $env.PROMPT_MULTILINE_INDICATOR = $PROMPT_MULTILINE_TEXT
}