herolib-code 0.3.13

//! AI-powered code analysis.
//!
//! This module provides AI-powered analysis of parsed Rust code structures.
//! Requires the `ai` feature to be enabled.

use crate::parser::{CodeBase, EnumInfo, MethodInfo, StructInfo};
use herolib_ai::{AiClient, AiResult, Model, PromptBuilderExt};

/// AI-powered code analyzer.
///
/// Uses AI to provide insights, documentation suggestions, and analysis
/// of parsed Rust code structures.
pub struct CodeAnalyzer {
    /// The AI client to use.
    client: AiClient,
    /// The model to use for analysis.
    model: Model,
}

impl Default for CodeAnalyzer {
    fn default() -> Self {
        Self::from_env()
    }
}

impl CodeAnalyzer {
    /// Creates a new CodeAnalyzer with the given AI client.
    pub fn new(client: AiClient) -> Self {
        Self {
            client,
            model: Model::default_coding(),
        }
    }

    /// Creates a new CodeAnalyzer from environment variables.
    pub fn from_env() -> Self {
        Self::new(AiClient::from_env())
    }

    /// Sets the model to use for analysis.
    pub fn with_model(mut self, model: Model) -> Self {
        self.model = model;
        self
    }

    /// Generates documentation for a struct.
    pub fn document_struct(&self, struct_info: &StructInfo) -> AiResult<String> {
        let struct_repr = self.struct_to_string(struct_info);

        self.client
            .prompt()
            .model(self.model)
            .system(
                "You are a Rust documentation expert. Generate clear, concise rustdoc \
                documentation for the given struct. Include:\n\
                - A brief description of the struct's purpose\n\
                - Documentation for each field\n\
                - Example usage if appropriate\n\
                Output only the documentation comments (with /// prefix), no code.",
            )
            .user(format!(
                "Generate documentation for this struct:\n\n{}",
                struct_repr
            ))
            .execute_content()
    }

    /// Generates documentation for an enum.
    pub fn document_enum(&self, enum_info: &EnumInfo) -> AiResult<String> {
        let enum_repr = self.enum_to_string(enum_info);

        self.client
            .prompt()
            .model(self.model)
            .system(
                "You are a Rust documentation expert. Generate clear, concise rustdoc \
                documentation for the given enum. Include:\n\
                - A brief description of the enum's purpose\n\
                - Documentation for each variant\n\
                - Example usage if appropriate\n\
                Output only the documentation comments (with /// prefix), no code.",
            )
            .user(format!(
                "Generate documentation for this enum:\n\n{}",
                enum_repr
            ))
            .execute_content()
    }

    /// Analyzes a codebase and provides a summary.
    pub fn summarize_codebase(&self, codebase: &CodeBase) -> AiResult<String> {
        let summary = self.codebase_summary(codebase);

        self.client
            .prompt()
            .model(self.model)
            .system(
                "You are a code analysis expert. Analyze the given Rust codebase summary \
                and provide:\n\
                - Overview of the codebase structure\n\
                - Main components and their purposes\n\
                - Notable patterns or architectural decisions\n\
                - Suggestions for improvements if any",
            )
            .user(format!("Analyze this Rust codebase:\n\n{}", summary))
            .execute_content()
    }

    /// Suggests improvements for a method.
    pub fn suggest_method_improvements(&self, method: &MethodInfo) -> AiResult<String> {
        let method_repr = self.method_to_string(method);

        self.client
            .prompt()
            .model(self.model)
            .system(
                "You are a Rust code review expert. Analyze the given method signature \
                and suggest improvements for:\n\
                - Naming conventions\n\
                - Parameter types (consider borrowing, generics)\n\
                - Return type (consider Result, Option)\n\
                - Documentation\n\
                Be concise and actionable.",
            )
            .user(format!("Review this method:\n\n{}", method_repr))
            .execute_content()
    }

    /// Generates a struct representation as a string.
    fn struct_to_string(&self, s: &StructInfo) -> String {
        let mut result = String::new();

        // Add existing doc comment if any
        if let Some(doc) = &s.doc_comment {
            result.push_str(&format!("// Existing doc: {}\n", doc));
        }

        // Add derives
        if !s.derives.is_empty() {
            result.push_str(&format!("#[derive({})]\n", s.derives.join(", ")));
        }

        // Add struct definition
        result.push_str(&format!("{} struct {} ", s.visibility, s.name));

        if !s.generics.is_empty() {
            result.push_str(&format!("<{}>", s.generics.join(", ")));
        }

        result.push_str("{\n");

        for field in &s.fields {
            if let Some(doc) = &field.doc_comment {
                result.push_str(&format!("    // {}\n", doc));
            }
            if let Some(name) = &field.name {
                result.push_str(&format!("    {}: {},\n", name, field.ty));
            }
        }

        result.push_str("}\n");

        // Add methods
        if !s.methods.is_empty() {
            result.push_str("\nimpl ");
            result.push_str(&s.name);
            result.push_str(" {\n");

            for method in &s.methods {
                result.push_str(&format!("    {}\n", self.method_signature(method)));
            }

            result.push_str("}\n");
        }

        result
    }

