# Prompt styling and context

autoload -Uz add-zsh-hook vcs_info

zstyle ':vcs_info:*' enable git
zstyle ':vcs_info:git:*' formats ' %b'

typeset -g PROMPT_SHOW_SPACER=0
typeset -gA PROMPT_PROJECT_MATCH_CACHE
typeset -gA PROMPT_VERSION_CACHE
typeset -gA PROMPT_SEPARATOR_CACHE
typeset -g PROMPT_FIRST_RENDER=1
typeset -g PROMPT_PROJECT_ROOT_CACHE_PWD=''
typeset -g PROMPT_PROJECT_ROOT_CACHE=''

prompt_preexec() {
    case "$1" in
        clear|clear\ *|reset|reset\ *)
            PROMPT_SHOW_SPACER=0
            ;;
        *)
            PROMPT_SHOW_SPACER=1
            ;;
    esac
}

get_command_version() {
    local command_name="$1"
    local binary_path=""
    local raw_version=""

    if [[ -n ${PROMPT_VERSION_CACHE[$command_name]+x} ]]; then
        print -r -- "${PROMPT_VERSION_CACHE[$command_name]}"
        return 0
    fi

    binary_path="$(whence -p "$command_name" 2>/dev/null || true)"
    if ! [ -n "$binary_path" ] || ! [ -x "$binary_path" ]; then
        PROMPT_VERSION_CACHE[$command_name]=""
        return 0
    fi

    case "$command_name" in
        python)
            raw_version="$($binary_path --version 2>&1 || true)"
            raw_version="${raw_version#Python }"
            ;;
        node)
            raw_version="$($binary_path -v 2>/dev/null || true)"
            raw_version="${raw_version#v}"
            ;;
        go)
            raw_version="$($binary_path version 2>/dev/null || true)"
            raw_version="${raw_version#go version go}"
            raw_version="${raw_version%% *}"
            ;;
        rustc)
            raw_version="$($binary_path -V 2>/dev/null || true)"
            raw_version="${raw_version#rustc }"
            raw_version="${raw_version%% *}"
            ;;
        *)
            PROMPT_VERSION_CACHE[$command_name]=""
            return 0
            ;;
    esac

    [ -n "$raw_version" ] || {
        PROMPT_VERSION_CACHE[$command_name]=""
        return 0
    }

    PROMPT_VERSION_CACHE[$command_name]="$raw_version"
    print -r -- "$raw_version"
}

get_project_root() {
    if [ "$PROMPT_PROJECT_ROOT_CACHE_PWD" = "$PWD" ] && [ -n "$PROMPT_PROJECT_ROOT_CACHE" ]; then
        print -r -- "$PROMPT_PROJECT_ROOT_CACHE"
        return 0
    fi

    PROMPT_PROJECT_ROOT_CACHE_PWD="$PWD"
    PROMPT_PROJECT_ROOT_CACHE="$(command git rev-parse --show-toplevel 2>/dev/null || print -r -- "$PWD")"
    print -r -- "$PROMPT_PROJECT_ROOT_CACHE"
}

project_has_pattern() {
    local project_root=""
    local cache_key=""
    local pattern
    local matches=()

    project_root="$(get_project_root)"
    cache_key="${project_root}::${(j:|:)@}"

    if [[ -n ${PROMPT_PROJECT_MATCH_CACHE[$cache_key]+x} ]]; then
        [[ ${PROMPT_PROJECT_MATCH_CACHE[$cache_key]} == 1 ]]
        return
    fi

    for pattern in "$@"; do
        matches=("${project_root}"/${~pattern}(N))
        if [ "${#matches[@]}" -gt 0 ]; then
            PROMPT_PROJECT_MATCH_CACHE[$cache_key]=1
            return 0
        fi
    done

    PROMPT_PROJECT_MATCH_CACHE[$cache_key]=0

    return 1
}

build_git_status_segment() {
    local git_status_output=""
    local line=""
    local staged=0
    local unstaged=0
    local untracked=0
    local status_parts=()

    command git rev-parse --is-inside-work-tree >/dev/null 2>&1 || return 0
    git_status_output="$(command git status --porcelain 2>/dev/null || true)"

    if [ -z "$git_status_output" ]; then
        print -r -- "%F{green}✓ clean%f"
        return 0
    fi

    while IFS= read -r line; do
        [ -n "$line" ] || continue

        if [ "${line:0:2}" = "??" ]; then
            ((untracked++))
            continue
        fi

        [ "${line:0:1}" = " " ] || ((staged++))
        [ "${line:1:1}" = " " ] || ((unstaged++))
    done <<< "$git_status_output"

    [ "$staged" -eq 0 ] || status_parts+=("+$staged")
    [ "$unstaged" -eq 0 ] || status_parts+=("~$unstaged")
    [ "$untracked" -eq 0 ] || status_parts+=("?$untracked")

    if [ "${#status_parts[@]}" -eq 0 ]; then
        print -r -- "%F{green}✓ clean%f"
    else
        print -r -- "%F{yellow}${(j: :)status_parts}%f"
    fi
}

