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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

use crate::storage::KvStore;
use config::{Config, File};
use dirs::config_dir;
use std::{error::Error, path::PathBuf};

pub fn initialize_storage() -> KvStore {
    let db_path = PathBuf::from("./data/db");

    if !db_path.parent().unwrap().exists() {
        std::fs::create_dir_all(db_path.parent().unwrap())
            .expect("Failed to create storage directory");
    }

    KvStore::new(db_path)
}

pub fn get_config() -> Config {
    let config_builder = Config::builder();
    let config_dir = config_dir().expect("Failed to get config directory");
    let config_path = config_dir.join(format!("{}/ai.toml", "askitty"));

    // if no config file exists, generate one
    if !config_path.exists() {
        let _ = config_setup_cli();
    }

    config_builder
        .add_source(File::from(config_path))
        .build()
        .expect("Failed to build config")
}

fn genetate_config(
    llm_host: &str,
    llm_api_key: &str,
    llm_model: &str,
    im_host: &str,
    im_api_key: &str,
    im_model: &str,
) {
    let config = format!(
        r#"
[llm]
host = "{}"
api_key = "{}"
llm_model = "{}"

[im]
host = "{}"
api_key = "{}"
image_model = "{}"
"#,
        llm_host, llm_api_key, llm_model, im_host, im_api_key, im_model
    );

    let config_dir = config_dir().expect("Failed to get config directory");
    let config_path = config_dir.join(format!("{}/ai.toml", "askitty"));
    if !config_path.parent().unwrap().exists() {
        std::fs::create_dir_all(config_path.parent().unwrap())
            .expect("Failed to create config directory");
    }

    std::fs::write(config_path, config).expect("Failed to write config file");
}

pub fn config_setup_cli() -> Result<(), Box<dyn Error>> {
    let mut llm_host = String::from("https://api.openai.com/v1");
    let mut llm_model = String::from("gpt-4");
    let llm_api_key;

    let mut im_host = String::from("https://api.openai.com/v1");
    let mut im_model = String::from("dall-e-3");
    let im_api_key;

    println!("Welcome to the askitty setup");
    println!("Let's start by setting up the language model");

    println!("Default LLM host: {}", llm_host);
    println!("Enter the LLM host (default: {}): ", llm_host);

    let mut input = String::new();
    std::io::stdin().read_line(&mut input).unwrap();
    input = input.trim().to_string();

    if !input.is_empty() {
        llm_host = input.clone();
    }

    println!("Default LLM model: {}", llm_model);
    println!("Enter the LLM model (default: {}): ", llm_model);

    input.clear();
    std::io::stdin().read_line(&mut input).unwrap();
    input = input.trim().to_string();

    if !input.is_empty() {
        llm_model = input.clone();
    }

    println!("Enter the LLM API key: ");
    input.clear();
    std::io::stdin().read_line(&mut input).unwrap();
    llm_api_key = input.trim().to_string();

    println!("Let's start by setting up the image model");

    println!("Default IM host: {}", im_host);
    println!("Enter the IM host (default: {}): ", im_host);

    input.clear();
    std::io::stdin().read_line(&mut input).unwrap();
    input = input.trim().to_string();

    if !input.is_empty() {
        im_host = input.trim().to_string();
    }

    println!("Default IM model: {}", im_model);
    print!("Enter the IM model (default: {}): ", im_model);
    input.clear();
    std::io::stdin().read_line(&mut input).unwrap();
    input = input.trim().to_string();

    if !input.is_empty() {
        im_model = input.trim().to_string();
    }

    println!("Enter the IM API key (default: {}): ", llm_api_key);
    input.clear();
    std::io::stdin().read_line(&mut input).unwrap();

    if !input.trim().is_empty() {
        im_api_key = input.trim().to_string();
    } else {
        im_api_key = llm_api_key.clone();
    }

    genetate_config(
        &llm_host,
        &llm_api_key,
        &llm_model,
        &im_host,
        &im_api_key,
        &im_model,
    );

    Ok(())
}