pforge-codegen 0.2.0

pub mod generator;

pub use generator::{generate_handler_registration, generate_param_struct, CodegenError};

use pforge_config::ForgeConfig;
use std::path::Path;

pub type Result<T> = std::result::Result<T, CodegenError>;

/// Generate all code from a ForgeConfig
pub fn generate_all(config: &ForgeConfig) -> Result<String> {
    let mut output = String::new();

    // Generate imports
    output.push_str("// Auto-generated by pforge\n");
    output.push_str("// DO NOT EDIT\n\n");
    output.push_str("use pforge_runtime::*;\n");
    output.push_str("use serde::{Deserialize, Serialize};\n");
    output.push_str("use schemars::JsonSchema;\n\n");

    // Generate parameter structs for each tool
    for tool in &config.tools {
        if let pforge_config::ToolDef::Native { name, params, .. } = tool {
            output.push_str(&generate_param_struct(name, params)?);
            output.push_str("\n\n");
        }
    }

    // Generate handler registration function
    output.push_str(&generate_handler_registration(config)?);

    Ok(output)
}

/// Write generated code to a file
pub fn write_generated_code(config: &ForgeConfig, output_path: &Path) -> Result<()> {
    let code = generate_all(config)?;
    std::fs::write(output_path, code)
        .map_err(|e| CodegenError::IoError(output_path.to_path_buf(), e))?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use pforge_config::*;
    use rustc_hash::FxHashMap;

    fn create_test_config() -> ForgeConfig {
        ForgeConfig {
            forge: ForgeMetadata {
                name: "test_server".to_string(),
                version: "1.0.0".to_string(),
                transport: TransportType::Stdio,
                optimization: OptimizationLevel::Debug,
            },
            tools: vec![ToolDef::Native {
                name: "test_tool".to_string(),
                description: "Test tool".to_string(),
                handler: HandlerRef {
                    path: "handlers::test_handler".to_string(),
                    inline: None,
                },
                params: ParamSchema {
                    fields: {
                        let mut map = FxHashMap::default();
                        map.insert("input".to_string(), ParamType::Simple(SimpleType::String));
                        map
                    },
                },
                timeout_ms: None,
            }],
            resources: vec![],
            prompts: vec![],
            state: None,
        }
    }

    #[test]
    fn test_generate_all() {
        let config = create_test_config();
        let result = generate_all(&config);

        assert!(result.is_ok());
        let code = result.unwrap();

        // Check for generated header
        assert!(code.contains("// Auto-generated by pforge"));
        assert!(code.contains("// DO NOT EDIT"));

        // Check for imports
        assert!(code.contains("use pforge_runtime::*"));
        assert!(code.contains("use serde::{Deserialize, Serialize}"));
        assert!(code.contains("use schemars::JsonSchema"));

        // Check for param struct
        assert!(code.contains("pub struct TestToolParams"));

        // Check for registration function
        assert!(code.contains("pub fn register_handlers"));
    }

    #[test]
    fn test_generate_all_empty_tools() {
        let config = ForgeConfig {
            forge: ForgeMetadata {
                name: "empty".to_string(),
                version: "1.0.0".to_string(),
                transport: TransportType::Stdio,
                optimization: OptimizationLevel::Debug,
            },
            tools: vec![],
            resources: vec![],
            prompts: vec![],
            state: None,
        };

        let result = generate_all(&config);
        assert!(result.is_ok());
        let code = result.unwrap();
        assert!(code.contains("pub fn register_handlers"));
    }

    #[test]
    fn test_write_generated_code() {
        let config = create_test_config();
        let temp_dir = std::env::temp_dir();
        let output_path = temp_dir.join("test_generated.rs");

        let result = write_generated_code(&config, &output_path);
        assert!(result.is_ok());

        // Verify file was created
        assert!(output_path.exists());

        // Verify content
        let content = std::fs::read_to_string(&output_path).unwrap();
        assert!(content.contains("pub struct TestToolParams"));

        // Cleanup
        std::fs::remove_file(&output_path).ok();
    }

    #[test]
    fn test_write_generated_code_invalid_path() {
        let config = create_test_config();
        let invalid_path = Path::new("/nonexistent/directory/test.rs");

        let result = write_generated_code(&config, invalid_path);
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), CodegenError::IoError(_, _)));
    }
}