__init__.py

"""
GraphCodeBERT with AST+DFG for JavaScript Vulnerability Detection
This module provides AST and DFG extraction using Tree-sitter and integrates with GraphCodeBERT
Includes visualization capabilities for AST and DFG
"""

import torch
import numpy as np
from tree_sitter import Language, Parser
import tree_sitter_javascript
from typing import List, Dict
from collections import defaultdict
import matplotlib.pyplot as plt
import networkx as nx
from matplotlib.patches import FancyBboxPatch
import os

# Initialize Tree-sitter parser for JavaScript
JS_LANGUAGE = Language(tree_sitter_javascript.language())


class ASTExtractor:
    """Extract Abstract Syntax Tree features from JavaScript code"""
    
    def __init__(self):
        self.parser = None
    
    def _get_parser(self):
        """Lazy initialization of parser for pickle compatibility"""
        if self.parser is None:
            self.parser = Parser(JS_LANGUAGE)
        return self.parser
    
    def parse_code(self, code: str):
        """Parse JavaScript code and return AST"""
        try:
            parser = self._get_parser()
            tree = parser.parse(bytes(code, "utf8"))
            return tree
        except Exception as e:
            return None
    
    def extract_nodes(self, tree) -> List[Dict]:
        """Extract nodes from AST with their types and relationships"""
        if tree is None:
            return []
        
        nodes = []
        node_id = 0
        
        def traverse(node, parent_id=None):
            nonlocal node_id
            current_id = node_id
            
            node_info = {
                'id': current_id,
                'type': node.type,
                'start_point': node.start_point,
                'end_point': node.end_point,
                'parent_id': parent_id,
                'text': node.text.decode('utf8') if node.text else ''
            }
            nodes.append(node_info)
            node_id += 1
            
            # Traverse children
            for child in node.children:
                traverse(child, current_id)
        
        traverse(tree.root_node)
        return nodes
    
    def get_ast_sequence(self, code: str) -> List[str]:
        """Get sequence of node types from AST (for GraphCodeBERT input)"""
        tree = self.parse_code(code)
        nodes = self.extract_nodes(tree)
        return [node['type'] for node in nodes]


class DFGExtractor:
    """Extract Data Flow Graph from JavaScript AST"""
    
    def __init__(self):
        self.ast_extractor = ASTExtractor()
        
    def extract_variables(self, tree) -> Dict[str, List]:
        """Extract variable definitions and uses"""
        if tree is None:
            return {'definitions': [], 'uses': []}
        
        definitions = []
        uses = []
        
        def traverse(node, scope_vars=None):
            if scope_vars is None:
                scope_vars = set()
            
            # Variable declarations
            if node.type in ['variable_declarator', 'formal_parameters', 'assignment_expression']:
                for child in node.children:
                    if child.type == 'identifier':
                        var_name = child.text.decode('utf8')
                        definitions.append({
                            'name': var_name,
                            'node_id': id(node),
                            'position': child.start_point
                        })
                        scope_vars.add(var_name)
            
            # Variable uses
            elif node.type == 'identifier':
                var_name = node.text.decode('utf8')
                # Check if it's a use (not a definition context)
                if var_name in scope_vars or node.parent.type in [
                    'member_expression', 'call_expression', 'binary_expression',
                    'return_statement', 'if_statement', 'arguments'
                ]:
                    uses.append({
                        'name': var_name,
                        'node_id': id(node),
                        'position': node.start_point
                    })
            
            # Recursively process children
            for child in node.children:
                traverse(child, scope_vars.copy())
        
        traverse(tree.root_node)
        return {'definitions': definitions, 'uses': uses}
    
    def build_dfg(self, code: str) -> Dict:
        """Build Data Flow Graph from code"""
        tree = self.ast_extractor.parse_code(code)
        if tree is None:
            return {'nodes': [], 'edges': []}
        
        var_info = self.extract_variables(tree)
        definitions = var_info['definitions']
        uses = var_info['uses']
        
        # Create mapping of variable names to their definitions
        var_defs = defaultdict(list)
        for def_info in definitions:
            var_defs[def_info['name']].append(def_info)
        
        # Build edges: from definition to use
        edges = []
        for use in uses:
            var_name = use['name']
            if var_name in var_defs:
                # Connect to the most recent definition before this use
                for def_info in var_defs[var_name]:
                    if def_info['position'] <= use['position']:
                        edges.append({
                            'from': def_info['node_id'],
                            'to': use['node_id'],
                            'variable': var_name,
                            'type': 'data_flow'
                        })
        
        # Create unique nodes
        all_nodes = definitions + uses
        unique_nodes = {n['node_id']: n for n in all_nodes}.values()
        
        return {
            'nodes': list(unique_nodes),
            'edges': edges
        }


