dotfiles

# Git worktree helpers.
# Layout:
#   <root>/.git          bare repository
#   <root>/<branch-slug> worktrees

const WT_MAIN_BRANCH = "wt.main-branch"

const DOT_GIT = ".git"

const DEFAULT_MAIN_BRANCH = "main"
const HEAD = "HEAD"
const REF_HEADS = "refs/heads/"
const DIRTY_REFUSE = "refuse"
const DIRTY_COPY = "copy"
const DIRTY_MODES = ["refuse" "copy"]
const STATUS_PORCELAIN = ["status" "--porcelain=v1"]
const STAGED_PATCH = "staged.patch"
const UNSTAGED_PATCH = "unstaged.patch"

const LINE_WORKTREE = "worktree "
const LINE_HEAD = "HEAD "
const LINE_BRANCH = "branch "
const LINE_DETACHED = "detached"
const LINE_BARE = "bare"
const LINE_LOCKED = "locked"
const LINE_PRUNABLE = "prunable"

# Run git and return completion record.
def git-complete [args: list<string>, cwd?: path] {
    if $cwd == null {
        ^git ...$args | complete
    } else {
        ^git -C $cwd ...$args | complete
    }
}

# Format a failed git invocation.
def git-failure [args: list<string>, stderr: string] {
    let msg = ($stderr | str trim)
    if ($msg | is-empty) {
        $"git failed: git ($args | str join ' ')"
    } else {
        $msg
    }
}

def choose-recovery [message: string, choices: list<record>, non_interactive: string] {
    if not $nu.is-interactive { error make {msg: $non_interactive} }

    print -e ""
    print -e $message
    let selected = ($choices | input list --display label "choose")
    if $selected == null or $selected.id == "cancel" { error make {msg: "cancelled"} }
    $selected.id
}

def remove-needs-force [message: string] {
    let modified = ($message | str contains "contains modified or untracked files")
    let submodules = ($message | str contains "working trees containing submodules")
    let dirty = ($message | str contains "is dirty")
    let not_clean = ($message | str contains "not clean")
    $modified or $submodules or $dirty or $not_clean
}

# Run git and return trimmed stdout; fail with stderr as Nushell error.
def git-output [args: list<string>, cwd?: path] {
    let res = if $cwd == null { git-complete $args } else { git-complete $args $cwd }

    if $res.exit_code != 0 {
        error make {msg: (git-failure $args $res.stderr)}
    }

    $res.stdout | str trim
}

# Compatibility wrappers used by the rest of this module.
def gitx [args: list<string>, cwd?: path] {
    if $cwd == null { git-output $args } else { git-output $args $cwd }
}

def gitc [args: list<string>, cwd?: path] {
    if $cwd == null { git-complete $args } else { git-complete $args $cwd }
}

def git-dir-args [git_dir: path, args: list<string>] {
    ["--git-dir" $git_dir] | append $args
}

def git-dir-output [git_dir: path, args: list<string>] {
    git-output (git-dir-args $git_dir $args)
}

def git-dir-complete [git_dir: path, args: list<string>] {
    git-complete (git-dir-args $git_dir $args)
}

def git-config [git_dir: path, key: string] {
    let res = (git-dir-complete $git_dir ["config" "--get" $key])
    if $res.exit_code == 0 { $res.stdout | str trim } else { null }
}

def git-config-values [git_dir: path, key: string] {
    let res = (git-dir-complete $git_dir ["config" "--get-all" $key])
    if $res.exit_code == 0 { $res.stdout | lines } else { [] }
}

def git-config-bool [git_dir: path, key: string] {
    let res = (git-dir-complete $git_dir ["config" "--bool" "--get" $key])
    if $res.exit_code == 0 { $res.stdout | str trim } else { null }
}

def git-config-set [git_dir: path, key: string, value: string] {
    git-dir-output $git_dir ["config" $key $value] | ignore
}

def git-config-unset [git_dir: path, key: string] {
    git-dir-complete $git_dir ["config" "--unset" $key] | ignore
}

def ensure-upstream [git_dir: path, branch: string] {
    if (git-config $git_dir "remote.origin.url") == null { return }

    if (git-config $git_dir "remote.origin.fetch") == null {
        git-config-set $git_dir "remote.origin.fetch" "+refs/heads/*:refs/remotes/origin/*"
    }

    if (git-config $git_dir $"branch.($branch).remote") == null {
        git-config-set $git_dir $"branch.($branch).remote" "origin"
        git-config-set $git_dir $"branch.($branch).merge" $"refs/heads/($branch)"
    }

    if (git-dir-complete $git_dir ["show-ref" "--verify" "--quiet" $"refs/heads/($branch)"]).exit_code == 0 {
        git-dir-output $git_dir ["update-ref" $"refs/remotes/origin/($branch)" $"refs/heads/($branch)"] | ignore
    }
}

# Branch name -> path-safe suffix.
# Slashes stay slashes, so `features/make-it-good` keeps its folder hierarchy.
# Every segment is normalized on its own: characters that win32 rejects in file
# names become `-`, and DOS device names get a trailing `-` so the path stays
# usable on both win32 and linux.
def branch-slug [branch: string] {
    let illegal = '[<>:"|?*\\\s-]+'
    let reserved = ["NUL" "CON" "PRN" "AUX" "COM1" "COM2" "COM3" "COM4" "COM5" "COM6" "COM7" "COM8" "COM9" "LPT1" "LPT2" "LPT3" "LPT4" "LPT5" "LPT6" "LPT7" "LPT8" "LPT9"]

    $branch
    | split row "/"
    | each {|segment|
        let cleaned = ($segment | str replace --all --regex $illegal '-' | str trim --char '-')
        if ($cleaned | is-empty) {
            null
        } else if (($cleaned | str uppercase) in $reserved) {
            $"($cleaned)-"
        } else {
            $cleaned
        }
    }
    | where {|segment| $segment != null }
    | str join "/"
}

