"""
Salesforce Audit and Compliance Framework for FoundryMatch
==========================================================
Comprehensive audit logging, compliance tracking, and rollback capabilities
for Salesforce data operations.
"""

import asyncio
import json
import logging
import hashlib
import tempfile
import zipfile
from dataclasses import dataclass, field, asdict
from datetime import datetime, timedelta
from enum import Enum
from pathlib import Path
from typing import Dict, List, Optional, Any, Tuple, Union
import uuid

import pandas as pd
import sqlite3
from cryptography.fernet import Fernet

from .core import SalesforceIntegration
from .workflows import WorkflowResult

log = logging.getLogger(__name__)

class AuditEventType(Enum):
    """Types of auditable events."""
    WORKFLOW_START = "workflow_start"
    WORKFLOW_COMPLETE = "workflow_complete"
    WORKFLOW_FAILED = "workflow_failed"
    DATA_BACKUP = "data_backup"
    DATA_RESTORE = "data_restore"
    VALIDATION_CHECK = "validation_check"
    API_CALL = "api_call"
    USER_ACTION = "user_action"
    SYSTEM_EVENT = "system_event"
    COMPLIANCE_CHECK = "compliance_check"

class ComplianceLevel(Enum):
    """Compliance requirement levels."""
    SOX = "sox"          # Sarbanes-Oxley
    GDPR = "gdpr"        # General Data Protection Regulation
    HIPAA = "hipaa"      # Health Insurance Portability and Accountability Act
    PCI_DSS = "pci_dss"  # Payment Card Industry Data Security Standard
    CUSTOM = "custom"    # Custom compliance requirements

@dataclass
class AuditEvent:
    """Represents a single audit event."""
    event_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    timestamp: datetime = field(default_factory=datetime.utcnow)
    event_type: AuditEventType = AuditEventType.SYSTEM_EVENT
    user_id: Optional[str] = None
    session_id: Optional[str] = None
    workflow_id: Optional[str] = None
    object_name: Optional[str] = None
    operation: Optional[str] = None
    records_affected: int = 0
    data_fingerprint: Optional[str] = None
    before_snapshot: Optional[str] = None
    after_snapshot: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    success: bool = True
    error_message: Optional[str] = None
    compliance_tags: List[ComplianceLevel] = field(default_factory=list)

@dataclass
class DataSnapshot:
    """Represents a point-in-time data snapshot."""
    snapshot_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    timestamp: datetime = field(default_factory=datetime.utcnow)
    workflow_id: str = ""
    object_name: str = ""
    operation_type: str = ""
    record_count: int = 0
    data_hash: str = ""
    backup_location: str = ""
    encryption_key_id: Optional[str] = None
    retention_period_days: int = 2555  # 7 years default
    metadata: Dict[str, Any] = field(default_factory=dict)

@dataclass
class ComplianceReport:
    """Compliance audit report."""
    report_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    generated_at: datetime = field(default_factory=datetime.utcnow)
    compliance_level: ComplianceLevel = ComplianceLevel.CUSTOM
    period_start: datetime = field(default_factory=lambda: datetime.utcnow() - timedelta(days=30))
    period_end: datetime = field(default_factory=datetime.utcnow)
    total_events: int = 0
    workflows_executed: int = 0
    records_processed: int = 0
    compliance_violations: List[Dict[str, Any]] = field(default_factory=list)
    recommendations: List[str] = field(default_factory=list)
    certification_status: str = "compliant"