class GraphBuilder:
    """Build graph representation combining AST and DFG for GraphCodeBERT"""
    
    def __init__(self, max_nodes=512):
        self.ast_extractor = ASTExtractor()
        self.dfg_extractor = DFGExtractor()
        self.max_nodes = max_nodes
    
    def build_graph(self, code: str) -> Dict:
        """Build combined AST+DFG graph"""
        # Extract AST
        tree = self.ast_extractor.parse_code(code)
        ast_nodes = self.ast_extractor.extract_nodes(tree)
        
        # Extract DFG
        dfg = self.dfg_extractor.build_dfg(code)
        
        # Combine into graph structure
        graph = {
            'ast_nodes': ast_nodes[:self.max_nodes],
            'dfg_nodes': dfg['nodes'][:self.max_nodes],
            'dfg_edges': dfg['edges'],
            'node_types': [n['type'] for n in ast_nodes[:self.max_nodes]]
        }
        
        return graph
    
    def create_adjacency_matrix(self, graph: Dict) -> np.ndarray:
        """Create adjacency matrix from graph for GNN processing"""
        num_nodes = len(graph['ast_nodes'])
        adj_matrix = np.zeros((num_nodes, num_nodes), dtype=np.float32)
        
        # Add AST edges (parent-child relationships) - BIDIRECTIONAL
        for node in graph['ast_nodes']:
            if node['parent_id'] is not None and node['parent_id'] < num_nodes:
                adj_matrix[node['parent_id']][node['id']] = 1  # Parent to child
                adj_matrix[node['id']][node['parent_id']] = 1  # Child to parent (bidirectional)
        
        # Add DFG edges (data flow) - DIRECTIONAL
        node_id_map = {n['node_id']: idx for idx, n in enumerate(graph['dfg_nodes'])}
        for edge in graph['dfg_edges']:
            if edge['from'] in node_id_map and edge['to'] in node_id_map:
                from_idx = node_id_map[edge['from']]
                to_idx = node_id_map[edge['to']]
                if from_idx < num_nodes and to_idx < num_nodes:
                    adj_matrix[from_idx][to_idx] = 2  # Mark DFG edges with value 2
        
        return adj_matrix
    
    def visualize_ast(self, code: str, save_path: str = None, max_depth: int = 4):
        """Visualize Abstract Syntax Tree"""
        tree = self.ast_extractor.parse_code(code)
        if tree is None:
            print("Failed to parse code")
            return
        
        G = nx.DiGraph()
        pos = {}
        labels = {}
        colors = []
        
        def add_nodes(node, parent_id=None, depth=0, x=0, width=1):
            if depth > max_depth:
                return x
            
            node_id = id(node)
            node_type = node.type
            
            # Simplify label for readability
            label = node_type[:15] + '...' if len(node_type) > 15 else node_type
            
            G.add_node(node_id)
            labels[node_id] = label
            
            # Assign color based on node type
            if node_type in ['function_declaration', 'arrow_function', 'function']:
                colors.append('#FFB6C1')  # Light red for functions
            elif node_type in ['if_statement', 'for_statement', 'while_statement']:
                colors.append('#87CEEB')  # Sky blue for control flow
            elif node_type in ['variable_declarator', 'identifier']:
                colors.append('#90EE90')  # Light green for variables
            else:
                colors.append('#FFE4B5')  # Moccasin for others
            
            if parent_id is not None:
                G.add_edge(parent_id, node_id)
            
            # Calculate position
            pos[node_id] = (x, -depth)
            
            # Process children
            num_children = len(node.children)
            if num_children > 0:
                child_width = width / num_children
                child_x = x - width/2 + child_width/2
                
                for child in node.children:
                    child_x = add_nodes(child, node_id, depth + 1, child_x, child_width)
                    child_x += child_width
            
            return x
        
        add_nodes(tree.root_node)
        
        # Create figure
        plt.figure(figsize=(16, 10))
        nx.draw(G, pos, labels=labels, node_color=colors, 
                node_size=2000, font_size=8, font_weight='bold',
                arrows=True, edge_color='gray', arrowsize=20,
                node_shape='s', alpha=0.9)
        
        plt.title(f'Abstract Syntax Tree (AST)\nDepth: {max_depth}', 
                  fontsize=14, fontweight='bold', pad=20)
        plt.axis('off')
        plt.tight_layout()
        
        if save_path:
            os.makedirs(os.path.dirname(save_path) if os.path.dirname(save_path) else '.', exist_ok=True)
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"✓ AST visualization saved to {save_path}")
        else:
            plt.show()
        
        plt.close()
        
        return G
    
    def visualize_dfg(self, code: str, save_path: str = None):
        """Visualize Data Flow Graph"""
        dfg = self.dfg_extractor.build_dfg(code)
        
        if not dfg['edges']:
            print("No data flow edges found")
            return
        
        G = nx.DiGraph()
        
        # Create mapping for node display
        node_labels = {}
        for node in dfg['nodes']:
            node_id = node['node_id']
            var_name = node['name']
            G.add_node(node_id)
            node_labels[node_id] = var_name
        
        # Add edges
        edge_labels = {}
        for edge in dfg['edges']:
            from_id = edge['from']
            to_id = edge['to']
            var = edge['variable']
            
            if from_id in G and to_id in G:
                G.add_edge(from_id, to_id)
                edge_labels[(from_id, to_id)] = var
        
        if len(G.nodes()) == 0:
            print("No nodes in DFG")
            return
        
        # Create layout
        try:
            pos = nx.spring_layout(G, k=2, iterations=50, seed=42)
        except:
            pos = nx.shell_layout(G)
        
        # Create figure
        plt.figure(figsize=(14, 10))
        
        # Draw nodes
        nx.draw_networkx_nodes(G, pos, node_color='#98D8C8', 
                               node_size=3000, alpha=0.9,
                               node_shape='o')
        
        # Draw edges with arrows
        nx.draw_networkx_edges(G, pos, edge_color='#FF6B6B', 
                               arrows=True, arrowsize=20,
                               arrowstyle='->', width=2,
                               connectionstyle='arc3,rad=0.1')
        
        # Draw labels
        nx.draw_networkx_labels(G, pos, node_labels, font_size=10, 
                                font_weight='bold', font_color='black')
        
        # Draw edge labels
        nx.draw_networkx_edge_labels(G, pos, edge_labels, 
                                      font_size=8, font_color='darkred')
        
        plt.title('Data Flow Graph (DFG)\nShowing variable definitions and uses', 
                  fontsize=14, fontweight='bold', pad=20)
        plt.axis('off')
        plt.tight_layout()
        
        if save_path:
            os.makedirs(os.path.dirname(save_path) if os.path.dirname(save_path) else '.', exist_ok=True)
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"✓ DFG visualization saved to {save_path}")
        else:
            plt.show()
        
        plt.close()
        
        return G
    
    def visualize_combined(self, code: str, save_path: str = None, max_ast_depth: int = 3):
        """Visualize both AST and DFG side by side"""
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 10))
        
        # AST visualization
        tree = self.ast_extractor.parse_code(code)
        if tree is not None:
            G_ast = nx.DiGraph()
            pos_ast = {}
            labels_ast = {}
            colors_ast = []
            
            def add_nodes(node, parent_id=None, depth=0, x=0, width=1):
                if depth > max_ast_depth:
                    return x
                
                node_id = id(node)
                node_type = node.type
                label = node_type[:12] + '...' if len(node_type) > 12 else node_type
                
                G_ast.add_node(node_id)
                labels_ast[node_id] = label
                
                if node_type in ['function_declaration', 'arrow_function']:
                    colors_ast.append('#FFB6C1')
                elif node_type in ['if_statement', 'for_statement']:
                    colors_ast.append('#87CEEB')
                elif node_type in ['variable_declarator', 'identifier']:
                    colors_ast.append('#90EE90')
                else:
                    colors_ast.append('#FFE4B5')
                
                if parent_id is not None:
                    G_ast.add_edge(parent_id, node_id)
                
                pos_ast[node_id] = (x, -depth)
                
                num_children = len(node.children)
                if num_children > 0:
                    child_width = width / num_children
                    child_x = x - width/2 + child_width/2
                    for child in node.children:
                        child_x = add_nodes(child, node_id, depth + 1, child_x, child_width)
                        child_x += child_width
                
                return x
            
            add_nodes(tree.root_node)
            
            plt.sca(ax1)
            nx.draw(G_ast, pos_ast, labels=labels_ast, node_color=colors_ast,
                   node_size=1500, font_size=7, arrows=True, ax=ax1,
                   edge_color='gray', arrowsize=15, node_shape='s', alpha=0.9)
            ax1.set_title('Abstract Syntax Tree (AST)', fontsize=12, fontweight='bold')
            ax1.axis('off')
        
        # DFG visualization
        dfg = self.dfg_extractor.build_dfg(code)
        if dfg['edges']:
            G_dfg = nx.DiGraph()
            node_labels_dfg = {}
            
            for node in dfg['nodes']:
                node_id = node['node_id']
                var_name = node['name']
                G_dfg.add_node(node_id)
                node_labels_dfg[node_id] = var_name
            
            edge_labels_dfg = {}
            for edge in dfg['edges']:
                from_id, to_id = edge['from'], edge['to']
                if from_id in G_dfg and to_id in G_dfg:
                    G_dfg.add_edge(from_id, to_id)
                    edge_labels_dfg[(from_id, to_id)] = edge['variable']
            
            if len(G_dfg.nodes()) > 0:
                try:
                    pos_dfg = nx.spring_layout(G_dfg, k=1.5, iterations=50, seed=42)
                except:
                    pos_dfg = nx.circular_layout(G_dfg)
                
                plt.sca(ax2)
                nx.draw_networkx_nodes(G_dfg, pos_dfg, node_color='#98D8C8',
                                      node_size=2000, alpha=0.9, ax=ax2)
                nx.draw_networkx_edges(G_dfg, pos_dfg, edge_color='#FF6B6B',
                                      arrows=True, arrowsize=15, width=2, ax=ax2,
                                      connectionstyle='arc3,rad=0.1')
                nx.draw_networkx_labels(G_dfg, pos_dfg, node_labels_dfg,
                                       font_size=9, font_weight='bold', ax=ax2)
                nx.draw_networkx_edge_labels(G_dfg, pos_dfg, edge_labels_dfg,
                                            font_size=7, font_color='darkred', ax=ax2)
                ax2.set_title('Data Flow Graph (DFG)', fontsize=12, fontweight='bold')
                ax2.axis('off')
        
        plt.suptitle('AST + DFG Visualization', fontsize=16, fontweight='bold', y=0.98)
        plt.tight_layout()
        
        if save_path:
            os.makedirs(os.path.dirname(save_path) if os.path.dirname(save_path) else '.', exist_ok=True)
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"✓ Combined AST+DFG visualization saved to {save_path}")
        else:
            plt.show()
        
        plt.close()


