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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

// use clap::Arg;
use clap::{value_parser, Arg, ArgMatches};

use crate::config::Config;

pub fn cli() -> clap::Command {
    clap::Command::new("energy")
        .about("Invokes commands related to database budgets")
        .args_conflicts_with_subcommands(true)
        .subcommand_required(true)
        .subcommands(get_energy_subcommands())
}

fn get_energy_subcommands() -> Vec<clap::Command> {
    vec![
        clap::Command::new("status")
            .about("Show current energy balance for an identity")
            .arg(
                Arg::new("identity")
                    .help("The identity to check the balance for")
                    .long_help(
                    "The identity to check the balance for. If no identity is provided, the default one will be used.",
                ),
            ),
        clap::Command::new("set-balance")
            .about("Update the current budget balance for a database")
            .arg(
                Arg::new("balance")
                    .required(true)
                    .value_parser(value_parser!(u64))
                    .help("The balance value to set"),
            )
            .arg(
                Arg::new("identity")
                    .help("The identity to set a balance for")
                    .long_help(
                        "The identity to set a balance for. If no identity is provided, the default one will be used.",
                    ),
            )
            .arg(
                Arg::new("quiet")
                    .long("quiet")
                    .short('q')
                    .action(clap::ArgAction::SetTrue)
                    .help("Runs command in silent mode"),
            ),
    ]
}

async fn exec_subcommand(config: Config, cmd: &str, args: &ArgMatches) -> Result<(), anyhow::Error> {
    match cmd {
        "status" => exec_status(config, args).await,
        "set-balance" => exec_update_balance(config, args).await,
        unknown => Err(anyhow::anyhow!("Invalid subcommand: {}", unknown)),
    }
}

pub async fn exec(config: Config, args: &ArgMatches) -> Result<(), anyhow::Error> {
    let (cmd, subcommand_args) = args.subcommand().expect("Subcommand required");
    exec_subcommand(config, cmd, subcommand_args).await
}

async fn exec_update_balance(config: Config, args: &ArgMatches) -> Result<(), anyhow::Error> {
    // let project_name = args.value_of("project name").unwrap();
    let hex_id = args.get_one::<String>("identity");
    let balance = *args.get_one::<u64>("balance").unwrap();
    let quiet = args.get_flag("quiet");

    let hex_id = hex_id_or_default(hex_id, &config);
    let res = set_balance(&reqwest::Client::new(), &config, hex_id, balance).await?;

    if !quiet {
        println!("{}", res.text().await?);
    }

    Ok(())
}

async fn exec_status(config: Config, args: &ArgMatches) -> Result<(), anyhow::Error> {
    // let project_name = args.value_of("project name").unwrap();
    let hex_id = args.get_one::<String>("identity");
    let hex_id = hex_id_or_default(hex_id, &config);

    let status = reqwest::Client::new()
        .get(format!("{}/energy/{}", config.get_host_url(), hex_id))
        .send()
        .await?
        .error_for_status()?
        .text()
        .await?;

    println!("{}", status);

    Ok(())
}

fn hex_id_or_default<'a>(hex_id: Option<&'a String>, config: &'a Config) -> &'a String {
    hex_id.unwrap_or_else(|| &config.get_default_identity_config().unwrap().identity)
}

pub(super) async fn set_balance(
    client: &reqwest::Client,
    config: &Config,
    hex_identity: &str,
    balance: u64,
) -> anyhow::Result<reqwest::Response> {
    // TODO: this really should be form data in POST body, not query string parameter, but gotham
    // does not support that on the server side without an extension.
    // see https://github.com/gotham-rs/gotham/issues/11
    client
        .post(format!("{}/energy/{}", config.get_host_url(), hex_identity))
        .query(&[("balance", balance)])
        .send()
        .await?
        .error_for_status()
        .map_err(|e| e.into())
}