opencrates 3.0.1

use anyhow::Result;
use async_trait::async_trait;
use std::collections::HashMap;
use tracing::{info, debug};

use crate::core::CrateContext;
use crate::providers::{OpenAIProvider, GenerationRequest, LegacyLLMProvider};
use crate::utils::templates::TemplateManager;
use super::{Stage, Architecture};

// TODO: document this
// TODO: document this
// TODO: document this
pub struct CrateGenerationStage {
    openai_provider: OpenAIProvider,
    template_manager: TemplateManager,
}

impl CrateGenerationStage {
    // TODO: document this
    // TODO: document this
    // TODO: document this
    // TODO: document this
    pub fn new(openai_provider: OpenAIProvider, template_manager: TemplateManager) -> Self {
        Self {
            openai_provider,
            template_manager,
        }
    }

    async fn generate_code_files(&self, architecture: &Architecture) -> Result<HashMap<String, String>> {
        let mut files = HashMap::new();
        
        // Generate lib.rs
        let lib_content = self.generate_lib_rs(architecture).await?;
        files.insert("src/lib.rs".to_string(), lib_content);
        
        // Generate module files
        for module in &architecture.modules {
            let module_content = self.generate_module_file(module, architecture).await?;
            files.insert(format!("src/{}.rs", module.name), module_content);
        }
        
        // Generate test files
        let test_content = self.generate_test_file(architecture).await?;
        files.insert("tests/integration_test.rs".to_string(), test_content);
        
        // Generate example files
        let example_content = self.generate_example_file(architecture).await?;
        files.insert("examples/basic_usage.rs".to_string(), example_content);
        
        Ok(files)
    }

    async fn generate_lib_rs(&self, architecture: &Architecture) -> Result<String> {
        let prompt = format!(
            "Generate a lib.rs file for a Rust crate named '{}' with modules: {:?}. 
            Include proper documentation, re-exports, and feature gates.",
            architecture.concept.spec.name,
            architecture.modules.iter().map(|m| &m.name).collect::<Vec<_>>()
        );

        let request = GenerationRequest {
            prompt,
            context: serde_json::to_string(architecture)?,
            max_tokens: Some(1024),
            temperature: Some(0.5),
            model: None,
        };

        let response = self.openai_provider.generate(request).await?;
        Ok(response.content)
    }

    async fn generate_module_file(&self, module: &super::Module, architecture: &Architecture) -> Result<String> {
        let prompt = format!(
            "Generate a Rust module file for '{}' with purpose: '{}'. 
            Exports: {:?}, Dependencies: {:?}. 
            Include proper documentation, error handling, and tests.",
            module.name, module.purpose, module.exports, module.dependencies
        );

        let request = GenerationRequest {
            prompt,
            context: format!("Module: {}\nArchitecture: {}", 
                           serde_json::to_string(module)?,
                           serde_json::to_string(architecture)?),
            max_tokens: Some(2048),
            temperature: Some(0.6),
            model: None,
        };

        let response = self.openai_provider.generate(request).await?;
        Ok(response.content)
    }

    async fn generate_test_file(&self, architecture: &Architecture) -> Result<String> {
        let prompt = format!(
            "Generate comprehensive integration tests for a Rust crate named '{}'. 
            Test all major functionality and edge cases.",
            architecture.concept.spec.name
        );

        let request = GenerationRequest {
            prompt,
            context: serde_json::to_string(architecture)?,
            max_tokens: Some(1536),
            temperature: Some(0.4),
            model: None,
        };

        let response = self.openai_provider.generate(request).await?;
        Ok(response.content)
    }

    async fn generate_example_file(&self, architecture: &Architecture) -> Result<String> {
        let prompt = format!(
            "Generate a basic usage example for a Rust crate named '{}'. 
            Show how to use the main functionality with clear comments.",
            architecture.concept.spec.name
        );

        let request = GenerationRequest {
            prompt,
            context: serde_json::to_string(architecture)?,
            max_tokens: Some(1024),
            temperature: Some(0.6),
            model: None,
        };

        let response = self.openai_provider.generate(request).await?;
        Ok(response.content)
    }

    async fn generate_documentation(&self, architecture: &Architecture) -> Result<HashMap<String, String>> {
        let mut docs = HashMap::new();
        
        // Generate README.md
        let readme_content = format!(
            "# {}\n\n{}\n\n## Features\n\n{}\n\n## Usage\n\n```rust\nuse {}::*;\n```",
            architecture.concept.spec.name,
            architecture.concept.spec.description,
            architecture.concept.spec.features.join("\n- "),
            architecture.concept.spec.name
        );
        docs.insert("README.md".to_string(), readme_content);
        
        // Generate CHANGELOG.md
        let changelog_content = format!(
            "# Changelog\n\n## [{}] - {}\n\n### Added\n- Initial release\n",
            architecture.concept.spec.version,
            chrono::Utc::now().format("%Y-%m-%d")
        );
        docs.insert("CHANGELOG.md".to_string(), changelog_content);
        
        Ok(docs)
    }
}

#[async_trait]
impl Stage for CrateGenerationStage {
    type Input = Architecture;
    type Output = CrateContext;

    async fn execute(&self, input: &Self::Input) -> Result<Self::Output> {
        info!("Starting code generation for crate: {}", input.concept.spec.name);
        
        let file_structure = self.generate_code_files(input).await?;
        let documentation = self.generate_documentation(input).await?;
        
        // Merge documentation into file structure
        let mut all_files = file_structure;
        all_files.extend(documentation);
        
        let dependencies = input.concept.spec.dependencies.clone();
        
        let mut metadata = HashMap::new();
        metadata.insert("generated_at".to_string(), chrono::Utc::now().to_rfc3339());
        metadata.insert("generator_version".to_string(), env!("CARGO_PKG_VERSION").to_string());
        
        let context = CrateContext {
            spec: input.concept.spec.clone(),
            file_structure: all_files,
            dependencies,
            metadata,
        };
        
        debug!("Code generation complete with {} files", context.file_structure.len());
        Ok(context)
    }

    fn name(&self) -> &str {
        "generation"
    }
}