1use crate::get_env;
2use git2::{DiffOptions, Error, ObjectType, Repository, StatusOptions, StatusShow};
3use serde::{Deserialize, Serialize};
4use std::{collections::BTreeMap, env, str};
5use tokio::process::Command;
6
7#[derive(Serialize, Deserialize, Debug, Default)]
8struct Prompt {
9    action: String,
10    branch: String,
11    remote: Vec<String>,
12    staged: bool,
13    status: String,
14    u_name: String,
15}
16
17pub fn render() {
18    if let Ok(path) = env::current_dir()
19        && let Ok(repo) = Repository::discover(path)
20    {
21        build_prompt(&repo);
22    }
23}
24
25fn build_prompt(repo: &Repository) {
26    let mut prompt = Prompt::default();
27
28    if let Ok(config) = repo.config() {
30        prompt.u_name = config
31            .get_string("user.name")
32            .unwrap_or_else(|_| String::new());
33    }
34
35    if let Ok(head) = repo.head() {
37        prompt.branch = head.shorthand().unwrap_or("(no branch)").to_string();
38    } else {
39        prompt.branch = "(no branch)".into();
40    }
41
42    if get_env("SLICK_PROMPT_GIT_FETCH") != "0" {
44        tokio::spawn(async move {
45            let mut cmd = Command::new("git");
46
47            cmd.env("GIT_TERMINAL_PROMPT", "0");
48
49            cmd.arg("-c")
50                .arg("gc.auto=0")
51                .arg("fetch")
52                .arg("--quiet")
53                .arg("--no-tags")
54                .arg("--no-recurse-submodules");
55
56            match cmd.output().await {
57                Ok(output) => {
58                    if !output.status.success() {
59                        eprintln!(
60                            "error: failed to execute git fetch: {}",
61                            str::from_utf8(&output.stderr).unwrap()
62                        );
63                    }
64                }
65                Err(e) => {
66                    eprintln!("error: failed to execute git fetch: {}", e);
67                }
68            }
69        });
70    }
71
72    let (ahead, behind) = is_ahead_behind_remote(repo);
74    if behind > 0 {
75        prompt.remote.push(format!(
76            "{}{}",
77            get_env("SLICK_PROMPT_GIT_REMOTE_BEHIND"),
78            behind
79        ));
80    }
81    if ahead > 0 {
82        prompt.remote.push(format!(
83            "{}{}",
84            get_env("SLICK_PROMPT_GIT_REMOTE_AHEAD"),
85            ahead
86        ));
87    }
88
89    if let Some(action) = get_action(repo) {
91        prompt.action = action;
92    }
93
94    if let Ok(status) = get_status(repo) {
96        prompt.status = status;
97    }
98
99    if let Ok(staged) = is_staged(repo) {
101        prompt.staged = staged;
102    }
103
104    if let Ok(serialized) = serde_json::to_string(&prompt) {
106        println!("{serialized}");
107    }
108}
109
110fn get_status(repo: &Repository) -> Result<String, Error> {
111    let mut status: Vec<String> = Vec::new();
112    let mut status_opt = StatusOptions::new();
113    status_opt
114        .show(StatusShow::IndexAndWorkdir)
115        .include_untracked(true)
116        .include_unmodified(false)
117        .no_refresh(false);
118
119    let statuses = repo.statuses(Some(&mut status_opt))?;
120    if !statuses.is_empty() {
121        let mut map: BTreeMap<&str, u32> = BTreeMap::new();
122        for entry in statuses.iter() {
123            let status = match entry.status() {
125                s if s.contains(git2::Status::INDEX_NEW)
126                    && s.contains(git2::Status::WT_MODIFIED) =>
127                {
128                    "AM"
129                }
130                s if s.contains(git2::Status::INDEX_MODIFIED)
131                    && s.contains(git2::Status::WT_MODIFIED) =>
132                {
133                    "MM"
134                }
135                s if s.contains(git2::Status::INDEX_MODIFIED)
136                    || s.contains(git2::Status::WT_MODIFIED) =>
137                {
138                    "M"
139                }
140                s if s.contains(git2::Status::INDEX_DELETED)
141                    || s.contains(git2::Status::WT_DELETED) =>
142                {
143                    "D"
144                }
145                s if s.contains(git2::Status::INDEX_RENAMED)
146                    || s.contains(git2::Status::WT_RENAMED) =>
147                {
148                    "R"
149                }
150                s if s.contains(git2::Status::INDEX_TYPECHANGE)
151                    || s.contains(git2::Status::WT_TYPECHANGE) =>
152                {
153                    "T"
154                }
155                s if s.contains(git2::Status::INDEX_NEW) => "A",
156                s if s.contains(git2::Status::WT_NEW) => "??",
157                s if s.contains(git2::Status::CONFLICTED) => "UU",
158                s if s.contains(git2::Status::IGNORED) => "!",
159                _ => "X",
160            };
161
162            *map.entry(status).or_insert(0) += 1;
163        }
164        for (k, v) in &map {
165            status.push(format!("{k} {v}"));
166        }
167    }
168    Ok(status.join(" "))
169}
170
171fn get_action(repo: &Repository) -> Option<String> {
172    let gitdir = repo.path();
173
174    for tmp in &[
175        gitdir.join("rebase-apply"),
176        gitdir.join("rebase"),
177        gitdir.join("..").join(".dotest"),
178    ] {
179        if tmp.join("rebasing").exists() {
180            return Some("rebase".to_string());
181        }
182        if tmp.join("applying").exists() {
183            return Some("am".to_string());
184        }
185        if tmp.exists() {
186            return Some("am/rebase".to_string());
187        }
188    }
189
190    for tmp in &[
191        gitdir.join("rebase-merge").join("interactive"),
192        gitdir.join(".dotest-merge").join("interactive"),
193    ] {
194        if tmp.exists() {
195            return Some("rebase-i".to_string());
196        }
197    }
198
199    for tmp in &[gitdir.join("rebase-merge"), gitdir.join(".dotest-merge")] {
200        if tmp.exists() {
201            return Some("rebase-m".to_string());
202        }
203    }
204
205    if gitdir.join("MERGE_HEAD").exists() {
206        return Some("merge".to_string());
207    }
208
209    if gitdir.join("BISECT_LOG").exists() {
210        return Some("bisect".to_string());
211    }
212
213    if gitdir.join("CHERRY_PICK_HEAD").exists() {
214        if gitdir.join("sequencer").exists() {
215            return Some("cherry-seq".to_string());
216        }
217        return Some("cherry".to_string());
218    }
219
220    if gitdir.join("sequencer").exists() {
221        return Some("cherry-or-revert".to_string());
222    }
223
224    None
225}
226
227fn is_ahead_behind_remote(repo: &Repository) -> (usize, usize) {
228    if let Ok(head) = repo.revparse_single("HEAD") {
229        let head = head.id();
230        if let Ok((upstream, _)) = repo.revparse_ext("@{u}") {
231            return match repo.graph_ahead_behind(head, upstream.id()) {
232                Ok((commits_ahead, commits_behind)) => (commits_ahead, commits_behind),
233                Err(_) => (0, 0),
234            };
235        }
236    }
237    (0, 0)
238}
239
240fn is_staged(repo: &Repository) -> Result<bool, Error> {
241    let mut opts = DiffOptions::new();
242    let obj = repo.head()?;
243    let tree = obj.peel(ObjectType::Tree)?;
244    let diff = repo.diff_tree_to_index(tree.as_tree(), None, Some(&mut opts))?;
245    let stats = diff.stats()?;
246    if stats.files_changed() > 0 || stats.insertions() > 0 || stats.deletions() > 0 {
247        return Ok(true);
248    }
249    Ok(false)
250    }