1use crate::config::{home_config_path, Config, ConfigError};
2use crate::openrouter::{ApiError, Client};
3use crate::ui;
4use colored::Colorize;
5
6fn print_splash_banner() {
7 let banner = r#"
8 ▄ ▄▄▄▄
9 ▀██████▀ ▄███████▄
10 ██ ▄ ▄ ██ ▀█▄ ▀█
11 ██ ▄███▄ ███▄███▄ ███▄███▄ ██ ▄█▀██ ██
12 ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ▄█
13 ▀█████▄▀███▀▄██ ██ ▀█▄██ ██ ▀█ ▀█▄ ▀▀▀▀▀▀
14 ▀██████▀▀
15"#;
16 println!("{}", banner.cyan().bold());
17}
18
19fn prompt_api_key(is_first_run: bool, existing_key: Option<&str>) -> String {
20 loop {
21 let mut prompt_text = "Enter OpenRouter API Key (sk-or-v1-...):".to_string();
22 if !is_first_run {
23 prompt_text.push_str(" (Leave blank to keep existing)");
24 }
25
26 let password = inquire::Password::new(&prompt_text)
27 .with_display_mode(inquire::PasswordDisplayMode::Masked)
28 .with_help_message("API key at https://openrouter.ai/keys")
29 .prompt()
30 .expect("User cancelled");
31
32 if password.is_empty() {
33 if is_first_run {
34 ui::error_message("API key is required on first run. Please enter your key.");
35 continue;
36 } else if let Some(key) = existing_key {
37 return key.to_string();
38 } else {
39 ui::error_message("No existing API key found. Please enter your key.");
40 continue;
41 }
42 }
43 break password;
44 }
45}
46
47pub fn run_setup_flow(is_first_run: bool) -> Result<Config, ConfigError> {
48 print_splash_banner();
49
50 let existing_key = if !is_first_run {
51 Config::load_from_path(&home_config_path())
52 .ok()
53 .map(|c| c.api_key)
54 } else {
55 None
56 };
57
58 let api_key = prompt_api_key(is_first_run, existing_key.as_deref());
59 let mut client = Client::new(api_key.clone());
60
61 let models = loop {
62 ui::fetching_models_message();
63 match client.fetch_models() {
64 Ok(m) => break m,
65 Err(ApiError::Unauthorized) => {
66 ui::error_message("Invalid API Key. Make sure you entered the correct key.");
67 let new_key = prompt_api_key(true, None);
68 client = Client::new(new_key);
69 continue;
70 }
71 Err(ApiError::RateLimited) => {
72 ui::rate_limited_message();
73 std::thread::sleep(std::time::Duration::from_secs(2));
74 continue;
75 }
76 Err(ApiError::EmptyResponse) => {
77 ui::error_message("No models available. Please try again.");
78 continue;
79 }
80 Err(ApiError::Forbidden) => {
81 return Err(ConfigError::ApiError(
82 "API key doesn't have access. Check permissions on OpenRouter.".into(),
83 ));
84 }
85 Err(e) => {
86 return Err(ConfigError::ApiError(format!(
87 "Failed to fetch models: {}. Please try again.",
88 e
89 )));
90 }
91 }
92 };
93
94 let model_ids: Vec<String> = models.iter().map(|m| m.id.clone()).collect();
95 ui::models_loaded(model_ids.len());
96
97 let selection = ui::model_select_prompt(&model_ids);
98 let model_id = if selection == "[ Type Manual Model ID... ]" {
99 ui::manual_model_prompt()
100 } else {
101 selection
102 };
103
104 let config = Config {
105 api_key,
106 model_id,
107 };
108
109 let path = home_config_path();
110 if let Err(e) = config.save(&path) {
111 eprintln!("Warning: Failed to save config: {}. Continuing anyway...", e);
112 }
113
114 ui::save_confirmation();
115 Ok(config)
116}