class AuditLogger:
    """Centralized audit logging system with encryption and retention policies."""

    def __init__(self, db_path: Optional[Path] = None, encryption_key: Optional[bytes] = None):
        self.db_path = db_path or Path.home() / ".foundryops" / "audit.db"
        self.db_path.parent.mkdir(parents=True, exist_ok=True)

        # Initialize encryption
        self.encryption_key = encryption_key or Fernet.generate_key()
        self.cipher = Fernet(self.encryption_key)

        # Initialize database
        self._init_database()

        # Session tracking
        self.current_session_id = str(uuid.uuid4())
        self.current_user_id = None

    def _init_database(self):
        """Initialize the audit database schema."""
        try:
            conn = sqlite3.connect(str(self.db_path))
            cursor = conn.cursor()

            # Audit events table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS audit_events (
                    event_id TEXT PRIMARY KEY,
                    timestamp TEXT NOT NULL,
                    event_type TEXT NOT NULL,
                    user_id TEXT,
                    session_id TEXT,
                    workflow_id TEXT,
                    object_name TEXT,
                    operation TEXT,
                    records_affected INTEGER DEFAULT 0,
                    data_fingerprint TEXT,
                    before_snapshot TEXT,
                    after_snapshot TEXT,
                    metadata TEXT,
                    success BOOLEAN DEFAULT TRUE,
                    error_message TEXT,
                    compliance_tags TEXT,
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP
                )
            """)

            # Data snapshots table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS data_snapshots (
                    snapshot_id TEXT PRIMARY KEY,
                    timestamp TEXT NOT NULL,
                    workflow_id TEXT,
                    object_name TEXT NOT NULL,
                    operation_type TEXT NOT NULL,
                    record_count INTEGER DEFAULT 0,
                    data_hash TEXT NOT NULL,
                    backup_location TEXT NOT NULL,
                    encryption_key_id TEXT,
                    retention_period_days INTEGER DEFAULT 2555,
                    metadata TEXT,
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP
                )
            """)

            # Compliance reports table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS compliance_reports (
                    report_id TEXT PRIMARY KEY,
                    generated_at TEXT NOT NULL,
                    compliance_level TEXT NOT NULL,
                    period_start TEXT NOT NULL,
                    period_end TEXT NOT NULL,
                    total_events INTEGER DEFAULT 0,
                    workflows_executed INTEGER DEFAULT 0,
                    records_processed INTEGER DEFAULT 0,
                    compliance_violations TEXT,
                    recommendations TEXT,
                    certification_status TEXT DEFAULT 'compliant',
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP
                )
            """)

            # Create indexes for better performance
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_events_timestamp ON audit_events(timestamp)")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_events_workflow ON audit_events(workflow_id)")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_events_user ON audit_events(user_id)")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_snapshots_timestamp ON data_snapshots(timestamp)")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_snapshots_workflow ON data_snapshots(workflow_id)")

            conn.commit()
            conn.close()

            log.info(f"Audit database initialized at {self.db_path}")

        except Exception as e:
            log.error(f"Failed to initialize audit database: {e}", exc_info=True)
            raise

    async def log_event(self, event: AuditEvent) -> bool:
        """Log an audit event to the database."""
        try:
            # Set session and user context if not provided
            if not event.session_id:
                event.session_id = self.current_session_id
            if not event.user_id:
                event.user_id = self.current_user_id

            # Encrypt sensitive data
            encrypted_metadata = self._encrypt_data(json.dumps(event.metadata))
            encrypted_before = self._encrypt_data(event.before_snapshot) if event.before_snapshot else None
            encrypted_after = self._encrypt_data(event.after_snapshot) if event.after_snapshot else None

            # Insert into database
            conn = sqlite3.connect(str(self.db_path))
            cursor = conn.cursor()

            cursor.execute("""
                INSERT INTO audit_events (
                    event_id, timestamp, event_type, user_id, session_id, workflow_id,
                    object_name, operation, records_affected, data_fingerprint,
                    before_snapshot, after_snapshot, metadata, success, error_message,
                    compliance_tags
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                event.event_id,
                event.timestamp.isoformat(),
                event.event_type.value,
                event.user_id,
                event.session_id,
                event.workflow_id,
                event.object_name,
                event.operation,
                event.records_affected,
                event.data_fingerprint,
                encrypted_before,
                encrypted_after,
                encrypted_metadata,
                event.success,
                event.error_message,
                json.dumps([tag.value for tag in event.compliance_tags])
            ))

            conn.commit()
            conn.close()

            log.debug(f"Logged audit event: {event.event_id}")
            return True

        except Exception as e:
            log.error(f"Failed to log audit event: {e}", exc_info=True)
            return False

    async def create_data_snapshot(self,
                                  data: pd.DataFrame,
                                  workflow_id: str,
                                  object_name: str,
                                  operation_type: str,
                                  metadata: Optional[Dict[str, Any]] = None) -> Optional[DataSnapshot]:
        """Create and store a data snapshot."""
        try:
            # Calculate data hash for integrity verification
            data_hash = self._calculate_data_hash(data)

            # Create backup file
            backup_location = await self._create_backup_file(data, workflow_id, object_name)

            # Create snapshot record
            snapshot = DataSnapshot(
                workflow_id=workflow_id,
                object_name=object_name,
                operation_type=operation_type,
                record_count=len(data),
                data_hash=data_hash,
                backup_location=backup_location,
                metadata=metadata or {}
            )

            # Store in database
            conn = sqlite3.connect(str(self.db_path))
            cursor = conn.cursor()

            cursor.execute("""
                INSERT INTO data_snapshots (
                    snapshot_id, timestamp, workflow_id, object_name, operation_type,
                    record_count, data_hash, backup_location, encryption_key_id,
                    retention_period_days, metadata
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                snapshot.snapshot_id,
                snapshot.timestamp.isoformat(),
                snapshot.workflow_id,
                snapshot.object_name,
                snapshot.operation_type,
                snapshot.record_count,
                snapshot.data_hash,
                snapshot.backup_location,
                None,  # encryption_key_id - could be implemented for key rotation
                snapshot.retention_period_days,
                self._encrypt_data(json.dumps(snapshot.metadata))
            ))

            conn.commit()
            conn.close()

            log.info(f"Created data snapshot: {snapshot.snapshot_id}")
            return snapshot

        except Exception as e:
            log.error(f"Failed to create data snapshot: {e}", exc_info=True)
            return None

    async def get_audit_trail(self,
                             workflow_id: Optional[str] = None,
                             start_date: Optional[datetime] = None,
                             end_date: Optional[datetime] = None,
                             event_types: Optional[List[AuditEventType]] = None) -> List[AuditEvent]:
        """Retrieve audit trail with optional filters."""
        try:
            conn = sqlite3.connect(str(self.db_path))
            cursor = conn.cursor()

            # Build query with filters
            query = "SELECT * FROM audit_events WHERE 1=1"
            params = []

            if workflow_id:
                query += " AND workflow_id = ?"
                params.append(workflow_id)

            if start_date:
                query += " AND timestamp >= ?"
                params.append(start_date.isoformat())

            if end_date:
                query += " AND timestamp <= ?"
                params.append(end_date.isoformat())

            if event_types:
                placeholders = ",".join("?" * len(event_types))
                query += f" AND event_type IN ({placeholders})"
                params.extend([et.value for et in event_types])

            query += " ORDER BY timestamp DESC"

            cursor.execute(query, params)
            rows = cursor.fetchall()
            conn.close()

            # Convert rows to AuditEvent objects
            events = []
            columns = [desc[0] for desc in cursor.description]

            for row in rows:
                row_dict = dict(zip(columns, row))

                # Decrypt sensitive data
                metadata = json.loads(self._decrypt_data(row_dict['metadata'])) if row_dict['metadata'] else {}
                before_snapshot = self._decrypt_data(row_dict['before_snapshot']) if row_dict['before_snapshot'] else None
                after_snapshot = self._decrypt_data(row_dict['after_snapshot']) if row_dict['after_snapshot'] else None
                compliance_tags = [ComplianceLevel(tag) for tag in json.loads(row_dict['compliance_tags'] or '[]')]

                event = AuditEvent(
                    event_id=row_dict['event_id'],
                    timestamp=datetime.fromisoformat(row_dict['timestamp']),
                    event_type=AuditEventType(row_dict['event_type']),
                    user_id=row_dict['user_id'],
                    session_id=row_dict['session_id'],
                    workflow_id=row_dict['workflow_id'],
                    object_name=row_dict['object_name'],
                    operation=row_dict['operation'],
                    records_affected=row_dict['records_affected'],
                    data_fingerprint=row_dict['data_fingerprint'],
                    before_snapshot=before_snapshot,
                    after_snapshot=after_snapshot,
                    metadata=metadata,
                    success=bool(row_dict['success']),
                    error_message=row_dict['error_message'],
                    compliance_tags=compliance_tags
                )
                events.append(event)

            return events

        except Exception as e:
            log.error(f"Failed to retrieve audit trail: {e}", exc_info=True)
            return []

    async def generate_compliance_report(self,
                                        compliance_level: ComplianceLevel,
                                        start_date: Optional[datetime] = None,
                                        end_date: Optional[datetime] = None) -> ComplianceReport:
        """Generate a compliance report for the specified period."""
        try:
            if not start_date:
                start_date = datetime.utcnow() - timedelta(days=30)
            if not end_date:
                end_date = datetime.utcnow()

            # Get relevant audit events
            events = await self.get_audit_trail(start_date=start_date, end_date=end_date)

            # Filter events by compliance level
            compliance_events = [
                event for event in events
                if compliance_level in event.compliance_tags or compliance_level == ComplianceLevel.CUSTOM
            ]

            # Analyze compliance
            violations = []
            total_records_processed = sum(event.records_affected for event in compliance_events)
            workflow_count = len(set(event.workflow_id for event in compliance_events if event.workflow_id))

            # Check for compliance violations based on level
            if compliance_level == ComplianceLevel.SOX:
                violations.extend(self._check_sox_compliance(compliance_events))
            elif compliance_level == ComplianceLevel.GDPR:
                violations.extend(self._check_gdpr_compliance(compliance_events))
            elif compliance_level == ComplianceLevel.HIPAA:
                violations.extend(self._check_hipaa_compliance(compliance_events))
            elif compliance_level == ComplianceLevel.PCI_DSS:
                violations.extend(self._check_pci_compliance(compliance_events))

            # Generate recommendations
            recommendations = self._generate_compliance_recommendations(compliance_level, violations, compliance_events)

            # Determine certification status
            certification_status = "non_compliant" if violations else "compliant"

            report = ComplianceReport(
                compliance_level=compliance_level,
                period_start=start_date,
                period_end=end_date,
                total_events=len(compliance_events),
                workflows_executed=workflow_count,
                records_processed=total_records_processed,
                compliance_violations=violations,
                recommendations=recommendations,
                certification_status=certification_status
            )

            # Store report in database
            await self._store_compliance_report(report)

            return report

        except Exception as e:
            log.error(f"Failed to generate compliance report: {e}", exc_info=True)
            return ComplianceReport(
                compliance_level=compliance_level,
                certification_status="error",
                recommendations=[f"Report generation failed: {str(e)}"]
            )

    async def rollback_workflow(self, workflow_id: str) -> bool:
        """Attempt to rollback a workflow using data snapshots."""
        try:
            log.info(f"Attempting rollback for workflow: {workflow_id}")

            # Get snapshots for this workflow
            snapshots = await self._get_workflow_snapshots(workflow_id)
            if not snapshots:
                log.error(f"No snapshots found for workflow {workflow_id}")
                return False

            # For now, we'll log the rollback attempt
            # In a full implementation, this would restore data from snapshots
            await self.log_event(AuditEvent(
                event_type=AuditEventType.DATA_RESTORE,
                workflow_id=workflow_id,
                operation="rollback",
                metadata={"snapshot_count": len(snapshots)},
                compliance_tags=[ComplianceLevel.SOX, ComplianceLevel.GDPR]
            ))

            log.info(f"Rollback simulation completed for workflow {workflow_id}")
            return True

        except Exception as e:
            log.error(f"Rollback failed for workflow {workflow_id}: {e}", exc_info=True)
            return False

    def _encrypt_data(self, data: str) -> str:
        """Encrypt sensitive data."""
        if not data:
            return ""
        return self.cipher.encrypt(data.encode()).decode()

    def _decrypt_data(self, encrypted_data: str) -> str:
        """Decrypt sensitive data."""
        if not encrypted_data:
            return ""
        return self.cipher.decrypt(encrypted_data.encode()).decode()

    def _calculate_data_hash(self, data: pd.DataFrame) -> str:
        """Calculate SHA-256 hash of DataFrame for integrity verification."""
        # Convert DataFrame to consistent string representation
        data_string = data.to_csv(index=False, sort=False)
        return hashlib.sha256(data_string.encode()).hexdigest()

    async def _create_backup_file(self, data: pd.DataFrame, workflow_id: str, object_name: str) -> str:
        """Create encrypted backup file."""
        try:
            # Create backup directory
            backup_dir = Path.home() / ".foundryops" / "backups"
            backup_dir.mkdir(parents=True, exist_ok=True)

            # Generate filename
            timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
            filename = f"{workflow_id}_{object_name}_{timestamp}.csv.enc"
            backup_path = backup_dir / filename

            # Encrypt and save data
            csv_data = data.to_csv(index=False)
            encrypted_data = self._encrypt_data(csv_data)

            with open(backup_path, 'w') as f:
                f.write(encrypted_data)

            return str(backup_path)

        except Exception as e:
            log.error(f"Failed to create backup file: {e}", exc_info=True)
            raise

    async def _get_workflow_snapshots(self, workflow_id: str) -> List[DataSnapshot]:
        """Get all snapshots for a workflow."""
        try:
            conn = sqlite3.connect(str(self.db_path))
            cursor = conn.cursor()

            cursor.execute("""
                SELECT * FROM data_snapshots
                WHERE workflow_id = ?
                ORDER BY timestamp DESC
            """, (workflow_id,))

            rows = cursor.fetchall()
            conn.close()

            snapshots = []
            columns = [desc[0] for desc in cursor.description]

            for row in rows:
                row_dict = dict(zip(columns, row))
                metadata = json.loads(self._decrypt_data(row_dict['metadata'])) if row_dict['metadata'] else {}

                snapshot = DataSnapshot(
                    snapshot_id=row_dict['snapshot_id'],
                    timestamp=datetime.fromisoformat(row_dict['timestamp']),
                    workflow_id=row_dict['workflow_id'],
                    object_name=row_dict['object_name'],
                    operation_type=row_dict['operation_type'],
                    record_count=row_dict['record_count'],
                    data_hash=row_dict['data_hash'],
                    backup_location=row_dict['backup_location'],
                    encryption_key_id=row_dict['encryption_key_id'],
                    retention_period_days=row_dict['retention_period_days'],
                    metadata=metadata
                )
                snapshots.append(snapshot)

            return snapshots

        except Exception as e:
            log.error(f"Failed to get workflow snapshots: {e}", exc_info=True)
            return []

    def _check_sox_compliance(self, events: List[AuditEvent]) -> List[Dict[str, Any]]:
        """Check Sarbanes-Oxley compliance requirements."""
        violations = []

        # Check for proper audit trail
        financial_events = [e for e in events if 'financial' in str(e.metadata).lower()]
        for event in financial_events:
            if not event.before_snapshot or not event.after_snapshot:
                violations.append({
                    "type": "missing_audit_trail",
                    "event_id": event.event_id,
                    "description": "Financial data changes must have complete before/after audit trail"
                })

        # Check for unauthorized access
        failed_events = [e for e in events if not e.success]
        if len(failed_events) > len(events) * 0.05:  # More than 5% failure rate
            violations.append({
                "type": "high_failure_rate",
                "description": f"High failure rate ({len(failed_events)/len(events):.1%}) may indicate unauthorized access attempts"
            })

        return violations

    def _check_gdpr_compliance(self, events: List[AuditEvent]) -> List[Dict[str, Any]]:
        """Check GDPR compliance requirements."""
        violations = []

        # Check for personal data processing justification
        personal_data_events = [e for e in events if 'personal' in str(e.metadata).lower() or 'email' in str(e.metadata).lower()]
        for event in personal_data_events:
            if 'legal_basis' not in event.metadata:
                violations.append({
                    "type": "missing_legal_basis",
                    "event_id": event.event_id,
                    "description": "Personal data processing requires documented legal basis"
                })

        # Check for data retention compliance
        old_snapshots = [e for e in events if e.timestamp < datetime.utcnow() - timedelta(days=365)]
        if old_snapshots:
            violations.append({
                "type": "data_retention_violation",
                "description": f"Found {len(old_snapshots)} events older than 1 year - check retention policies"
            })

        return violations

    def _check_hipaa_compliance(self, events: List[AuditEvent]) -> List[Dict[str, Any]]:
        """Check HIPAA compliance requirements."""
        violations = []

        # Check for PHI access logging
        phi_events = [e for e in events if 'health' in str(e.metadata).lower() or 'medical' in str(e.metadata).lower()]
        for event in phi_events:
            if not event.user_id:
                violations.append({
                    "type": "missing_user_identification",
                    "event_id": event.event_id,
                    "description": "PHI access must be associated with specific user"
                })

        return violations

    def _check_pci_compliance(self, events: List[AuditEvent]) -> List[Dict[str, Any]]:
        """Check PCI DSS compliance requirements."""
        violations = []

        # Check for payment data handling
        payment_events = [e for e in events if 'payment' in str(e.metadata).lower() or 'card' in str(e.metadata).lower()]
        for event in payment_events:
            if not event.data_fingerprint:
                violations.append({
                    "type": "missing_data_fingerprint",
                    "event_id": event.event_id,
                    "description": "Payment data operations must have data fingerprint for integrity"
                })

        return violations

    def _generate_compliance_recommendations(self,
                                           compliance_level: ComplianceLevel,
                                           violations: List[Dict[str, Any]],
                                           events: List[AuditEvent]) -> List[str]:
        """Generate compliance recommendations."""
        recommendations = []

        if violations:
            recommendations.append(f"Address {len(violations)} compliance violations identified in the report")

        if compliance_level == ComplianceLevel.SOX:
            recommendations.extend([
                "Ensure all financial data changes have complete audit trails",
                "Implement automated controls for financial data access",
                "Regular review of audit logs for unauthorized access attempts"
            ])

        elif compliance_level == ComplianceLevel.GDPR:
            recommendations.extend([
                "Document legal basis for all personal data processing",
                "Implement automated data retention policies",
                "Provide data subject access and deletion capabilities"
            ])

        elif compliance_level == ComplianceLevel.HIPAA:
            recommendations.extend([
                "Ensure all PHI access is logged with user identification",
                "Implement encryption for all PHI at rest and in transit",
                "Regular access reviews and authorization audits"
            ])

        elif compliance_level == ComplianceLevel.PCI_DSS:
            recommendations.extend([
                "Implement strong encryption for all payment data",
                "Regular vulnerability scans and penetration testing",
                "Maintain secure development practices"
            ])

        # General recommendations based on event patterns
        if len(events) == 0:
            recommendations.append("No audit events found - ensure audit logging is properly configured")

        failed_rate = len([e for e in events if not e.success]) / len(events) if events else 0
        if failed_rate > 0.02:  # More than 2% failure rate
            recommendations.append(f"High failure rate ({failed_rate:.1%}) - review system reliability and access controls")

        return recommendations

    async def _store_compliance_report(self, report: ComplianceReport):
        """Store compliance report in database."""
        try:
            conn = sqlite3.connect(str(self.db_path))
            cursor = conn.cursor()

            cursor.execute("""
                INSERT INTO compliance_reports (
                    report_id, generated_at, compliance_level, period_start, period_end,
                    total_events, workflows_executed, records_processed,
                    compliance_violations, recommendations, certification_status
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                report.report_id,
                report.generated_at.isoformat(),
                report.compliance_level.value,
                report.period_start.isoformat(),
                report.period_end.isoformat(),
                report.total_events,
                report.workflows_executed,
                report.records_processed,
                self._encrypt_data(json.dumps(report.compliance_violations)),
                self._encrypt_data(json.dumps(report.recommendations)),
                report.certification_status
            ))

            conn.commit()
            conn.close()

        except Exception as e:
            log.error(f"Failed to store compliance report: {e}", exc_info=True)
            raise

    def set_user_context(self, user_id: str):
        """Set the current user context for audit logging."""
        self.current_user_id = user_id

    def start_new_session(self) -> str:
        """Start a new audit session."""
        self.current_session_id = str(uuid.uuid4())
        return self.current_session_id