# URL/path -> repo name.
def repo-name-from [source: string] {
    $source | path basename | str replace --regex '\.git$' ''
}

# Current layout record. Fails outside co/wt layout.
def layout [] {
    let git_dir = (git-output ["rev-parse" "--path-format=absolute" "--git-common-dir"])

    if (git-config-bool $git_dir "core.bare") != "true" {
        error make {msg: "not inside co/wt layout: git common dir is not bare"}
    }

    let root = ($git_dir | path dirname)
    let repo = ($root | path basename)
    let main = (git-config $git_dir $WT_MAIN_BRANCH | default $DEFAULT_MAIN_BRANCH)

    {root: $root, git_dir: $git_dir, repo: $repo, main: $main}
}

# <root>/<branch-slug> as a native path.
# The slug may contain `/`, so join it segment-wise instead of embedding a
# foreign separator into a win32 path.
def branch-path [root: path, branch: string] {
    [$root] | append (branch-slug $branch | split row "/") | path join
}

# Top-level entry that a branch's worktree occupies inside the layout root.
def branch-root-entry [branch: string] {
    branch-slug $branch | split row "/" | first
}

# Computed conventional worktree path for branch.
def wt-path [branch: string] {
    branch-path (layout).root $branch
}

def empty-worktree-record [] {
    {path: null, head: null, branch: null, detached: false, bare: false, locked: false, lock_reason: null, prunable: false}
}

def strip-prefix [line: string, prefix: string] {
    $line | str replace $prefix ""
}

def parse-lock-line [line: string] {
    let reason = (strip-prefix $line $LINE_LOCKED | str trim)
    if ($reason | is-empty) { null } else { $reason }
}

def parse-worktree-line [rec: record, line: string] {
    if ($line | str starts-with $LINE_WORKTREE) {
        return ($rec | upsert path (strip-prefix $line $LINE_WORKTREE))
    }
    if ($line | str starts-with $LINE_HEAD) {
        return ($rec | upsert head (strip-prefix $line $LINE_HEAD))
    }
    if ($line | str starts-with $LINE_BRANCH) {
        let branch = (strip-prefix $line $LINE_BRANCH | str replace $REF_HEADS "")
        return ($rec | upsert branch $branch)
    }
    if $line == $LINE_DETACHED {
        return ($rec | upsert detached true)
    }
    if $line == $LINE_BARE {
        return ($rec | upsert bare true)
    }
    if ($line | str starts-with $LINE_LOCKED) {
        return ($rec | upsert locked true | upsert lock_reason (parse-lock-line $line))
    }
    if ($line | str starts-with $LINE_PRUNABLE) {
        return ($rec | upsert prunable true)
    }

    $rec
}

def parse-worktree-block [block: string] {
    $block | lines | reduce --fold (empty-worktree-record) {|line, rec| parse-worktree-line $rec $line }
}

# Parse `git worktree list --porcelain` into records.
def parse-worktree-list [text: string] {
    $text
    | split row "\n\n"
    | where {|block| not (($block | str trim) | is-empty) }
    | each {|block| parse-worktree-block $block }
}

# Worktree table, no dirty check.
def worktrees-basic [] {
    let l = (layout)
    parse-worktree-list (git-dir-output $l.git_dir ["worktree" "list" "--porcelain"])
}

# Find actual worktree path by branch; fallback to convention path.
def path-for-branch [branch: string] {
    let found = (worktrees-basic | where branch == $branch | first)
    if ($found == null) { wt-path $branch } else { $found.path }
}

def branch-exists [git_dir: path, branch: string] {
    (git-dir-complete $git_dir ["show-ref" "--verify" "--quiet" $"($REF_HEADS)($branch)"]).exit_code == 0
}

def add-candidate-display [branch: string, source: string, base: string] {
    $"($branch)  ($source)  ($base)"
}

def local-add-candidates [git_dir: path] {
    git-dir-output $git_dir ["for-each-ref" "--format=%(refname:short)" "refs/heads"]
    | lines
    | where {|branch| not (($branch | str trim) | is-empty) }
    | each {|branch| {branch: $branch, base: $branch, source: "local", display: (add-candidate-display $branch "local" $branch)} }
}

def remote-add-candidates [git_dir: path] {
    git-dir-output $git_dir ["for-each-ref" "--format=%(refname:short)" "refs/remotes"]
    | lines
    | where {|ref| not (($ref | str trim) | is-empty) }
    | where {|ref| not ($ref | str ends-with "/HEAD") }
    | each {|ref|
        let branch = ($ref | split row "/" | skip 1 | str join "/")
        {branch: $branch, base: $ref, source: "remote", display: (add-candidate-display $branch "remote" $ref)}
    }
}

def git-remotes [git_dir: path] {
    git-dir-output $git_dir ["remote"]
    | lines
    | where {|name| not (($name | str trim) | is-empty) }
    | each {|name| {name: $name, url: (git-dir-output $git_dir ["remote" "get-url" $name])} }
}

