1use anyhow::Result;
2use futures_util::StreamExt;
3use inquire::Text;
4use inquire::ui::{Attributes, Color, RenderConfig, StyleSheet, Styled};
5use std::collections::HashMap;
6use std::io::{self, Write};
7
8use super::command::{Input, SlashCommand, SlashCommandCompleter, parse_input};
9use super::ui;
10use crate::config::{CustomStyle, ResolvedConfig};
11use crate::style;
12use crate::translation::{TranslationClient, TranslationRequest};
13use crate::ui::{Spinner, Style};
14
15#[derive(Debug, Clone)]
19pub struct SessionConfig {
20 pub resolved: ResolvedConfig,
22 pub custom_styles: HashMap<String, CustomStyle>,
24}
25
26impl SessionConfig {
27 #[allow(clippy::missing_const_for_fn)] pub fn new(resolved: ResolvedConfig, custom_styles: HashMap<String, CustomStyle>) -> Self {
30 Self {
31 resolved,
32 custom_styles,
33 }
34 }
35}
36
37pub struct ChatSession {
41 config: SessionConfig,
42 client: TranslationClient,
43}
44
45impl ChatSession {
46 pub fn new(config: SessionConfig) -> Self {
48 let client = TranslationClient::new(
49 config.resolved.endpoint.clone(),
50 config.resolved.api_key.clone(),
51 );
52 Self { config, client }
53 }
54
55 pub async fn run(&mut self) -> Result<()> {
56 ui::print_header();
57
58 let prompt_style = Styled::new("❯")
59 .with_fg(Color::LightBlue)
60 .with_attr(Attributes::BOLD);
61 let mut render_config = RenderConfig::default()
62 .with_prompt_prefix(prompt_style)
63 .with_answered_prompt_prefix(prompt_style);
64
65 render_config.option = StyleSheet::new().with_fg(Color::Grey);
67 render_config.selected_option = Some(StyleSheet::new().with_fg(Color::DarkMagenta));
69
70 loop {
71 let input = Text::new("")
72 .with_render_config(render_config)
73 .with_autocomplete(SlashCommandCompleter)
74 .with_help_message("Type text to translate, /help for commands, Ctrl+C to quit")
75 .prompt();
76
77 match input {
78 Ok(line) => match parse_input(&line) {
79 Input::Empty => {}
80 Input::Command(cmd) => {
81 if !self.handle_command(cmd) {
82 break;
83 }
84 }
85 Input::Text(text) => {
86 self.translate_and_print(&text).await?;
87 }
88 },
89 Err(
90 inquire::InquireError::OperationCanceled
91 | inquire::InquireError::OperationInterrupted,
92 ) => {
93 println!(); break;
95 }
96 Err(e) => return Err(e.into()),
97 }
98 }
99
100 ui::print_goodbye();
101 Ok(())
102 }
103
104 fn handle_command(&mut self, cmd: SlashCommand) -> bool {
105 match cmd {
106 SlashCommand::Config => {
107 ui::print_config(&self.config);
108 true
109 }
110 SlashCommand::Help => {
111 ui::print_help();
112 true
113 }
114 SlashCommand::Quit => false,
115 SlashCommand::Set { key, value } => {
116 self.handle_set(&key, value.as_deref());
117 true
118 }
119 SlashCommand::Unknown(cmd) => {
120 ui::print_error(&format!("Unknown command: /{cmd}"));
121 true
122 }
123 }
124 }
125
126 fn handle_set(&mut self, key: &str, value: Option<&str>) {
127 match key {
128 "style" => self.set_style(value),
129 "to" => self.set_to(value),
130 "model" => self.set_model(value),
131 "" => {
132 println!("Usage: /set <key> <value>");
133 println!("Keys: style, to, model");
134 }
135 _ => {
136 ui::print_error(&format!("Unknown setting: {key}"));
137 println!("Available: style, to, model");
138 }
139 }
140 }
141
142 fn set_style(&mut self, value: Option<&str>) {
143 let Some(key) = value else {
144 self.config.resolved.style_name = None;
146 self.config.resolved.style_prompt = None;
147 println!("{} Style cleared", Style::success("✓"));
148 return;
149 };
150
151 let resolved = match style::resolve_style(key, &self.config.custom_styles) {
153 Ok(r) => r,
154 Err(e) => {
155 ui::print_error(&e.to_string());
156 return;
157 }
158 };
159
160 self.config.resolved.style_name = Some(key.to_string());
161 self.config.resolved.style_prompt = Some(resolved.prompt().to_string());
162 println!(
163 "{} Style set to {}\n",
164 Style::success("✓"),
165 Style::value(key)
166 );
167 }
168
169 fn set_to(&mut self, value: Option<&str>) {
170 match value {
171 None => {
172 ui::print_error("Usage: /set to <language>");
173 }
174 Some(lang) => {
175 self.config.resolved.target_language = lang.to_string();
176 println!(
177 "{} Target language set to {}",
178 Style::success("✓"),
179 Style::value(lang)
180 );
181 }
182 }
183 }
184
185 fn set_model(&mut self, value: Option<&str>) {
186 match value {
187 None => {
188 ui::print_error("Usage: /set model <name>");
189 }
190 Some(model) => {
191 self.config.resolved.model = model.to_string();
192 println!(
193 "{} Model set to {}",
194 Style::success("✓"),
195 Style::value(model)
196 );
197 }
198 }
199 }
200
201 async fn translate_and_print(&self, text: &str) -> Result<()> {
202 let request = TranslationRequest {
203 source_text: text.to_string(),
204 target_language: self.config.resolved.target_language.clone(),
205 model: self.config.resolved.model.clone(),
206 endpoint: self.config.resolved.endpoint.clone(),
207 style: self.config.resolved.style_prompt.clone(),
208 };
209
210 let spinner = Spinner::new("Translating...");
211
212 let mut stream = self.client.translate_stream(&request).await?;
213 let mut first_chunk = true;
214
215 while let Some(chunk_result) = stream.next().await {
216 let chunk = chunk_result?;
217
218 if first_chunk {
219 spinner.stop();
220 first_chunk = false;
221 }
222
223 print!("{chunk}");
224 io::stdout().flush()?;
225 }
226
227 if first_chunk {
228 spinner.stop();
229 }
230
231 println!();
232 println!();
233 Ok(())
234 }
235}
236
237#[cfg(test)]
238mod tests {
239 use super::*;
240
241 #[test]
242 fn test_session_config_new() {
243 let mut custom_styles = HashMap::new();
244 custom_styles.insert(
245 "my_style".to_string(),
246 CustomStyle {
247 description: "My description".to_string(),
248 prompt: "My custom prompt".to_string(),
249 },
250 );
251
252 let resolved = ResolvedConfig {
253 provider_name: "ollama".to_string(),
254 endpoint: "http://localhost:11434".to_string(),
255 model: "gemma3:12b".to_string(),
256 api_key: None,
257 target_language: "ja".to_string(),
258 style_name: Some("casual".to_string()),
259 style_prompt: Some("Use a casual tone.".to_string()),
260 };
261
262 let config = SessionConfig::new(resolved, custom_styles);
263
264 assert_eq!(config.resolved.provider_name, "ollama");
265 assert_eq!(config.resolved.endpoint, "http://localhost:11434");
266 assert_eq!(config.resolved.model, "gemma3:12b");
267 assert!(config.resolved.api_key.is_none());
268 assert_eq!(config.resolved.target_language, "ja");
269 assert_eq!(config.resolved.style_name, Some("casual".to_string()));
270 assert_eq!(
271 config.resolved.style_prompt,
272 Some("Use a casual tone.".to_string())
273 );
274 assert_eq!(
275 config.custom_styles.get("my_style").map(|s| &s.prompt),
276 Some(&"My custom prompt".to_string())
277 );
278 }
279}