class SalesforceAuditIntegration:
    """Integration layer for Salesforce workflow auditing."""

    def __init__(self, audit_logger: AuditLogger, sf_integration: SalesforceIntegration):
        self.audit_logger = audit_logger
        self.sf_integration = sf_integration

    async def audit_workflow_start(self,
                                  workflow_id: str,
                                  object_name: str,
                                  operation: str,
                                  data: pd.DataFrame,
                                  compliance_levels: List[ComplianceLevel] = None) -> bool:
        """Audit the start of a Salesforce workflow."""
        try:
            # Create data snapshot
            snapshot = await self.audit_logger.create_data_snapshot(
                data, workflow_id, object_name, operation
            )

            # Log workflow start event
            event = AuditEvent(
                event_type=AuditEventType.WORKFLOW_START,
                workflow_id=workflow_id,
                object_name=object_name,
                operation=operation,
                records_affected=len(data),
                data_fingerprint=self.audit_logger._calculate_data_hash(data),
                before_snapshot=snapshot.snapshot_id if snapshot else None,
                metadata={
                    "data_columns": list(data.columns),
                    "org_id": self.sf_integration.credentials.org_id if self.sf_integration.credentials else None,
                    "instance_url": self.sf_integration.credentials.instance_url if self.sf_integration.credentials else None
                },
                compliance_tags=compliance_levels or []
            )

            return await self.audit_logger.log_event(event)

        except Exception as e:
            log.error(f"Failed to audit workflow start: {e}", exc_info=True)
            return False

    async def audit_workflow_completion(self,
                                       workflow_result: WorkflowResult,
                                       compliance_levels: List[ComplianceLevel] = None) -> bool:
        """Audit the completion of a Salesforce workflow."""
        try:
            event = AuditEvent(
                event_type=AuditEventType.WORKFLOW_COMPLETE if workflow_result.success else AuditEventType.WORKFLOW_FAILED,
                workflow_id=workflow_result.workflow_id,
                records_affected=workflow_result.records_processed,
                success=workflow_result.success,
                error_message=workflow_result.error_message,
                metadata={
                    "records_updated": workflow_result.records_updated,
                    "records_created": workflow_result.records_created,
                    "records_merged": workflow_result.records_merged,
                    "api_calls_used": workflow_result.api_calls_used,
                    "execution_time": workflow_result.execution_time,
                    "backup_file": workflow_result.backup_file
                },
                compliance_tags=compliance_levels or []
            )

            return await self.audit_logger.log_event(event)

        except Exception as e:
            log.error(f"Failed to audit workflow completion: {e}", exc_info=True)
            return False

    async def audit_api_call(self,
                            operation: str,
                            object_name: str,
                            records_affected: int = 0,
                            success: bool = True,
                            error_message: str = None) -> bool:
        """Audit a Salesforce API call."""
        try:
            event = AuditEvent(
                event_type=AuditEventType.API_CALL,
                object_name=object_name,
                operation=operation,
                records_affected=records_affected,
                success=success,
                error_message=error_message,
                metadata={
                    "api_version": "59.0",
                    "org_id": self.sf_integration.credentials.org_id if self.sf_integration.credentials else None
                }
            )

            return await self.audit_logger.log_event(event)

        except Exception as e:
            log.error(f"Failed to audit API call: {e}", exc_info=True)
            return False