def gh-pr-add-candidates [remotes: list<record>] {
    if (which gh | is-empty) { return [] }
    if (^gh auth status | complete).exit_code != 0 { return [] }

    $remotes
    | each {|remote|
        let res = (^gh pr list --repo $remote.url --json number,headRefName,title --limit 100 | complete)
        if $res.exit_code != 0 or (($res.stdout | str trim) | is-empty) {
            []
        } else {
            $res.stdout
            | from json
            | each {|pr|
                let branch = $pr.headRefName
                let base = $"pull/($pr.number)/head"
                {branch: $branch, base: $base, source: "gh-pr", remote: $remote.name, number: $pr.number, display: (add-candidate-display $branch $"gh ($remote.name) #($pr.number)" $pr.title)}
            }
        }
    }
    | flatten
}

def parse-fj-pr-line [line: string] {
    let trimmed = ($line | str trim)
    let first = ($trimmed | split row " " | first | str replace "#" "")
    if not ($first =~ '^\d+$') { return null }

    let title = (
        $trimmed
        | str replace --regex '^#?\d+\s+' ''
        | str replace --regex '\s+\(by .+\)$' ''
        | str trim
    )

    let number = ($first | into int)
    {branch: $"pr-($number)", base: $"pull/($number)/head", source: "fj-pr", number: $number, title: $title, display: (add-candidate-display $"pr-($number)" $"fj #($number)" $title)}
}

def fj-pr-add-candidates [remotes: list<record>] {
    if (which fj | is-empty) { return [] }
    if (^fj auth list | complete).exit_code != 0 { return [] }

    $remotes
    | each {|remote|
        let res = (^fj --style minimal pr -R $remote.name search --state open | complete)
        if $res.exit_code != 0 or (($res.stdout | str trim) | is-empty) {
            []
        } else {
            $res.stdout
            | lines
            | each {|line|
                let parsed = (parse-fj-pr-line $line)
                if $parsed == null { null } else { $parsed | upsert remote $remote.name | upsert display (add-candidate-display $parsed.branch $"fj ($remote.name) #($parsed.number)" $parsed.title) }
            }
            | where {|candidate| $candidate != null }
        }
    }
    | flatten
}

def dedupe-add-candidates [] {
    $in | sort-by branch | uniq-by branch
}

# Candidate branches for `wt add`.
# PR lookups spawn `gh`/`fj` with network round trips, so they stay opt-in.
def raw-add-candidates [include_prs: bool] {
    let l = (layout)
    let occupied = (worktrees-basic | where branch != null | get --optional branch)

    let candidates = if $include_prs {
        let remotes = (git-remotes $l.git_dir)
        gh-pr-add-candidates $remotes | append (fj-pr-add-candidates $remotes)
    } else {
        local-add-candidates $l.git_dir | append (remote-add-candidates $l.git_dir)
    }

    $candidates | where {|candidate|
        (not ($candidate.branch in $occupied)) and (not ((wt-path $candidate.branch) | path exists))
    }
}

def available-add-candidates [include_prs: bool] {
    raw-add-candidates $include_prs | dedupe-add-candidates
}

def add-candidate-matches [candidate: record, query: string] {
    let branch = ($candidate.branch | str contains --ignore-case $query)
    let base = ($candidate.base | str contains --ignore-case $query)
    let display = ($candidate.display | str contains --ignore-case $query)
    $branch or $base or $display
}

def pick-add-candidate [query: any, include_prs: bool] {
    let candidates = (available-add-candidates $include_prs)
    let matches = if $query == null {
        $candidates
    } else {
        $candidates | where {|candidate| add-candidate-matches $candidate $query }
    }

    if ($matches | is-empty) {
        if $query != null {
            let action = (choose-recovery
                $"No local branch, remote branch, or PR matched `($query)`."
                [
                    {id: "create", label: $"Create new branch `($query)` from HEAD"}
                    {id: "cancel", label: "Cancel"}
                ]
                $"no branch/remote/PR matching `($query)`. Non-interactive: create explicitly with `git wt new ($query)` or pass `git wt add ($query) --base <base>`."
            )
            if $action == "create" { return {branch: $query, base: $HEAD, source: "new", display: $"($query)  new  HEAD"} }
        }
        error make {msg: "no branches or PRs available; use `wt new` to create a branch"}
    }

    if ($matches | length) == 1 { return ($matches | first) }

    let selected = ($matches | input list --fuzzy --display display "branch")
    if $selected == null { error make {msg: "no branch selected"} }
    $selected
}

def ensure-fetch-refspec [git_dir: path, remote: string, refspec: string] {
    let key = $"remote.($remote).fetch"
    if not ($refspec in (git-config-values $git_dir $key)) {
        git-dir-output $git_dir ["config" "--add" $key $refspec] | ignore
    }
}

def split-remote-branch [git_dir: path, ref: string] {
    let parts = ($ref | split row "/")
    if ($parts | length) < 2 { return null }

    let remote = ($parts | first)
    let branch = ($parts | skip 1 | str join "/")
    if (git-dir-complete $git_dir ["show-ref" "--verify" "--quiet" $"refs/remotes/($remote)/($branch)"]).exit_code != 0 { return null }
    {remote: $remote, branch: $branch}
}

def configure-branch-upstream [git_dir: path, branch: string, remote: string, merge_ref: string] {
    git-config-set $git_dir $"branch.($branch).remote" $remote
    git-config-set $git_dir $"branch.($branch).merge" $merge_ref
}

def configure-upstream-from-base [git_dir: path, branch: string, base: string] {
    let remote_branch = (split-remote-branch $git_dir $base)
    if $remote_branch != null {
        configure-branch-upstream $git_dir $branch $remote_branch.remote $"refs/heads/($remote_branch.branch)"
    }
}

