#!/usr/bin/env python3
"""
Initialize SIGMA rule templates for enhanced rule generation
"""

import json
import yaml
import os
from pathlib import Path
from datetime import datetime
from main import SessionLocal, RuleTemplate, Base, engine

# Create tables if they don't exist
Base.metadata.create_all(bind=engine)

def load_templates_from_files():
    """Load SIGMA rule templates from YAML files in the templates directory"""
    templates = []
    templates_dir = Path(__file__).parent / "templates"
    
    if not templates_dir.exists():
        print(f"Templates directory not found: {templates_dir}")
        return templates
    
    # Load all YAML files from templates directory
    for template_file in templates_dir.glob("*.yaml"):
        try:
            with open(template_file, 'r', encoding='utf-8') as f:
                template_data = yaml.safe_load(f)
                templates.append(template_data)
                print(f"Loaded template: {template_data['template_name']}")
        except Exception as e:
            print(f"Error loading template from {template_file}: {e}")
    
    return templates

def initialize_templates():
    """Initialize rule templates in the database"""
    db = SessionLocal()
    
    try:
        # Load templates from YAML files
        templates = load_templates_from_files()
        
        if not templates:
            print("No templates found to initialize")
            return
        
        # Clear existing templates
        db.query(RuleTemplate).delete()
        
        # Add new templates
        for template_data in templates:
            template = RuleTemplate(
                template_name=template_data["template_name"],
                template_content=template_data["template_content"],
                applicable_product_patterns=template_data["applicable_product_patterns"],
                description=template_data["description"]
            )
            db.add(template)
        
        db.commit()
        print(f"Successfully initialized {len(templates)} rule templates")
        
    except Exception as e:
        db.rollback()
        print(f"Error initializing templates: {e}")
        raise
    finally:
        db.close()

def list_available_templates():
    """List all available template files and their details"""
    templates = load_templates_from_files()
    
    if not templates:
        print("No templates found")
        return
    
    print(f"\nFound {len(templates)} available templates:")
    print("=" * 60)
    
    for i, template in enumerate(templates, 1):
        print(f"{i}. {template['template_name']}")
        print(f"   Description: {template['description']}")
        print(f"   Applicable patterns: {', '.join(template['applicable_product_patterns'])}")
        print()

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1 and sys.argv[1] == "list":
        list_available_templates()
    else:
        initialize_templates()