# Helper functions for integration

async def create_audit_logger(config: Dict[str, Any] = None) -> AuditLogger:
    """Create and configure an audit logger."""
    config = config or {}

    db_path = Path(config.get('db_path', Path.home() / ".foundryops" / "audit.db"))
    encryption_key = config.get('encryption_key')

    return AuditLogger(db_path=db_path, encryption_key=encryption_key)

async def generate_audit_export(audit_logger: AuditLogger,
                               start_date: datetime,
                               end_date: datetime,
                               export_path: Path) -> bool:
    """Export audit data for external analysis or compliance."""
    try:
        # Get audit events
        events = await audit_logger.get_audit_trail(start_date=start_date, end_date=end_date)

        # Convert to DataFrame for export
        events_data = []
        for event in events:
            event_dict = asdict(event)
            event_dict['timestamp'] = event_dict['timestamp'].isoformat()
            event_dict['event_type'] = event_dict['event_type'].value
            event_dict['compliance_tags'] = [tag.value for tag in event_dict['compliance_tags']]
            events_data.append(event_dict)

        df = pd.DataFrame(events_data)

        # Export to multiple formats
        df.to_csv(export_path.with_suffix('.csv'), index=False)
        df.to_excel(export_path.with_suffix('.xlsx'), index=False)

        # Create JSON export with full metadata
        with open(export_path.with_suffix('.json'), 'w') as f:
            json.dump(events_data, f, indent=2, default=str)

        log.info(f"Audit data exported to {export_path}")
        return True

    except Exception as e:
        log.error(f"Failed to export audit data: {e}", exc_info=True)
        return False

def calculate_audit_integrity_score(events: List[AuditEvent]) -> float:
    """Calculate integrity score for audit trail."""
    if not events:
        return 0.0

    score = 100.0

    # Check for gaps in timeline
    sorted_events = sorted(events, key=lambda e: e.timestamp)
    for i in range(1, len(sorted_events)):
        time_gap = (sorted_events[i].timestamp - sorted_events[i-1].timestamp).total_seconds()
        if time_gap > 86400:  # More than 24 hours
            score -= 5.0

    # Check for incomplete events
    incomplete_events = [e for e in events if not e.success and not e.error_message]
    if incomplete_events:
        score -= len(incomplete_events) / len(events) * 20

    # Check for missing metadata
    events_without_metadata = [e for e in events if not e.metadata]
    if events_without_metadata:
        score -= len(events_without_metadata) / len(events) * 10

    return max(0.0, score)
