opengrep 1.1.0

//! AST context extraction and formatting
//!
//! This module handles extracting contextual information from AST nodes
//! and formatting it for display.

use super::AstNode;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;

/// AST context for a code location
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AstContext {
    /// Context nodes from outer to inner
    pub nodes: Vec<ContextNode>,
    /// All lines that should be shown
    pub visible_lines: HashSet<usize>,
    /// Header lines for each scope
    pub headers: Vec<HeaderInfo>,
    /// Context summary
    pub summary: String,
}

/// A node in the context hierarchy
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextNode {
    /// The AST node
    pub node: AstNode,
    /// Formatted display text
    pub display_text: String,
    /// Indentation level
    pub indent: usize,
    /// Whether this is a direct match container
    pub is_match_container: bool,
    /// Whether this node is collapsed
    pub is_collapsed: bool,
}

/// Header information for a scope
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HeaderInfo {
    /// Start line of the header
    pub start_line: usize,
    /// End line of the header
    pub end_line: usize,
    /// The node this header represents
    pub node_index: usize,
    /// Header text
    pub text: String,
}

impl AstContext {
    /// Create a new empty context
    pub fn new() -> Self {
        Self {
            nodes: Vec::new(),
            visible_lines: HashSet::new(),
            headers: Vec::new(),
            summary: String::new(),
        }
    }
    
    /// Add a node to the context
    pub fn add_node(&mut self, node: &AstNode, all_nodes: &[AstNode]) {
        // Walk up the parent chain to get full context
        let mut chain = vec![node.clone()];
        let mut current = node.parent;
        
        while let Some(parent_idx) = current {
            if let Some(parent) = all_nodes.get(parent_idx) {
                // Only add significant parent nodes
                if self.is_significant_node(parent) {
                    chain.push(parent.clone());
                }
                current = parent.parent;
            } else {
                break;
            }
        }
        
        // Reverse to get outer-to-inner order
        chain.reverse();
        
        // Convert to context nodes
        let chain_len = chain.len();
        for (indent, node) in chain.into_iter().enumerate() {
            let is_match_container = indent == chain_len - 1;
            
            let context_node = ContextNode {
                display_text: self.format_node(&node),
                indent,
                is_match_container,
                is_collapsed: false,
                node,
            };
            
            // Avoid duplicates
            if !self.nodes.iter().any(|cn| {
                cn.node.range == context_node.node.range && 
                cn.node.kind == context_node.node.kind
            }) {
                self.nodes.push(context_node);
            }
        }
    }
    
    /// Build the final context with visible lines
    pub fn build(&mut self, source: &str) {
        let lines: Vec<&str> = source.lines().collect();
        
        // Sort nodes by indentation for proper hierarchy
        self.nodes.sort_by_key(|n| n.indent);
        
        // Add header lines for each scope
        for (idx, context_node) in self.nodes.iter().enumerate() {
            let node = &context_node.node;
            
            // Determine how many lines to show for this node
            let header_lines = self.calculate_header_lines(node, &lines);
            
            for line in header_lines {
                self.visible_lines.insert(line);
            }
            
            // Create header info
            let header_text = self.extract_header_text(node, &lines);
            self.headers.push(HeaderInfo {
                start_line: node.start.line,
                end_line: node.end.line.min(node.start.line + 3),
                node_index: idx,
                text: header_text,
            });
        }
        
        // Generate context summary
        self.summary = self.generate_summary();
    }
    
    /// Check if a node is significant for context
    fn is_significant_node(&self, node: &AstNode) -> bool {
        // Include nodes that provide meaningful context
        matches!(
            node.kind.as_str(),
            "function_item" | "function_definition" | "function_declaration" |
            "method_definition" | "method_declaration" |
            "class_definition" | "class_declaration" |
            "struct_item" | "struct_specifier" |
            "enum_item" | "enum_declaration" |
            "trait_item" | "trait_declaration" |
            "impl_item" | "implementation" |
            "mod_item" | "module" |
            "interface_declaration" | "namespace_declaration"
        ) || node.metadata.is_scope || node.metadata.is_definition
    }
    
    /// Format a node for display
    fn format_node(&self, node: &AstNode) -> String {
        let name = node.name.as_deref().unwrap_or("<anonymous>");
        
        // Add visibility if available
        let visibility = node.metadata.visibility
            .as_ref()
            .map(|v| format!("{} ", v))
            .unwrap_or_default();
        
        // Format based on node type
        match node.kind.as_str() {
            "function_item" | "function_definition" | "function_declaration" => {
                format!("{}fn {}", visibility, name)
            },
            "method_definition" | "method_declaration" => {
                format!("{}method {}", visibility, name)
            },
            "class_definition" | "class_declaration" => {
                format!("{}class {}", visibility, name)
            },
            "struct_item" | "struct_specifier" => {
                format!("{}struct {}", visibility, name)
            },
            "enum_item" | "enum_declaration" => {
                format!("{}enum {}", visibility, name)
            },
            "trait_item" | "trait_declaration" => {
                format!("{}trait {}", visibility, name)
            },
            "impl_item" => {
                format!("{}impl {}", visibility, name)
            },
            "interface_declaration" => {
                format!("{}interface {}", visibility, name)
            },
            _ => format!("{} {}", node.kind, name),
        }
    }
    
