1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

use comfy_table::Table;
use utils::error::{Error, Result};

#[derive(Debug)]
pub enum GPTResult {
    Success(GPTResponse),
    Failure(String),
}

#[derive(Debug)]
pub struct GPTResponse {
    pub status: ResponseStatus,
    pub explanation: String,
    pub commands: Vec<String>,
}

#[derive(Debug)]
pub enum ResponseStatus {
    Success,
    NotValid,
    NotPossible,
}

pub fn decode_gpt_response(response: String) -> Result<GPTResult> {
    let response = response.trim();

    if response.starts_with("NOT_VALID:") {
        let explanation = response.replace("NOT_VALID:", "");

        return Ok(GPTResult::Success(GPTResponse {
            status: ResponseStatus::NotValid,
            explanation,
            commands: vec![],
        }));
    }

    if response.starts_with("NOT_POSSIBLE:") {
        let explanation = response.replace("NOT_POSSIBLE:", "");

        return Ok(GPTResult::Success(GPTResponse {
            status: ResponseStatus::NotPossible,
            explanation,
            commands: vec![],
        }));
    }

    if response.starts_with("git") {
        // Sanitize commands
        let commands: Vec<String> = response
            .split("&&")
            .map(|command| {
                if command.trim().starts_with("git") {
                    Ok(command.trim().replace("git ", ""))
                } else {
                    Err(Error::new("One or more commands do not start with 'git'"))
                }
            })
            .collect::<Result<Vec<String>>>()?;

        return Ok(GPTResult::Success(GPTResponse {
            status: ResponseStatus::Success,
            explanation: String::from(""),
            commands,
        }));
    }

    Err(Error::new("Could not decode GPT response"))
}

// implement display for GPTResult
// use comfy_table::Table to display the commands
impl std::fmt::Display for GPTResult {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            GPTResult::Success(response) => {
                let mut table = Table::new();

                table.set_header(vec!["#", "Command"]);

                response
                    .commands
                    .iter()
                    .enumerate()
                    .for_each(|(i, command)| {
                        table.add_row(vec![i.to_string(), format!("git {}", command.to_string())]);
                    });

                write!(f, "{}", table)
            }
            GPTResult::Failure(error) => write!(f, "{}", error),
        }
    }
}