    /// Generates an enum representation as a string.
    fn enum_to_string(&self, e: &EnumInfo) -> String {
        let mut result = String::new();

        if let Some(doc) = &e.doc_comment {
            result.push_str(&format!("// Existing doc: {}\n", doc));
        }

        if !e.derives.is_empty() {
            result.push_str(&format!("#[derive({})]\n", e.derives.join(", ")));
        }

        result.push_str(&format!("{} enum {} ", e.visibility, e.name));

        if !e.generics.is_empty() {
            result.push_str(&format!("<{}>", e.generics.join(", ")));
        }

        result.push_str("{\n");

        for variant in &e.variants {
            if let Some(doc) = &variant.doc_comment {
                result.push_str(&format!("    // {}\n", doc));
            }
            result.push_str(&format!("    {},\n", variant.name));
        }

        result.push_str("}\n");

        result
    }

    /// Generates a method representation as a string.
    fn method_to_string(&self, m: &MethodInfo) -> String {
        let mut result = String::new();

        if let Some(doc) = &m.doc_comment {
            result.push_str(&format!("// {}\n", doc));
        }

        result.push_str(&self.method_signature(m));

        result
    }

    /// Generates a method signature.
    fn method_signature(&self, m: &MethodInfo) -> String {
        let mut sig = format!("{} ", m.visibility);

        if m.is_async {
            sig.push_str("async ");
        }
        if m.is_const {
            sig.push_str("const ");
        }
        if m.is_unsafe {
            sig.push_str("unsafe ");
        }

        sig.push_str(&format!("fn {}", m.name));

        if !m.generics.is_empty() {
            sig.push_str(&format!("<{}>", m.generics.join(", ")));
        }

        sig.push('(');

        let mut params = Vec::new();

        if let Some(recv) = &m.receiver {
            params.push(format!("{:?}", recv).to_lowercase());
        }

        for param in &m.parameters {
            params.push(format!("{}: {}", param.name, param.ty));
        }

        sig.push_str(&params.join(", "));
        sig.push(')');

        if let Some(ret) = &m.return_type {
            sig.push_str(&format!(" -> {}", ret));
        }

        sig
    }

    /// Generates a codebase summary.
    fn codebase_summary(&self, codebase: &CodeBase) -> String {
        let mut summary = String::new();

        summary.push_str(&format!(
            "Codebase Statistics:\n\
            - {} structs\n\
            - {} enums\n\
            - {} files\n\n",
            codebase.structs.len(),
            codebase.enums.len(),
            codebase.files.len()
        ));

        summary.push_str("Structs:\n");
        for s in &codebase.structs {
            summary.push_str(&format!(
                "- {} ({} fields, {} methods)\n",
                s.name,
                s.fields.len(),
                s.methods.len()
            ));
        }

        summary.push_str("\nEnums:\n");
        for e in &codebase.enums {
            summary.push_str(&format!("- {} ({} variants)\n", e.name, e.variants.len()));
        }

        summary
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::{FieldInfo, Visibility};

    fn sample_struct() -> StructInfo {
        StructInfo {
            name: "User".to_string(),
            doc_comment: Some("A user in the system".to_string()),
            file_path: "src/models.rs".to_string(),
            line_number: 10,
            visibility: Visibility::Public,
            generics: vec![],
            derives: vec!["Debug".to_string(), "Clone".to_string()],
            attributes: vec![],
            fields: vec![
                FieldInfo {
                    name: Some("id".to_string()),
                    ty: "u64".to_string(),
                    doc_comment: Some("Unique identifier".to_string()),
                    visibility: Visibility::Public,
                    attributes: vec![],
                },
                FieldInfo {
                    name: Some("name".to_string()),
                    ty: "String".to_string(),
                    doc_comment: None,
                    visibility: Visibility::Public,
                    attributes: vec![],
                },
            ],
            methods: vec![],
        }
    }

    #[test]
    fn test_struct_to_string() {
        let analyzer = CodeAnalyzer::from_env();
        let s = sample_struct();
        let repr = analyzer.struct_to_string(&s);

        assert!(repr.contains("struct User"));
        assert!(repr.contains("id: u64"));
        assert!(repr.contains("name: String"));
        assert!(repr.contains("#[derive(Debug, Clone)]"));
    }
}