def configure-upstream-matching-remote [git_dir: path, branch: string] {
    for remote in (git-remotes $git_dir) {
        if (git-dir-complete $git_dir ["show-ref" "--verify" "--quiet" $"refs/remotes/($remote.name)/($branch)"]).exit_code == 0 {
            configure-branch-upstream $git_dir $branch $remote.name $"refs/heads/($branch)"
            return
        }
    }
}

def ensure-branch-points-at [git_dir: path, branch: string, target_ref: string, label: string] {
    if not (branch-exists $git_dir $branch) { return }

    let local = (git-dir-output $git_dir ["rev-parse" $branch])
    let target = (git-dir-output $git_dir ["rev-parse" $target_ref])
    if $local == $target { return }

    let action = (choose-recovery
        $"Local branch `($branch)` already exists but does not point at selected ($label)."
        [
            {id: "reset", label: $"Reset `($branch)` to ($label)"}
            {id: "cancel", label: "Cancel"}
        ]
        $"local branch `($branch)` differs from selected ($label). Non-interactive: delete/reset it first, then rerun `git wt add ($branch)`."
    )
    if $action == "reset" { git-dir-output $git_dir ["branch" "-f" $branch $target_ref] | ignore }
}

def create-picked-worktree [candidate: record, submodules: bool] {
    let l = (layout)
    if $candidate.source == "remote" {
        ensure-branch-points-at $l.git_dir $candidate.branch $candidate.base $"remote branch ($candidate.base)"
        let path = (create-worktree $candidate.branch $candidate.base $submodules)
        configure-upstream-from-base $l.git_dir $candidate.branch $candidate.base
        return $path
    }

    if $candidate.source in ["gh-pr" "fj-pr"] {
        let tracking_ref = $"refs/remotes/($candidate.remote)/pull/($candidate.number)"
        let refspec = $"+refs/pull/($candidate.number)/head:($tracking_ref)"
        ensure-fetch-refspec $l.git_dir $candidate.remote $refspec
        git-dir-output $l.git_dir ["fetch" $candidate.remote $"refs/pull/($candidate.number)/head:($tracking_ref)"] | ignore
        ensure-branch-points-at $l.git_dir $candidate.branch $tracking_ref $"PR #($candidate.number) from ($candidate.remote)"
        let path = (create-worktree $candidate.branch $tracking_ref $submodules)
        configure-branch-upstream $l.git_dir $candidate.branch $candidate.remote $"refs/pull/($candidate.number)/head"
        return $path
    }

    let path = (create-worktree $candidate.branch $candidate.base $submodules)
    if $candidate.source == "local" { configure-upstream-matching-remote $l.git_dir $candidate.branch }
    $path
}

# Init + checkout submodules inside a freshly added worktree.
# No-op when disabled or the worktree has no .gitmodules. Best-effort:
# failure (e.g. no network/auth) warns but leaves the worktree intact.
def init-submodules [path: path, enabled: bool] {
    if not $enabled { return }
    if not (([$path ".gitmodules"] | path join) | path exists) { return }

    let res = (git-complete ["submodule" "update" "--init" "--recursive"] $path)
    if $res.exit_code != 0 {
        print -e $"warning: submodule init failed in ($path); run `git -C ($path) submodule update --init --recursive` to retry"
        let detail = ($res.stderr | str trim)
        if not ($detail | is-empty) { print -e $detail }
    }
}

def add-worktree-needs-fetch [message: string] {
    let invalid = ($message | str contains "invalid reference")
    let unknown = ($message | str contains "unknown revision")
    let not_commit = ($message | str contains "not a commit")
    let no_remote = ($message | str contains "couldn't find remote ref")
    $invalid or $unknown or $not_commit or $no_remote
}

def add-worktree [git_dir: path, path: path, branch: string, base: string, submodules: bool] {
    let args = if (branch-exists $git_dir $branch) {
        ["worktree" "add" $path $branch]
    } else {
        ["worktree" "add" "-b" $branch $path $base]
    }

    let res = (git-dir-complete $git_dir $args)
    if $res.exit_code != 0 {
        let msg = (git-failure $args $res.stderr)
        if not (add-worktree-needs-fetch $msg) { error make {msg: $msg} }

        let action = (choose-recovery
            $"Git could not create this worktree.\n\n($msg)\n\nOften the base branch/ref is stale or has not been fetched yet."
            [
                {id: "fetch", label: "Fetch all remotes, then retry"}
                {id: "cancel", label: "Cancel"}
            ]
            $"($msg)\nNon-interactive: run `git fetch --all --prune`, then retry the `git wt` command."
        )
        if $action == "fetch" {
            git-dir-output $git_dir ["fetch" "--all" "--prune"] | ignore
            git-dir-output $git_dir $args | ignore
        }
    }

    init-submodules $path $submodules
}