    /// Calculate which lines to show for a node
    fn calculate_header_lines(&self, node: &AstNode, _lines: &[&str]) -> Vec<usize> {
        let start_line = node.start.line;
        let end_line = node.end.line;
        
        if start_line == end_line {
            // Single line node
            vec![start_line]
        } else if end_line - start_line <= 5 {
            // Small node, show all lines
            (start_line..=end_line).collect()
        } else {
            // Large node, show first few lines and closing line
            let mut lines_to_show = vec![start_line];
            
            // Add a few lines after the start
            for i in 1..=3 {
                if start_line + i < end_line {
                    lines_to_show.push(start_line + i);
                }
            }
            
            // Add the closing line if it's different
            if end_line > start_line + 3 {
                lines_to_show.push(end_line);
            }
            
            lines_to_show
        }
    }
    
    /// Extract meaningful header text from a node
    fn extract_header_text(&self, node: &AstNode, lines: &[&str]) -> String {
        if node.start.line < lines.len() {
            let line = lines[node.start.line];
            
            // Try to extract a meaningful signature
            if line.len() > 100 {
                // Truncate long lines
                format!("{}...", &line[..97])
            } else {
                line.to_string()
            }
        } else {
            self.format_node(node)
        }
    }
    
    /// Generate a summary of the context
    fn generate_summary(&self) -> String {
        if self.nodes.is_empty() {
            return "No context available".to_string();
        }
        
        let mut parts = Vec::new();
        for node in &self.nodes {
            if node.indent < 3 {  // Only include top-level context
                parts.push(self.format_node(&node.node));
            }
        }
        
        if parts.is_empty() {
            "Global scope".to_string()
        } else {
            parts.join(" > ")
        }
    }
    
    /// Merge with another context
    pub fn merge(&mut self, other: &AstContext) {
        // Merge nodes, avoiding duplicates
        let mut existing_keys = HashSet::new();
        for node in &self.nodes {
            existing_keys.insert((node.node.start, node.node.end, node.node.kind.clone()));
        }
        
        for context_node in &other.nodes {
            let key = (context_node.node.start, context_node.node.end, context_node.node.kind.clone());
            if !existing_keys.contains(&key) {
                self.nodes.push(context_node.clone());
                existing_keys.insert(key);
            }
        }
        
        // Merge visible lines
        self.visible_lines.extend(&other.visible_lines);
        
        // Merge headers
        self.headers.extend(other.headers.clone());
        
        // Update summary
        self.summary = self.generate_summary();
    }
    
    /// Get context path as a string
    pub fn get_path(&self) -> String {
        self.nodes
            .iter()
            .take(3)  // Only show first 3 levels
            .map(|n| {
                n.node.name.as_deref()
                    .unwrap_or(&n.node.kind)
            })
            .collect::<Vec<_>>()
            .join("::")
    }
    
    /// Get the deepest context node
    pub fn get_deepest_node(&self) -> Option<&ContextNode> {
        self.nodes.iter().max_by_key(|n| n.indent)
    }
    
    /// Get nodes at a specific indentation level
    pub fn get_nodes_at_level(&self, level: usize) -> Vec<&ContextNode> {
        self.nodes.iter()
            .filter(|n| n.indent == level)
            .collect()
    }
    
    /// Check if context contains a specific node type
    pub fn contains_node_type(&self, node_type: &str) -> bool {
        self.nodes.iter().any(|n| n.node.kind == node_type)
    }
}

impl Default for AstContext {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::{NodeMetadata, Position};
    use std::ops::Range;
    
    fn create_test_node(kind: &str, name: Option<&str>, start_line: usize, end_line: usize) -> AstNode {
        AstNode {
            kind: kind.to_string(),
            name: name.map(|s| s.to_string()),
            start: Position { line: start_line, column: 0, offset: 0 },
            end: Position { line: end_line, column: 0, offset: 0 },
            range: Range { start: 0, end: 100 },
            depth: 0,
            parent: None,
            children: Vec::new(),
            metadata: NodeMetadata::default(),
        }
    }
    
    #[test]
    fn test_ast_context_creation() {
        let mut context = AstContext::new();
        assert!(context.nodes.is_empty());
        assert!(context.visible_lines.is_empty());
        assert!(context.headers.is_empty());
    }
    
    #[test]
    fn test_format_node() {
        let context = AstContext::new();
        let node = create_test_node("function_item", Some("test_fn"), 0, 5);
        
        let formatted = context.format_node(&node);
        assert_eq!(formatted, "fn test_fn");
    }
    
    #[test]
    fn test_context_path() {
        let mut context = AstContext::new();
        
        let node1 = create_test_node("class_definition", Some("MyClass"), 0, 10);
        let node2 = create_test_node("function_definition", Some("my_method"), 2, 8);
        
        context.add_node(&node1, &[]);
        context.add_node(&node2, &[]);
        
        let path = context.get_path();
        assert!(path.contains("MyClass"));
    }
    
    #[test]
    fn test_significant_node() {
        let context = AstContext::new();
        
        let func_node = create_test_node("function_item", Some("test"), 0, 5);
        assert!(context.is_significant_node(&func_node));
        
        let expr_node = create_test_node("expression", None, 1, 1);
        assert!(!context.is_significant_node(&expr_node));
    }
}