build_prompt_separator_line() {
    local last_status="$1"
    local line_width="${COLUMNS:-80}"
    local cache_key="${line_width}:${last_status}"
    local separator_line=""
    local color_prefix=""
    local color_suffix=""
    local i=0

    if [[ -n ${PROMPT_SEPARATOR_CACHE[$cache_key]+x} ]]; then
        print -r -- "${PROMPT_SEPARATOR_CACHE[$cache_key]}"
        return 0
    fi

    if [ "$line_width" -lt 1 ]; then
        line_width=1
    fi

    if [ "$last_status" -ne 0 ]; then
        color_prefix="%F{red}"
        color_suffix="%f"
    fi

    while [ "$i" -lt "$line_width" ]; do
        separator_line+="─"
        ((i++))
    done

    PROMPT_SEPARATOR_CACHE[$cache_key]="${color_prefix}${separator_line}${color_suffix}"
    print -r -- "${PROMPT_SEPARATOR_CACHE[$cache_key]}"
}

build_prompt() {
    local last_status=$?
    local git_segment=""
    local git_status_segment=""
    local env_segment=""
    local project_symbols=()
    local tool_versions=()
    local python_version=""
    local node_version=""
    local go_version=""
    local rust_version=""
    local context_line=""
    local newline=$'\n'
    local prompt_spacer=""
    local separator="::"

    if [ "$last_status" -ne 0 ]; then
        separator="%F{red}::%f"
    fi

    if [ "$PROMPT_FIRST_RENDER" -eq 1 ]; then
        PROMPT_FIRST_RENDER=0
        PROMPT="%B%~%b ${separator} "
        RPROMPT="%n@%m"
        return 0
    fi

    vcs_info
    git_segment="${vcs_info_msg_0_}"
    git_status_segment="$(build_git_status_segment)"

    if [ -n "${CONDA_DEFAULT_ENV:-}" ]; then
        env_segment=" ${CONDA_DEFAULT_ENV}"
    elif [ -n "${DIRENV_DIR:-}" ] || [ -n "${DIRENV_FILE:-}" ]; then
        env_segment=" direnv"
    fi

    if project_has_pattern '**/docker-compose*.yml' '**/docker-compose*.yaml' '**/compose*.yml' '**/compose*.yaml'; then
        project_symbols+=("%F{blue}%f")
    fi

    if project_has_pattern '**/Makefile' '**/makefile' '**/GNUmakefile' '**/Justfile' '**/justfile' '**/.justfile'; then
        project_symbols+=("%F{yellow}%f")
    fi

    if project_has_pattern '**/pyproject.toml' '**/requirements.txt' '**/.python-version' '**/setup.py'; then
        python_version="$(get_command_version python)"
        [ -z "$python_version" ] || tool_versions+=("%F{magenta} ${python_version}%f")
    fi

    if project_has_pattern '**/package.json' '**/.nvmrc' '**/pnpm-workspace.yaml' '**/yarn.lock' '**/package-lock.json' '**/bun.lock' '**/bun.lockb'; then
        node_version="$(get_command_version node)"
        [ -z "$node_version" ] || tool_versions+=("%F{green} ${node_version}%f")
    fi

    if project_has_pattern '**/go.mod' '**/go.work'; then
        go_version="$(get_command_version go)"
        [ -z "$go_version" ] || tool_versions+=("%F{cyan} ${go_version}%f")
    fi

    if project_has_pattern '**/Cargo.toml' '**/Cargo.lock' '**/rust-toolchain.toml' '**/rust-toolchain'; then
        rust_version="$(get_command_version rustc)"
        [ -z "$rust_version" ] || tool_versions+=("%F{yellow} ${rust_version}%f")
    fi

    if [ -n "$git_segment" ]; then
        context_line="%F{cyan}${git_segment}%f"
        [ -z "$git_status_segment" ] || context_line+="  ${git_status_segment}"
    fi

    if [ -n "$env_segment" ]; then
        [ -z "$context_line" ] || context_line+="  %F{240}·%f  "
        context_line+="%F{green}${env_segment}%f"
    fi

    if [ "${#project_symbols[@]}" -gt 0 ]; then
        [ -z "$context_line" ] || context_line+="  %F{240}·%f  "
        context_line+="${(j: :)project_symbols}"
    fi

    if [ "${#tool_versions[@]}" -gt 0 ]; then
        [ -z "$context_line" ] || context_line+="  %F{240}·%f  "
        context_line+="${(j:  :)tool_versions}"
    fi

    if [ -n "$context_line" ]; then
        context_line+="${newline}"
    fi

    if [ "$PROMPT_SHOW_SPACER" -eq 1 ]; then
        prompt_spacer="$(build_prompt_separator_line "$last_status")${newline}"
    fi

    PROMPT="${prompt_spacer}${context_line}%B%~%b ${separator} "
    RPROMPT="%n@%m"
}

add-zsh-hook preexec prompt_preexec
add-zsh-hook precmd build_prompt
build_prompt