# Create branch + worktree, or add existing branch.
def create-worktree [branch: string, base: string, submodules: bool] {
    let l = (layout)
    let existing = (worktrees-basic | where branch == $branch | first)
    if $existing != null {
        let action = (choose-recovery
            $"Branch `($branch)` already has a worktree:\n($existing.path)"
            [
                {id: "use", label: "Use existing worktree"}
                {id: "cancel", label: "Cancel"}
            ]
            $"branch `($branch)` already has worktree ($existing.path). Non-interactive: use `git wt go ($branch)`."
        )
        if $action == "use" { return $existing.path }
    }

    let path = (wt-path $branch)

    if ($path | path exists) {
        let action = (choose-recovery
            $"Target path already exists: ($path)\nGit cannot create a worktree on top of it."
            [
                {id: "use", label: "Use existing path; do not create anything"}
                {id: "remove", label: "Delete existing path, then create worktree"}
                {id: "cancel", label: "Cancel"}
            ]
            $"target exists: ($path)\nNon-interactive: remove it first (`rm -rf ($path)`) or use it manually."
        )
        if $action == "use" { return $path }
        rm -rf $path
    }

    # Create parent directories for nested worktree paths (e.g. features/make-it-good).
    let parent = ($path | path dirname)
    if ($parent != $path) {
        mkdir $parent | ignore
    }

    add-worktree $l.git_dir $path $branch $base $submodules
    configure-upstream-from-base $l.git_dir $branch $base
    $path
}

# Current branch, refusing detached HEAD.
def current-branch [] {
    let branch = (git-output ["rev-parse" "--abbrev-ref" $HEAD])
    if $branch == $HEAD { error make {msg: "detached HEAD is not supported"} }
    $branch
}

def current-repo-top [] {
    try { git-output ["rev-parse" "--show-toplevel"] | path expand } catch { "" }
}

def existing-main-or-root [layout: record] {
    let main_path = (path-for-branch $layout.main)
    if ($main_path | path exists) { $main_path } else { $layout.root }
}

def --env cd-away-from-target [layout: record, target: path, branch: string] {
    if $target != (current-repo-top) { return }

    if $branch == $layout.main {
        cd $layout.root
    } else {
        cd (existing-main-or-root $layout)
    }
}

def remove-worktree-args [target: path, force: bool] {
    if $force { ["worktree" "remove" "--force" $target] } else { ["worktree" "remove" $target] }
}

def remove-worktree-git [git_dir: path, target: path, force: bool, command_hint: string] {
    let res = (git-dir-complete $git_dir (remove-worktree-args $target $force))
    if $res.exit_code == 0 { return }

    let msg = (git-failure (remove-worktree-args $target $force) $res.stderr)
    if $force or not (remove-needs-force $msg) { error make {msg: $msg} }

    let action = (choose-recovery
        $"Git refused to remove this worktree.\n\n($msg)\n\nUsually this means dirty files or submodules are present. Force removal deletes that worktree directory."
        [
            {id: "force", label: "Force remove worktree"}
            {id: "cancel", label: "Cancel"}
        ]
        $"($msg)\nNon-interactive: rerun with `($command_hint) --force`."
    )
    if $action == "force" { git-dir-output $git_dir (remove-worktree-args $target true) | ignore }
}

def delete-branch [git_dir: path, branch: string, force: bool, command_hint: string] {
    let flag = if $force { "-D" } else { "-d" }
    let res = (git-dir-complete $git_dir ["branch" $flag $branch])
    if $res.exit_code == 0 { return }

    let msg = (git-failure ["branch" $flag $branch] $res.stderr)
    if $force { error make {msg: $msg} }

    let action = (choose-recovery
        $"Git refused to delete branch `($branch)`.\n\n($msg)\n\nUsually this means commits are not merged upstream."
        [
            {id: "force", label: $"Force delete branch `($branch)`"}
            {id: "keep", label: "Keep branch"}
            {id: "cancel", label: "Cancel"}
        ]
        $"($msg)\nNon-interactive: rerun with `($command_hint) --force` to force-delete, or use `git wt rm ($branch)` to keep the branch."
    )
    if $action == "force" { git-dir-output $git_dir ["branch" "-D" $branch] | ignore }
}

# Remove worktree, maybe branch, and cd away if removing current worktree.
def --env remove-worktree [branch: string, delete_branch: bool, force: bool, command_hint: string] {
    let l = (layout)
    let target = (path-for-branch $branch | path expand)

    cd-away-from-target $l $target $branch
    remove-worktree-git $l.git_dir $target $force $command_hint

    if $delete_branch {
        delete-branch $l.git_dir $branch $force $command_hint
    }

    $target
}



def choose-fresh-branch [git_dir: path, branch: any] {
    if (($branch == null) or ($branch | is-empty)) {
        git-dir-output $git_dir ["symbolic-ref" "--short" $HEAD]
    } else {
        $branch
    }
}

def selected-branch [branch: any] {
    if (($branch == null) or ($branch | is-empty)) { current-branch } else { $branch }
}

# Create fresh layout from remote/local repo URL/path.
def --env co-fresh [source: string, name?: string, branch?: string, dir?: path, submodules: bool = true] {
    let repo = if ($name == null) { repo-name-from $source } else { $name }
    let parent = if ($dir == null) { "." } else { $dir }
    let root = ([$parent $repo] | path join | path expand)
    let git_dir = ([$root $DOT_GIT] | path join)

    if ($root | path exists) {
        let action = (choose-recovery
            $"Clone target already exists: ($root)"
            [
                {id: "remove", label: "Delete target directory, then clone"}
                {id: "cancel", label: "Cancel"}
            ]
            $"target exists: ($root). Non-interactive: choose another name/dir or remove it first."
        )
        if $action == "remove" { rm -rf $root }
    }

    mkdir $root
    git-output ["clone" "--bare" $source $git_dir] | ignore

    let chosen_branch = (choose-fresh-branch $git_dir $branch)
    git-config-set $git_dir $WT_MAIN_BRANCH $chosen_branch
    ensure-upstream $git_dir $chosen_branch
    git-dir-complete $git_dir ["fetch" "--all" "--prune"] | ignore

    let worktree = (branch-path $root $chosen_branch)
    git-dir-output $git_dir ["worktree" "add" $worktree $chosen_branch] | ignore
    init-submodules $worktree $submodules
    cd $worktree

    {root: $root, worktree: $worktree, branch: $chosen_branch}
}

