use crate::cache::Cache; use crate::files::{FileCategory, OrganizationPlan}; use crate::gemini_errors::GeminiError; use reqwest::Client; use serde::Deserialize; use serde_json::json; use std::path::Path; use std::time::Duration; #[derive(Deserialize, Default)] struct GeminiResponse { candidates: Vec, } #[derive(Deserialize)] struct Candidate { content: Content, } #[derive(Deserialize)] struct Content { parts: Vec, } #[derive(Deserialize)] struct Part { text: String, } #[derive(Deserialize)] struct FileCategoryResponse { filename: String, category: String, } #[derive(Deserialize)] struct OrganizationPlanResponse { files: Vec, } pub struct GeminiClient { api_key: String, client: Client, base_url: String, model: String, timeout: Duration, } impl GeminiClient { pub fn new(api_key: String) -> Self { Self::with_model(api_key, "gemini-3-flash-preview".to_string()) } pub fn with_model(api_key: String, model: String) -> Self { Self { api_key, client: Client::builder() .timeout(Duration::from_secs(30)) .build() .unwrap_or_default(), base_url: format!( "https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent", model ), model, timeout: Duration::from_secs(30), } } /// Takes a list of filenames and asks Gemini to categorize them pub async fn organize_files( &self, filenames: Vec, ) -> Result { self.organize_files_with_cache(filenames, None, None).await } /// Takes a list of filenames and asks Gemini to categorize them with caching support pub async fn organize_files_with_cache( &self, filenames: Vec, mut cache: Option<&mut Cache>, base_path: Option<&Path>, ) -> Result { let url = format!("{}?key={}", self.base_url, self.api_key); // Check cache first if available if let (Some(cache_ref), Some(base_path)) = (cache.as_ref(), base_path) && let Some(cached_response) = cache_ref.get_cached_response(&filenames, base_path) { return Ok(cached_response); } // 1. Construct the Prompt let file_list_str = filenames.join(", "); let prompt = format!( "I have these files in my Downloads folder: [{}]. \ Categorize them into these folders: 'Images', 'Documents', 'Installers', 'Music', 'Archives', 'Code', 'Misc'. \ Return ONLY a JSON object with this structure: {{ 'files': [ {{ 'filename': 'name', 'category': 'folder' }} ] }}", file_list_str ); // 2. Build Request with JSON Mode enforced let request_body = json!({ "contents": [{ "parts": [{ "text": prompt }] }], "generationConfig": { "response_mime_type": "application/json" } }); // 3. Send with retry logic let res = self.send_request_with_retry(&url, &request_body).await?; // 4. Parse if res.status().is_success() { let gemini_response: GeminiResponse = res.json().await.map_err(GeminiError::NetworkError)?; // Extract raw JSON string from Gemini using proper structs let raw_text = &gemini_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()))? .text; // Deserialize into our temporary response struct let plan_response: OrganizationPlanResponse = serde_json::from_str(raw_text)?; // Manually map to the final OrganizationPlan let plan = OrganizationPlan { files: plan_response .files .into_iter() .map(|f| FileCategory { filename: f.filename, category: f.category, sub_category: String::new(), // Initialize with empty sub_category }) .collect(), }; // Cache the response if cache is available if let (Some(cache), Some(base_path)) = (cache.as_mut(), base_path) { cache.cache_response(&filenames, plan.clone(), base_path); } Ok(plan) } else { Err(GeminiError::from_response(res).await) } } /// Send request with retry logic for retryable errors async fn send_request_with_retry( &self, url: &str, request_body: &serde_json::Value, ) -> Result { let mut attempts = 0; let max_attempts = 3; 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_attempts { let delay = error.retry_delay().unwrap_or(base_delay); println!( "API Error: {}. Retrying in {} seconds (attempt {}/{})", error, delay.as_secs(), attempts, max_attempts ); tokio::time::sleep(delay).await; base_delay *= 2; continue; } return Err(error); } Err(e) => { if attempts < max_attempts { println!( "Network error: {}. Retrying in {} seconds (attempt {}/{})", e, base_delay.as_secs(), attempts, max_attempts ); tokio::time::sleep(base_delay).await; base_delay *= 2; continue; } return Err(GeminiError::NetworkError(e)); } } } } pub async fn get_ai_sub_category( &self, filename: &str, parent_category: &str, content: &str, ) -> String { let url = format!("{}?key={}", self.base_url, self.api_key); let prompt = format!( "I have a file named '{}' inside the '{}' folder. Here is the first 1000 characters of the content:\n---\n{}\n---\nBased on this, suggest a single short sub-folder name (e.g., 'Invoices', 'Notes', 'Config'). Return ONLY the name of the sub-folder. Do not use markdown or explanations.", filename, parent_category, content ); let request_body = json!({ "contents": [{ "parts": [{ "text": 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(); } }; if res.status().is_success() { 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(); } }; let sub_category = gemini_response .candidates .first() .and_then(|c| c.content.parts.first()) .map(|p| p.text.trim()) .unwrap_or("General") .to_string(); if sub_category.is_empty() { "General".to_string() } else { sub_category } } else { eprintln!( "Warning: API returned error for {}: {}", filename, res.status() ); "General".to_string() } } }