import json
import re
from datasets import Dataset, DatasetDict
import ast

def check_syntax(code):
    """Check if Python code has syntax errors"""
    try:
        ast.parse(code)
        return True, None
    except SyntaxError as e:
        return False, str(e)

def create_enhanced_dataset():
    """Create better dataset with syntax errors marked"""
    
    samples = []
    
    # ===== VULNERABLE CODE (with proper syntax) =====
    vulnerable_examples = [
        # SQL Injection (proper syntax)
        {
            "code": """def get_user(user_id):\n    query = f"SELECT * FROM users WHERE id = {user_id}"\n    return execute_query(query)""",
            "label": 1,
            "type": "sql_injection",
            "explanation": "SQL injection: User input directly in query string"
        },
        {
            "code": """cursor.execute(f"DELETE FROM logs WHERE date = '{user_date}'")""",
            "label": 1,
            "type": "sql_injection",
            "explanation": "SQL injection in DELETE statement"
        },
        
        # Hardcoded Secrets
        {
            "code": """API_KEY = "EX_SAMPLE_KEY_NOT_REAL_12345" """,
            "label": 1,
            "type": "hardcoded_secret",
            "explanation": "Hardcoded API key in source code"
        },
        {
            "code": """db_password = "SuperSecret123!" """,
            "label": 1,
            "type": "hardcoded_secret", 
            "explanation": "Hardcoded database password"
        },
        
        # XSS Vulnerabilities
        {
            "code": """def render_comment(comment):\n    return f"<div>{comment}</div>" """,
            "label": 1,
            "type": "xss",
            "explanation": "XSS: User input in HTML without sanitization"
        },
        
        # Insecure Deserialization
        {
            "code": """import pickle\ndata = pickle.loads(user_input)""",
            "label": 1,
            "type": "insecure_deserialization",
            "explanation": "Pickle deserialization of user input is dangerous"
        },
        
        # Command Injection
        {
            "code": """import os\nos.system(f"echo {user_input}")""",
            "label": 1,
            "type": "command_injection",
            "explanation": "Command injection: User input in system command"
        }
    ]
    
    # ===== SAFE CODE =====
    safe_examples = [
        # Safe SQL
        {
            "code": """def get_user(user_id):\n    query = "SELECT * FROM users WHERE id = %s"\n    return execute_query(query, (user_id,))""",
            "label": 0,
            "type": "safe_sql",
            "explanation": "Parameterized query prevents SQL injection"
        },
        
        # Safe Secrets
        {
            "code": """import os\nAPI_KEY = os.getenv("API_KEY")""",
            "label": 0,
            "type": "safe_secret",
            "explanation": "API key from environment variable"
        },
        
        # Safe XSS
        {
            "code": """import html\ndef render_comment(comment):\n    return f"<div>{html.escape(comment)}</div>" """,
            "label": 0,
            "type": "safe_xss",
            "explanation": "HTML escaped user input prevents XSS"
        },
        
        # Safe Deserialization
        {
            "code": """import json\ndata = json.loads(user_input)""",
            "label": 0,
            "type": "safe_deserialization",
            "explanation": "JSON deserialization is generally safe"
        },
        
        # Safe Command Execution
        {
            "code": """import subprocess\nsubprocess.run(["echo", "static_text"], check=True)""",
            "label": 0,
            "type": "safe_command",
            "explanation": "Command with static arguments"
        }
    ]
    
    # ===== SYNTAX ERRORS (NEW!) =====
    syntax_error_examples = [
        {
            "code": """def test()\n    print("hello")""",  # Missing colon
            "label": 1,  # Mark as vulnerable/problematic
            "type": "syntax_error",
            "explanation": "Syntax error: Missing colon after function definition"
        },
        {
            "code": """x = [1, 2, 3""",  # Missing closing bracket
            "label": 1,
            "type": "syntax_error",
            "explanation": "Syntax error: Missing closing bracket"
        },
        {
            "code": """print("hello'""",  # Quote mismatch
            "label": 1,
            "type": "syntax_error", 
            "explanation": "Syntax error: Mismatched quotes"
        }
    ]
    
    # Combine all examples
    all_examples = []
    
    # Add multiple copies for balance
    for _ in range(5):  # 5 copies of each
        all_examples.extend(vulnerable_examples)
        all_examples.extend(safe_examples)
    
    # Add syntax errors
    all_examples.extend(syntax_error_examples)
    
    # Check syntax of all examples
    print("🔍 Checking syntax of all examples...")
    for i, example in enumerate(all_examples):
        is_valid, error = check_syntax(example["code"])
        example["has_syntax_error"] = not is_valid
        if error:
            example["syntax_error"] = error
    
    # Save to files
    import os
    os.makedirs("enhanced_data", exist_ok=True)
    
    # Save as JSON
    with open("enhanced_data/dataset.json", "w") as f:
        json.dump(all_examples, f, indent=2)
    
    # Create Hugging Face Dataset
    dataset = Dataset.from_list(all_examples)
    
    # Split into train/test
    dataset = dataset.train_test_split(test_size=0.2, seed=42)
    dataset.save_to_disk("enhanced_data/hf_dataset")
    
    # Statistics
    print("\n📊 DATASET STATISTICS:")
    print(f"Total examples: {len(all_examples)}")
    print(f"Vulnerable examples: {sum(1 for e in all_examples if e['label'] == 1)}")
    print(f"Safe examples: {sum(1 for e in all_examples if e['label'] == 0)}")
    print(f"With syntax errors: {sum(1 for e in all_examples if e.get('has_syntax_error', False))}")
    
    print("\n✅ Enhanced dataset created!")
    print("📁 Saved to: enhanced_data/")
    
    return all_examples

if __name__ == "__main__":
    create_enhanced_dataset()