def status-porcelain [cwd: path] {
    git-output $STATUS_PORCELAIN $cwd
}

def stash-count [cwd: path] {
    let res = (git-complete ["stash" "list"] $cwd)
    if $res.exit_code == 0 { $res.stdout | lines | length } else { 0 }
}

def assert-dirty-mode [dirty: string] {
    if not ($dirty in $DIRTY_MODES) {
        error make {msg: $"--dirty must be one of: ($DIRTY_MODES | str join ', ')"}
    }
}

def migration-plan [dirty_mode: string, branch: any] {
    assert-dirty-mode $dirty_mode

    let src = (git-output ["rev-parse" "--show-toplevel"] | path expand)
    let chosen_branch = (selected-branch $branch)
    let target = (branch-path $src $chosen_branch)
    let status_before = (status-porcelain $src)
    let is_dirty = (not ($status_before | is-empty))

    {
        mode: "current-checkout"
        source: $src
        target_root: $src
        target_git_dir: ([$src $DOT_GIT] | path join)
        target_worktree: $target
        repo: ($src | path basename)
        branch: $chosen_branch
        dirty: $is_dirty
        dirty_mode: $dirty_mode
        stashes: (stash-count $src)
        destructive: true
        status_before: $status_before
    }
}

def public-migration [plan: record] {
    $plan | reject status_before
}

def resolve-migration-dirty [plan: record] {
    if not ($plan.dirty and $plan.dirty_mode == $DIRTY_REFUSE) { return $plan }

    let action = (choose-recovery
        "Working tree has uncommitted changes. Migration can copy them into the new worktree, or you can cancel and commit/stash manually."
        [
            {id: "copy", label: "Copy dirty state into migrated worktree"}
            {id: "cancel", label: "Cancel"}
        ]
        "working tree has changes. Non-interactive: rerun with `git wt migrate --dirty copy`, or commit/stash first."
    )
    if $action == "copy" { return ($plan | upsert dirty_mode $DIRTY_COPY) }
    $plan
}

def assert-migration-ready [plan: record] {
    if $plan.dirty and $plan.dirty_mode == $DIRTY_REFUSE {
        error make {msg: "working tree has changes; rerun with `--dirty copy` or commit/stash first"}
    }
}

def patch-paths [tmp: path] {
    {staged: ([$tmp $STAGED_PATCH] | path join), unstaged: ([$tmp $UNSTAGED_PATCH] | path join)}
}

def write-dirty-patches [source: path, tmp: path] {
    let patches = (patch-paths $tmp)
    ^git -C $source diff --binary --cached | save --force $patches.staged
    ^git -C $source diff --binary | save --force $patches.unstaged
    $patches
}

def patch-has-data [patch: path] {
    not ((open --raw $patch) | is-empty)
}

def apply-patch [target: path, args: list<string>, patch: path] {
    if (patch-has-data $patch) {
        git-output (["apply"] | append $args | append $patch) $target | ignore
    }
}

def apply-dirty-patches [target: path, patches: record] {
    apply-patch $target ["--index"] $patches.staged
    apply-patch $target [] $patches.unstaged
}

def copy-untracked [source: path, target: path, --exclude: list<string> = []] {
    mkdir $target
    let skip = ([$DOT_GIT] | append $exclude)
    for entry in (ls -a $source | get name) {
        if (($entry | path basename) in $skip) { continue }
        cp -r --no-dereference --no-clobber $entry $target
    }
}

def copy-dirty-state [plan: record, patches: record] {
    apply-dirty-patches $plan.target_worktree $patches
    copy-untracked $plan.source $plan.target_worktree --exclude [($plan.target_worktree | path basename)]

    let status_after = (status-porcelain $plan.target_worktree)
    if $plan.status_before != $status_after {
        error make {msg: $"dirty-copy verification failed; target left for inspection at ($plan.target_root)"}
    }
}

# Copy stash refs from old .git to new bare .git.
# `git clone --bare` does not transport refs/stash.
def copy-git-config [old_git: path, new_git: path] {
    let src = ([$old_git "config"] | path join)
    let dst = ([$new_git "config"] | path join)
    if ($src | path exists) { cp --force $src $dst }
}

def preserve-stash-refs [old_git: path, new_git: path] {
    for sub in ["refs/stash" "logs/refs/stash"] {
        let src = ([$old_git $sub] | path join)
        let dst = ([$new_git $sub] | path join)
        if ($src | path exists) {
            mkdir ($dst | path dirname)
            cp -r --no-dereference --no-clobber $src $dst
        }
    }
}

# Remove top-level entries from `root` except those in `keep`.
def clear-stale-files [root: path, keep: list<string>] {
    let keep_set = ($keep | append [".", ".."])
    for entry in (ls -a $root) {
        if (($entry.name | path basename) not-in $keep_set) {
            rm -rf $entry.name
        }
    }
}

