1use crate::code_generator::CodeGenerator;
16use crate::code_quality_enforcer::CodeQualityEnforcer;
17use crate::code_validator::CodeValidator;
18use crate::conflict_detector::ConflictDetector;
19use crate::conflict_resolver::{ConflictResolver, ConflictStrategy};
20use crate::error::GenerationError;
21use crate::output_writer::OutputWriter;
22use crate::prompt_builder::PromptBuilder;
23use crate::report_generator::{GenerationResult, GenerationStats, ReportGenerator};
24use crate::review_engine::ReviewEngine;
25use crate::spec_processor::SpecProcessor;
26use crate::templates::TemplateEngine;
27use ricecoder_providers::provider::Provider;
28use ricecoder_specs::models::Spec;
29use std::path::PathBuf;
30use std::time::Instant;
31
32#[derive(Debug, Clone)]
34pub struct GenerationManagerConfig {
35 pub project_root: PathBuf,
37 pub validate: bool,
39 pub review: bool,
41 pub dry_run: bool,
43 pub conflict_strategy: ConflictStrategy,
45 pub max_retries: usize,
47 pub use_templates: bool,
49}
50
51impl Default for GenerationManagerConfig {
52 fn default() -> Self {
53 Self {
54 project_root: PathBuf::from("."),
55 validate: true,
56 review: false,
57 dry_run: false,
58 conflict_strategy: ConflictStrategy::Prompt,
59 max_retries: 3,
60 use_templates: false,
61 }
62 }
63}
64
65pub struct GenerationManager {
76 config: GenerationManagerConfig,
77 spec_processor: SpecProcessor,
78 prompt_builder: PromptBuilder,
79 code_generator: CodeGenerator,
80 #[allow(dead_code)]
81 template_engine: TemplateEngine,
82 code_quality_enforcer: CodeQualityEnforcer,
83 code_validator: CodeValidator,
84 conflict_detector: ConflictDetector,
85 #[allow(dead_code)]
86 conflict_resolver: ConflictResolver,
87 output_writer: OutputWriter,
88 review_engine: ReviewEngine,
89 #[allow(dead_code)]
90 report_generator: ReportGenerator,
91}
92
93impl std::fmt::Debug for GenerationManager {
94 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95 f.debug_struct("GenerationManager")
96 .field("config", &self.config)
97 .finish()
98 }
99}
100
101impl GenerationManager {
102 pub fn new(project_root: PathBuf) -> Self {
104 Self::with_config(GenerationManagerConfig {
105 project_root: project_root.clone(),
106 ..Default::default()
107 })
108 }
109
110 pub fn with_config(config: GenerationManagerConfig) -> Self {
112 let project_root = config.project_root.clone();
113
114 Self {
115 config,
116 spec_processor: SpecProcessor::new(),
117 prompt_builder: PromptBuilder::new(project_root.clone()),
118 code_generator: CodeGenerator::new(),
119 template_engine: TemplateEngine::new(),
120 code_quality_enforcer: CodeQualityEnforcer::new(),
121 code_validator: CodeValidator::new(),
122 conflict_detector: ConflictDetector::new(),
123 conflict_resolver: ConflictResolver::new(),
124 output_writer: OutputWriter::new(),
125 review_engine: ReviewEngine::new(),
126 report_generator: ReportGenerator,
127 }
128 }
129
130 #[allow(clippy::too_many_arguments)]
160 pub async fn generate(
161 &self,
162 spec: &Spec,
163 target_path: PathBuf,
164 _language: String,
165 provider: Option<&dyn Provider>,
166 model: &str,
167 temperature: f32,
168 max_tokens: usize,
169 ) -> Result<GenerationResult, GenerationError> {
170 let start_time = Instant::now();
171
172 let plan = self
175 .spec_processor
176 .process(spec)
177 .map_err(|e| GenerationError::SpecError(format!("Failed to process spec: {}", e)))?;
178
179 let prompt = self
182 .prompt_builder
183 .build(&plan, None, None)
184 .map_err(|e| GenerationError::PromptError(format!("Failed to build prompts: {}", e)))?;
185
186 let mut generated_files = if self.config.use_templates {
189 let provider = provider.ok_or_else(|| {
193 GenerationError::GenerationFailed(
194 "AI provider required for code generation".to_string(),
195 )
196 })?;
197 self.code_generator
198 .generate(provider, &prompt, model, temperature, max_tokens)
199 .await
200 .map_err(|e| {
201 GenerationError::GenerationFailed(format!("Code generation failed: {}", e))
202 })?
203 } else {
204 let provider = provider.ok_or_else(|| {
206 GenerationError::GenerationFailed(
207 "AI provider required for AI-based generation".to_string(),
208 )
209 })?;
210 self.code_generator
211 .generate(provider, &prompt, model, temperature, max_tokens)
212 .await
213 .map_err(|e| {
214 GenerationError::GenerationFailed(format!("AI generation failed: {}", e))
215 })?
216 };
217
218 generated_files = self
221 .code_quality_enforcer
222 .enforce(generated_files)
223 .map_err(|e| {
224 GenerationError::GenerationFailed(format!("Quality enforcement failed: {}", e))
225 })?;
226
227 let validation_result = if self.config.validate {
230 self.code_validator
231 .validate(&generated_files)
232 .map_err(|e| {
233 GenerationError::ValidationFailed(format!("Validation failed: {}", e))
234 })?
235 } else {
236 crate::models::ValidationResult::default()
237 };
238
239 let conflicts = self
242 .conflict_detector
243 .detect(&generated_files, &target_path)
244 .map_err(|e| {
245 GenerationError::GenerationFailed(format!("Conflict detection failed: {}", e))
246 })?;
247
248 let review_result = if self.config.review {
251 Some(
252 self.review_engine
253 .review(&generated_files, spec)
254 .map_err(|e| {
255 GenerationError::GenerationFailed(format!("Review failed: {}", e))
256 })?,
257 )
258 } else {
259 None
260 };
261
262 if !self.config.dry_run && validation_result.valid {
265 self.output_writer
266 .write(&generated_files, &target_path, &conflicts)
267 .map_err(|e| GenerationError::WriteFailed(format!("Write failed: {}", e)))?;
268 }
269
270 let elapsed = start_time.elapsed();
272 let stats = GenerationStats {
273 tokens_used: prompt.estimated_tokens,
274 time_elapsed: elapsed,
275 files_generated: generated_files.len(),
276 lines_generated: generated_files
277 .iter()
278 .map(|f| f.content.lines().count())
279 .sum(),
280 conflicts_detected: conflicts.len(),
281 conflicts_resolved: conflicts.len(), };
283
284 let mut result =
286 GenerationResult::new(generated_files, validation_result, conflicts, stats);
287
288 if let Some(review) = review_result {
289 result = result.with_review(review);
290 }
291
292 Ok(result)
293 }
294
295 #[allow(clippy::too_many_arguments)]
317 pub async fn generate_with_retries(
318 &self,
319 spec: &Spec,
320 target_path: PathBuf,
321 language: String,
322 provider: Option<&dyn Provider>,
323 model: &str,
324 temperature: f32,
325 max_tokens: usize,
326 ) -> Result<GenerationResult, GenerationError> {
327 let mut last_error = None;
328
329 for attempt in 0..self.config.max_retries {
330 match self
331 .generate(
332 spec,
333 target_path.clone(),
334 language.clone(),
335 provider,
336 model,
337 temperature,
338 max_tokens,
339 )
340 .await
341 {
342 Ok(result) => return Ok(result),
343 Err(e) => {
344 last_error = Some(e);
345 if attempt < self.config.max_retries - 1 {
346 let backoff_ms = 100 * (2_u64.pow(attempt as u32));
348 tokio::time::sleep(tokio::time::Duration::from_millis(backoff_ms)).await;
349 }
350 }
351 }
352 }
353
354 Err(last_error.unwrap_or_else(|| {
355 GenerationError::GenerationFailed("Generation failed after all retries".to_string())
356 }))
357 }
358
359 pub fn config(&self) -> &GenerationManagerConfig {
361 &self.config
362 }
363
364 pub fn set_config(&mut self, config: GenerationManagerConfig) {
366 self.config = config;
367 }
368}
369
370#[cfg(test)]
371mod tests {
372 use super::*;
373
374 #[test]
375 fn test_generation_manager_creation() {
376 let manager = GenerationManager::new(PathBuf::from("."));
377 assert_eq!(manager.config.validate, true);
378 assert_eq!(manager.config.dry_run, false);
379 }
380
381 #[test]
382 fn test_generation_manager_with_config() {
383 let config = GenerationManagerConfig {
384 project_root: PathBuf::from("."),
385 validate: false,
386 review: true,
387 dry_run: true,
388 conflict_strategy: ConflictStrategy::Skip,
389 max_retries: 5,
390 use_templates: true,
391 };
392 let manager = GenerationManager::with_config(config.clone());
393 assert_eq!(manager.config.validate, false);
394 assert_eq!(manager.config.review, true);
395 assert_eq!(manager.config.dry_run, true);
396 assert_eq!(manager.config.max_retries, 5);
397 assert_eq!(manager.config.use_templates, true);
398 }
399}