use crate::gemini::errors::GeminiError; use crate::gemini::prompt::PromptBuilder; use crate::gemini::types::{GeminiResponse, OrganizationPlanResponse}; use crate::models::OrganizationPlan; use crate::storage::Cache; use reqwest::Client; use serde_json::json; use std::path::Path; use std::time::Duration; const DEFAULT_MODEL: &str = "gemini-3-flash-preview"; const DEFAULT_TIMEOUT_SECS: u64 = 30; const MAX_RETRIES: u32 = 3; pub struct GeminiClient { api_key: String, client: Client, base_url: String, #[allow(dead_code)] model: String, #[allow(dead_code)] timeout: Duration, } impl GeminiClient { pub fn new(api_key: String) -> Self { Self::with_model(api_key, DEFAULT_MODEL.to_string()) } pub fn with_model(api_key: String, model: String) -> Self { let timeout = Duration::from_secs(DEFAULT_TIMEOUT_SECS); let client = Self::build_client(timeout); let base_url = Self::build_base_url(&model); Self { api_key, client, base_url, model, timeout, } } fn build_client(timeout: Duration) -> Client { Client::builder() .timeout(timeout) .build() .unwrap_or_default() } fn build_base_url(model: &str) -> String { format!( "https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent", model ) } pub async fn organize_files( &self, filenames: Vec, ) -> Result { self.organize_files_with_cache(filenames, None, None).await } pub async fn organize_files_with_cache( &self, filenames: Vec, mut cache: Option<&mut Cache>, base_path: Option<&Path>, ) -> Result { let url = self.build_url(); if let (Some(cache), Some(base_path)) = (cache.as_ref(), base_path) && let Some(cached_response) = cache.get_cached_response(&filenames, base_path) { return Ok(cached_response); } let prompt = PromptBuilder::new(filenames.clone()).build_categorization_prompt(); let request_body = self.build_categorization_request(&prompt); let res = self.send_request_with_retry(&url, &request_body).await?; let plan = self.parse_categorization_response(res).await?; if let (Some(cache), Some(base_path)) = (cache.as_mut(), base_path) { cache.cache_response(&filenames, plan.clone(), base_path); } Ok(plan) } fn build_url(&self) -> String { format!("{}?key={}", self.base_url, self.api_key) } fn build_categorization_request(&self, prompt: &str) -> serde_json::Value { json!({ "contents": [{ "parts": [{ "text": prompt }] }], "generationConfig": { "response_mime_type": "application/json" } }) } async fn parse_categorization_response( &self, res: reqwest::Response, ) -> Result { if !res.status().is_success() { return Err(GeminiError::from_response(res).await); } let gemini_response: GeminiResponse = res.json().await.map_err(GeminiError::NetworkError)?; let raw_text = self.extract_text_from_response(&gemini_response)?; let plan_response: OrganizationPlanResponse = serde_json::from_str(&raw_text)?; Ok(plan_response.to_organization_plan()) } fn extract_text_from_response(&self, response: &GeminiResponse) -> Result { response .candidates .first() .ok_or_else(|| GeminiError::InvalidResponse("No candidates in response".to_string()))? .content .parts .first() .ok_or_else(|| GeminiError::InvalidResponse("No parts in content".to_string())) .map(|p| p.text.clone()) } async fn send_request_with_retry( &self, url: &str, request_body: &serde_json::Value, ) -> Result { let mut attempts = 0; let mut base_delay = Duration::from_secs(2); loop { attempts += 1; match self.client.post(url).json(request_body).send().await { Ok(response) => { if response.status().is_success() { return Ok(response); } let error = GeminiError::from_response(response).await; if error.is_retryable() && attempts < MAX_RETRIES { let delay = error.retry_delay().unwrap_or(base_delay); self.print_retry_message(&error, delay, attempts); tokio::time::sleep(delay).await; base_delay *= 2; continue; } return Err(error); } Err(e) => { if attempts < MAX_RETRIES { self.print_network_retry(&e, base_delay, attempts); tokio::time::sleep(base_delay).await; base_delay *= 2; continue; } return Err(GeminiError::NetworkError(e)); } } } } fn print_retry_message(&self, error: &GeminiError, delay: Duration, attempt: u32) { println!( "API Error: {}. Retrying in {} seconds (attempt {}/{})", error, delay.as_secs(), attempt, MAX_RETRIES ); } fn print_network_retry(&self, error: &reqwest::Error, delay: Duration, attempt: u32) { println!( "Network error: {}. Retrying in {} seconds (attempt {}/{})", error, delay.as_secs(), attempt, MAX_RETRIES ); } pub async fn get_ai_sub_category( &self, filename: &str, parent_category: &str, content: &str, ) -> String { let url = self.build_url(); let prompt = PromptBuilder::build_subcategory_prompt(filename, parent_category, content); let request_body = self.build_subcategory_request(&prompt); let res = match self.client.post(&url).json(&request_body).send().await { Ok(res) => res, Err(e) => { eprintln!( "Warning: Failed to get sub-category for {}: {}", filename, e ); return "General".to_string(); } }; self.parse_subcategory_response(res, filename).await } fn build_subcategory_request(&self, prompt: &str) -> serde_json::Value { json!({ "contents": [{ "parts": [{ "text": prompt }] }] }) } async fn parse_subcategory_response(&self, res: reqwest::Response, filename: &str) -> String { if !res.status().is_success() { eprintln!( "Warning: API returned error for {}: {}", filename, res.status() ); return "General".to_string(); } let gemini_response: GeminiResponse = match res.json().await { Ok(r) => r, Err(e) => { eprintln!("Warning: Failed to parse response for {}: {}", filename, e); return "General".to_string(); } }; self.extract_subcategory_from_response(&gemini_response, filename) } fn extract_subcategory_from_response( &self, response: &GeminiResponse, _filename: &str, ) -> String { match self.extract_text_from_response(response) { Ok(text) => { let sub_category = text.trim(); if sub_category.is_empty() { "General".to_string() } else { sub_category.to_string() } } Err(_) => "General".to_string(), } } }