# Convert current checkout in place to bare+worktree layout.
# Original tracked files and untracked files are removed; dirty state is
# restored from patches, untracked files are copied, stashes are preserved.
def --env co-current [branch?: string, plan: bool = false, dirty: string = $DIRTY_REFUSE, submodules: bool = true] {
    let migration = (resolve-migration-dirty (migration-plan $dirty $branch))
    if $plan { return (public-migration $migration) }

    assert-migration-ready $migration

    let src = $migration.source
    let old_git = ([$src $DOT_GIT] | path join)
    let old_git_backup = ([$src ".git.wt-backup"] | path join)
    let new_git = ([$src ".git.wt-new"] | path join)
    let tmp = (mktemp -d)

    if ($old_git_backup | path exists) { error make {msg: $"backup exists: ($old_git_backup)"} }
    if ($new_git | path exists) { error make {msg: $"temporary git dir exists: ($new_git)"} }

    try {
        let patches = if $migration.dirty { write-dirty-patches $src $tmp } else { null }

        # Keep .git backup on same filesystem: cross-device mv copies sockets and can ENXIO.
        git-output ["clone" "--bare" $src $new_git] | ignore
        mv $old_git $old_git_backup
        mv $new_git $old_git

        copy-git-config $old_git_backup $migration.target_git_dir
        git-config-set $migration.target_git_dir "core.bare" "true"
        git-config-unset $migration.target_git_dir "core.worktree"
        preserve-stash-refs $old_git_backup $migration.target_git_dir
        git-config-set $migration.target_git_dir $WT_MAIN_BRANCH $migration.branch
        ensure-upstream $migration.target_git_dir $migration.branch
        git-dir-output $migration.target_git_dir ["worktree" "add" $migration.target_worktree $migration.branch] | ignore

        if $migration.dirty {
            copy-dirty-state $migration $patches
        }

        clear-stale-files $src [$DOT_GIT (branch-root-entry $migration.branch) ".git.wt-backup"]
        init-submodules $migration.target_worktree $submodules

        rm -rf $old_git_backup
        rm -rf $tmp
    } catch {|e|
        if ($new_git | path exists) { rm -rf $new_git }
        if ($old_git_backup | path exists) {
            rm -rf $old_git
            mv $old_git_backup $old_git
        }
        rm -rf $tmp
        error make {msg: ($e | get msg)}
    }

    cd $migration.target_worktree
    public-migration $migration
}

# Bootstrap a new bare+worktree layout from a URL/path.
export def --env "wt clone" [
    source: string             # Git URL/path to clone into bare+worktree layout
    name?: string              # Root directory name; default derives from source basename
    --branch (-b): string = "" # Initial branch to check out; default remote HEAD
    --dir (-d): path = "."     # Parent directory for the new layout root
    --no-submodules            # Skip submodule init in the initial worktree
] {
    co-fresh $source $name $branch $dir (not $no_submodules)
}

# Migrate the current checkout in place to a bare+worktree layout.
# Destructive: original tracked files are removed; dirty state, untracked
# files, and stashes are preserved.
export def --env "wt migrate" [
    --branch (-b): string = ""      # Branch to make the first worktree; default current branch
    --plan                          # Print migration plan only; do not modify files
    --dirty: string = $DIRTY_REFUSE # Dirty handling: `refuse` or `copy`; prompt if dirty and omitted
    --no-submodules                 # Skip submodule init in the migrated worktree
] {
    co-current $branch $plan $dirty (not $no_submodules)
}

# Summary.
export def wt [] { wt ls }

export def "wt root" [] { (layout).root }
export def "wt repo" [] { (layout).repo }
export def "wt path" [
    branch: string # Branch name to map to its conventional worktree path
] { wt-path $branch }

def worktree-dirty [row: record] {
    if (($row.path | path exists) and (not $row.bare) and ($row.branch != null)) {
        not ((status-porcelain $row.path) | is-empty)
    } else {
        false
    }
}

export def "wt ls" [] {
    let cwd_top = (current-repo-top)

    worktrees-basic | where bare == false | each {|row|
        let path = ($row.path | path expand)
        $row
        | upsert current ($path == $cwd_top)
        | upsert dirty (worktree-dirty $row)
    }
}

export def --env "wt new" [
    branch: string                 # New or existing local branch to create a worktree for
    --base (-b): string = $HEAD    # Base ref for a new branch; ignored when branch already exists
    --switch (-s)                  # cd into the created or existing worktree
    --no-submodules                # Skip `git submodule update --init --recursive`
] {
    let path = (create-worktree $branch $base (not $no_submodules))
    if $switch { cd $path }
    $path
}

export def --env "wt go" [
    branch: string                 # Branch/worktree to cd into
    --create (-c)                  # Create missing worktree without prompting
    --base (-b): string = $HEAD    # Base ref when creating a missing branch/worktree
    --no-submodules                # With create: skip submodule init
] {
    let path = (path-for-branch $branch)
    if not ($path | path exists) {
        if not $create {
            let action = (choose-recovery
                $"No worktree exists for branch `($branch)`."
                [
                    {id: "create", label: $"Create worktree for `($branch)`"}
                    {id: "cancel", label: "Cancel"}
                ]
                $"no worktree for branch `($branch)`. Non-interactive: rerun with `git wt go ($branch) --create`."
            )
            if $action != "create" { error make {msg: "cancelled"} }
        }
        create-worktree $branch $base (not $no_submodules) | ignore
    }
    cd $path
    $path
}

export def --env "wt rm" [
    name?: string                  # Branch/worktree to remove; default current branch
    --branch (-b)                  # Also delete the local branch after removing the worktree
    --force (-f)                   # Skip prompts and force removal/deletion where Git allows
] {
    let b = if ($name == null) { current-branch } else { $name }
    remove-worktree $b $branch $force $"git wt rm ($b)"
}