class GraphCodeBERTProcessor:
    """Process code for GraphCodeBERT model with AST+DFG features"""
    
    def __init__(self, model_name='microsoft/graphcodebert-base', max_length=512):
        from transformers import RobertaTokenizer
        
        self.tokenizer = RobertaTokenizer.from_pretrained(model_name)
        self.max_length = max_length
        self.graph_builder = GraphBuilder(max_nodes=max_length)
    
    def process_code(self, code: str) -> Dict[str, torch.Tensor]:
        """Process code into GraphCodeBERT input format with AST+DFG"""
        # Tokenize code
        tokens = self.tokenizer.encode_plus(
            code,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        
        # Extract graph features
        graph = self.graph_builder.build_graph(code)
        adj_matrix = self.graph_builder.create_adjacency_matrix(graph)
        
        # Pad adjacency matrix to max_length
        padded_adj = np.zeros((self.max_length, self.max_length), dtype=np.float32)
        actual_size = min(adj_matrix.shape[0], self.max_length)
        padded_adj[:actual_size, :actual_size] = adj_matrix[:actual_size, :actual_size]
        
        # Create position embeddings for graph nodes
        node_positions = np.arange(self.max_length)
        
        return {
            'input_ids': tokens['input_ids'],
            'attention_mask': tokens['attention_mask'],
            'adjacency_matrix': torch.tensor(padded_adj, dtype=torch.float32),
            'position_ids': torch.tensor(node_positions, dtype=torch.long),
            'graph': graph  # Keep original graph for reference
        }
    
    def process_batch(self, codes: List[str]) -> Dict[str, torch.Tensor]:
        """Process batch of code samples"""
        batch_inputs = {
            'input_ids': [],
            'attention_mask': [],
            'adjacency_matrix': [],
            'position_ids': []
        }
        
        for code in codes:
            processed = self.process_code(code)
            batch_inputs['input_ids'].append(processed['input_ids'])
            batch_inputs['attention_mask'].append(processed['attention_mask'])
            batch_inputs['adjacency_matrix'].append(processed['adjacency_matrix'])
            batch_inputs['position_ids'].append(processed['position_ids'])
        
        # Stack tensors
        return {
            'input_ids': torch.cat(batch_inputs['input_ids'], dim=0),
            'attention_mask': torch.cat(batch_inputs['attention_mask'], dim=0),
            'adjacency_matrix': torch.stack(batch_inputs['adjacency_matrix'], dim=0),
            'position_ids': torch.stack(batch_inputs['position_ids'], dim=0)
        }


# Export main classes
__all__ = [
    'ASTExtractor',
    'DFGExtractor', 
    'GraphBuilder',
    'GraphCodeBERTProcessor',
    'JS_LANGUAGE'
]


# Convenience function
def process_javascript_code(code: str, model_name='microsoft/graphcodebert-base', max_length=512):
    """
    Convenience function to process JavaScript code with AST+DFG extraction
    
    Args:
        code: JavaScript source code string
        model_name: GraphCodeBERT model name
        max_length: Maximum sequence length
    
    Returns:
        Dictionary with processed inputs ready for GraphCodeBERT
    """
    processor = GraphCodeBERTProcessor(model_name=model_name, max_length=max_length)
    return processor.process_code(code)