"""
Database Models for CVE-SIGMA Auto Generator CLI

Maintains database model definitions for migration and data processing purposes.
Used by CLI migration tools to export data from legacy web application database.
"""

import uuid
from datetime import datetime
from sqlalchemy import create_engine, Column, String, Text, DECIMAL, TIMESTAMP, Boolean, ARRAY, Integer, JSON
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.dialects.postgresql import UUID
import os

# Database setup
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://cve_user:cve_password@localhost:5432/cve_sigma_db")
engine = create_engine(DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()

# Database Models (for migration purposes)
class CVE(Base):
    __tablename__ = "cves"
    
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    cve_id = Column(String(20), unique=True, nullable=False)
    description = Column(Text)
    cvss_score = Column(DECIMAL(3, 1))
    severity = Column(String(20))
    published_date = Column(TIMESTAMP)
    modified_date = Column(TIMESTAMP)
    affected_products = Column(ARRAY(String))
    reference_urls = Column(ARRAY(String))
    # Bulk processing fields
    data_source = Column(String(20), default='nvd_api')  # 'nvd_api', 'nvd_bulk', 'manual'
    nvd_json_version = Column(String(10), default='2.0')
    bulk_processed = Column(Boolean, default=False)
    # nomi-sec PoC fields
    poc_count = Column(Integer, default=0)
    poc_data = Column(JSON)  # Store nomi-sec PoC metadata
    # Reference data fields
    reference_data = Column(JSON)  # Store extracted reference content and analysis
    reference_sync_status = Column(String(20), default='pending')  # 'pending', 'processing', 'completed', 'failed'
    reference_last_synced = Column(TIMESTAMP)
    created_at = Column(TIMESTAMP, default=datetime.utcnow)
    updated_at = Column(TIMESTAMP, default=datetime.utcnow)

class SigmaRule(Base):
    __tablename__ = "sigma_rules"
    
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    cve_id = Column(String(20))
    rule_name = Column(String(255), nullable=False)
    rule_content = Column(Text, nullable=False)
    detection_type = Column(String(50))
    log_source = Column(String(100))
    confidence_level = Column(String(20))
    auto_generated = Column(Boolean, default=True)
    exploit_based = Column(Boolean, default=False)
    github_repos = Column(ARRAY(String))
    exploit_indicators = Column(Text)  # JSON string of extracted indicators
    # Enhanced fields for new data sources
    poc_source = Column(String(20), default='github_search')  # 'github_search', 'nomi_sec', 'manual'
    poc_quality_score = Column(Integer, default=0)  # Based on star count, activity, etc.
    nomi_sec_data = Column(JSON)  # Store nomi-sec PoC metadata
    created_at = Column(TIMESTAMP, default=datetime.utcnow)
    updated_at = Column(TIMESTAMP, default=datetime.utcnow)

class RuleTemplate(Base):
    __tablename__ = "rule_templates"
    
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    template_name = Column(String(255), nullable=False)
    template_content = Column(Text, nullable=False)
    applicable_product_patterns = Column(ARRAY(String))
    description = Column(Text)
    created_at = Column(TIMESTAMP, default=datetime.utcnow)

def get_db():
    """Get database session for migration purposes"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

# Create all tables (for migration purposes)
def create_tables():
    """Create database tables"""
    Base.metadata.create_all(bind=engine)