export def --env "wt done" [
    branch?: string                # Branch/worktree to finish; default current branch
    --force (-f)                   # Skip prompts and force worktree removal + branch deletion
] {
    let b = if ($branch == null) { current-branch } else { $branch }
    remove-worktree $b true $force $"git wt done ($b)"
}

export def "wt prune" [] {
    let l = (layout)
    let res = (git-dir-complete $l.git_dir ["worktree" "prune"])
    if $res.exit_code == 0 { return }
    error make {msg: (git-failure ["worktree" "prune"] $res.stderr)}
}

export def "wt repair" [] {
    let l = (layout)
    let res = (git-dir-complete $l.git_dir ["worktree" "repair"])
    if $res.exit_code == 0 { return }
    error make {msg: (git-failure ["worktree" "repair"] $res.stderr)}
}

# Thin native aliases for muscle memory / git parity.
export def --env "wt add" [
    branch?: string                # Filter local/remote/PR picker; omitted opens full picker
    --base (-b): string = $HEAD    # Explicitly create a new branch from this base instead of picking
    --switch (-s)                  # cd into the created worktree
    --no-submodules                # Skip submodule init in the created worktree
    --pr                           # Only show PR candidates (gh/fj); local/remote branches excluded
] {
    if $branch != null and $base != $HEAD {
        return (wt new $branch --base $base --switch=$switch --no-submodules=$no_submodules)
    }

    let selected = (pick-add-candidate $branch $pr)
    let path = (create-picked-worktree $selected (not $no_submodules))
    if $switch { cd $path }
    $path
}

export def "wt list" [] { wt ls }

export def "wt lock" [
    branch?: string                # Worktree/branch to lock; default current branch
    reason?: string                # Optional lock reason stored by Git
] {
    let b = if ($branch == null) { current-branch } else { $branch }
    let path = (path-for-branch $b)
    let l = (layout)

    if not ($path | path exists) {
        let action = (choose-recovery
            $"No worktree path exists for `($b)`: ($path)"
            [
                {id: "create", label: "Create worktree, then lock it"}
                {id: "cancel", label: "Cancel"}
            ]
            $"no worktree path for `($b)`. Non-interactive: run `git wt go ($b) --create`, then `git wt lock ($b)`."
        )
        if $action == "create" { create-worktree $b $HEAD true | ignore }
    }

    let args = if ($reason == null) { ["worktree" "lock" $path] } else { ["worktree" "lock" "--reason" $reason $path] }
    let res = (git-dir-complete $l.git_dir $args)
    if $res.exit_code == 0 { return }

    let msg = (git-failure $args $res.stderr)
    if ($msg | str contains "already locked") {
        let action = (choose-recovery
            $"Worktree `($b)` is already locked.\n\n($msg)"
            [
                {id: "relock", label: "Unlock, then lock again"}
                {id: "cancel", label: "Cancel"}
            ]
            $"($msg)\nNon-interactive: run `git wt unlock ($b)`, then retry `git wt lock ($b)`."
        )
        if $action == "relock" {
            git-dir-output $l.git_dir ["worktree" "unlock" $path] | ignore
            git-dir-output $l.git_dir $args | ignore
            return
        }
    }

    error make {msg: $msg}
}

export def "wt unlock" [
    branch?: string                # Worktree/branch to unlock; default current branch
] {
    let b = if ($branch == null) { current-branch } else { $branch }
    let path = (path-for-branch $b)
    let l = (layout)
    let res = (git-dir-complete $l.git_dir ["worktree" "unlock" $path])
    if $res.exit_code == 0 { return }

    let msg = (git-failure ["worktree" "unlock" $path] $res.stderr)
    if ($msg | str contains "not locked") {
        if $nu.is-interactive { print -e $"Worktree `($b)` is not locked. Nothing to do." }
        return
    }
    error make {msg: $msg}
}

export def "wt move" [
    branch: string                 # Worktree/branch to move
    dest: path                     # New worktree path
] {
    let l = (layout)
    let source = (path-for-branch $branch)

    if ($dest | path exists) {
        let action = (choose-recovery
            $"Destination already exists: ($dest)"
            [
                {id: "remove", label: "Delete destination, then move"}
                {id: "cancel", label: "Cancel"}
            ]
            $"destination exists: ($dest). Non-interactive: remove it first, then rerun `git wt move ($branch) ($dest)`."
        )
        if $action == "remove" { rm -rf $dest }
    }

    let res = (git-dir-complete $l.git_dir ["worktree" "move" $source $dest])
    if $res.exit_code == 0 { return $dest }

    let msg = (git-failure ["worktree" "move" $source $dest] $res.stderr)
    if ($msg | str contains "working trees containing submodules") {
        let action = (choose-recovery
            $"Git refuses to move worktrees containing submodules.\n\n($msg)\n\nSafe workaround: filesystem move, then `git worktree repair`."
            [
                {id: "repair", label: "Move directory, then repair worktree metadata"}
                {id: "cancel", label: "Cancel"}
            ]
            $"($msg)\nNon-interactive: run `mv ($source) ($dest)` then `git wt repair`."
        )
        if $action == "repair" {
            mv $source $dest
            git-dir-output $l.git_dir ["worktree" "repair" $dest] | ignore
            return $dest
        }
    }

    error make {msg: $msg}
}

export def --env "wt remove" [
    name?: string                  # Branch/worktree to remove; default current branch
    --branch (-b)                  # Also delete the local branch after removing the worktree
    --force (-f)                   # Skip prompts and force removal/deletion where Git allows
] {
    wt rm $name --branch=$branch --force=$force
}