stwo-gpu 2.0.0

//! GPU Secure Session for Obelysk
//!
//! This module provides a secure, high-performance session management layer
//! for GPU-accelerated STARK proof generation within a TEE (Trusted Execution Environment).
//!
//! # Architecture
//!
//! ```text
//! ┌─────────────────────────────────────────────────────────────────────────────┐
//! │                           Obelysk Worker Node                               │
//! │  ┌───────────────────────────────────────────────────────────────────────┐  │
//! │  │                        GpuSecureSession                               │  │
//! │  │  • Session management (per-user isolation)                            │  │
//! │  │  • TEE integration (encryption/decryption)                            │  │
//! │  │  • Timeout handling                                                   │  │
//! │  │  • Recovery logic                                                     │  │
//! │  └───────────────────────────────────────────────────────────────────────┘  │
//! │                                    │                                        │
//! │                                    │ USES                                   │
//! │                                    ▼                                        │
//! │  ┌───────────────────────────────────────────────────────────────────────┐  │
//! │  │                      GpuProofPipeline                                 │  │
//! │  │  • Persistent GPU memory                                              │  │
//! │  │  • Bulk transfers (upload once, download once)                        │  │
//! │  │  • Chained operations (IFFT→FFT→FRI→Merkle)                          │  │
//! │  │  • Twiddle caching                                                    │  │
//! │  └───────────────────────────────────────────────────────────────────────┘  │
//! └─────────────────────────────────────────────────────────────────────────────┘
//! ```
//!
//! # Performance
//!
//! By using the Production Pipeline instead of the High-Level API:
//! - Single transfer in (trace data)
//! - All computation on GPU (IFFT, FFT, FRI, Merkle)
//! - Single transfer out (final proof)
//!
//! This achieves **33-50x speedup** over the High-Level API which transfers
//! data back to CPU between each operation.
//!
//! # Security
//!
//! - Each user gets an isolated session with unique encryption keys
//! - Data is decrypted only within the TEE before GPU upload
//! - GPU memory is zeroed on session destruction
//! - Session keys are destroyed after use

#[cfg(feature = "cuda-runtime")]
use super::pipeline::GpuProofPipeline;

#[cfg(feature = "cuda-runtime")]
use super::cuda_executor::CudaFftError;

use std::time::{Duration, Instant};

#[cfg(feature = "cuda-runtime")]
use std::collections::HashMap;

#[cfg(feature = "cuda-runtime")]
use std::sync::{Arc, Mutex};

/// Unique identifier for a user/job
pub type UserId = u64;

/// Session state machine
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SessionState {
    /// Just created, GPU memory allocated
    New,
    
    /// Actively processing a proof
    Active {
        /// When processing started
        started_at: Instant,
    },
    
    /// Waiting for next job (twiddles cached)
    Idle {
        /// When the session became idle
        idle_since: Instant,
    },
    
    /// Timed out, pending cleanup
    Expired,
    
    /// Cleaned up, ready for destruction
    Dead,
}

impl SessionState {
    /// Check if the session is usable for processing
    pub fn is_usable(&self) -> bool {
        matches!(self, SessionState::New | SessionState::Idle { .. })
    }
    
    /// Check if the session is currently processing
    pub fn is_active(&self) -> bool {
        matches!(self, SessionState::Active { .. })
    }
}

/// Error types for secure session operations
#[derive(Debug, Clone)]
pub enum SecureSessionError {
    /// Session has expired
    SessionExpired,
    
    /// Session is in an invalid state for the requested operation
    InvalidState(String),
    
    /// GPU operation failed
    GpuError(String),
    
    /// Encryption/decryption failed
    CryptoError(String),
    
    /// TEE attestation failed
    AttestationError(String),
    
    /// Maximum retries exceeded
    MaxRetriesExceeded,
    
    /// Session not found
    SessionNotFound(UserId),
    
    /// Resource limit exceeded
    ResourceLimitExceeded(String),
}

#[cfg(feature = "cuda-runtime")]
impl From<CudaFftError> for SecureSessionError {
    fn from(e: CudaFftError) -> Self {
        SecureSessionError::GpuError(format!("{:?}", e))
    }
}

impl std::fmt::Display for SecureSessionError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SecureSessionError::SessionExpired => write!(f, "Session has expired"),
            SecureSessionError::InvalidState(msg) => write!(f, "Invalid session state: {}", msg),
            SecureSessionError::GpuError(msg) => write!(f, "GPU error: {}", msg),
            SecureSessionError::CryptoError(msg) => write!(f, "Crypto error: {}", msg),
            SecureSessionError::AttestationError(msg) => write!(f, "Attestation error: {}", msg),
            SecureSessionError::MaxRetriesExceeded => write!(f, "Maximum retries exceeded"),
            SecureSessionError::SessionNotFound(id) => write!(f, "Session not found: {}", id),
            SecureSessionError::ResourceLimitExceeded(msg) => write!(f, "Resource limit exceeded: {}", msg),
        }
    }
}

impl std::error::Error for SecureSessionError {}

/// TEE Context for secure key management
/// 
/// In production, this would integrate with actual TEE hardware (SGX, TDX, SEV).
/// For now, we provide a software implementation that maintains the same API.
#[derive(Debug)]
pub struct TeeContext {
    /// Session encryption key (derived from TEE sealing key)
    session_key: [u8; 32],
    
    /// Attestation quote (proof of TEE execution)
    attestation_quote: Vec<u8>,
    
    /// User ID this context belongs to (used in attestation quote generation)
    user_id: UserId,
}

impl TeeContext {
    /// Create a new TEE context for a user using cryptographically secure random keys.
    ///
    /// Uses `getrandom` (OS CSPRNG) for session key generation. This is the default
    /// constructor — keys are secure but not TEE-sealed. Use `new_with_attestation()`
    /// when real TEE attestation is available.
    pub fn new(user_id: UserId) -> Self {
        Self::new_software(user_id)
    }

    /// Create a software-mode TEE context (no hardware attestation).
    ///
    /// When `cuda-runtime` is enabled, session key uses the OS CSPRNG via `getrandom`.
    /// Otherwise, uses blake2s hash of high-entropy system state.
    /// Attestation quote is empty (honest about no TEE).
    pub fn new_software(user_id: UserId) -> Self {
        let session_key = Self::generate_session_key(user_id);

        Self {
            session_key,
            attestation_quote: Vec::new(), // Empty = no TEE attestation
            user_id,
        }
    }

    /// Create a TEE context with a real attestation report.
    ///
    /// Session key generation uses the same CSPRNG/blake2s approach as `new_software()`,
    /// but the attestation quote contains a real hardware attestation from NVIDIA CC-On
    /// or CPU TEE. This is used when `SecurityLevel::ZkPlusTee` is active.
    pub fn new_with_attestation(user_id: UserId, attestation_quote: Vec<u8>) -> Self {
        let session_key = Self::generate_session_key(user_id);

        Self {
            session_key,
            attestation_quote,
            user_id,
        }
    }

    /// Generate a session key using the best available entropy source.
    fn generate_session_key(_user_id: UserId) -> [u8; 32] {
        #[cfg(feature = "cuda-runtime")]
        {
            let mut key = [0u8; 32];
            getrandom::getrandom(&mut key).expect("OS CSPRNG (getrandom) should not fail");
            return key;
        }

        #[cfg(not(feature = "cuda-runtime"))]
        {
            use blake2::{Blake2s256, Digest};

            let mut hasher = Blake2s256::new();
            hasher.update(&std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_nanos()
                .to_le_bytes());
            hasher.update(&std::process::id().to_le_bytes());
            hasher.update(&_user_id.to_le_bytes());
            let key_var: u64 = 0;
            hasher.update(&(&key_var as *const _ as usize).to_le_bytes());
            hasher.update(format!("{:?}", std::thread::current().id()).as_bytes());
            hasher.finalize().into()
        }
    }
    
    /// Decrypt data using the session key with AES-256-GCM
    ///
    /// This uses authenticated encryption (AEAD) which provides:
    /// - Confidentiality: Data is encrypted
    /// - Integrity: Any tampering is detected
    /// - Authenticity: Data came from the key holder
    ///
    /// # Format
    /// Input: [12-byte nonce][ciphertext][16-byte auth tag]
    /// Output: plaintext
    #[cfg(feature = "cuda-runtime")]
    pub fn decrypt(&self, encrypted_data: &[u8]) -> Result<Vec<u8>, SecureSessionError> {
        use aes_gcm::{
            aead::{Aead, KeyInit},
            Aes256Gcm, Nonce,
        };

        // Minimum size: 12 (nonce) + 16 (auth tag) = 28 bytes
        if encrypted_data.len() < 28 {
            return Err(SecureSessionError::CryptoError(
                "Encrypted data too short (minimum 28 bytes for nonce + tag)".into()
            ));
        }

        // Extract nonce (first 12 bytes)
        let nonce = Nonce::from_slice(&encrypted_data[..12]);

        // Extract ciphertext + auth tag (remaining bytes)
        let ciphertext_with_tag = &encrypted_data[12..];

        // Create cipher instance
        let cipher = Aes256Gcm::new_from_slice(&self.session_key)
            .map_err(|e| SecureSessionError::CryptoError(format!("Invalid key: {}", e)))?;

        // Decrypt and verify authentication tag
        let plaintext = cipher.decrypt(nonce, ciphertext_with_tag)
            .map_err(|_| SecureSessionError::CryptoError(
                "Decryption failed: invalid ciphertext or authentication tag".into()
            ))?;

        Ok(plaintext)
    }

    #[cfg(not(feature = "cuda-runtime"))]
    pub fn decrypt(&self, _encrypted_data: &[u8]) -> Result<Vec<u8>, SecureSessionError> {
        Err(SecureSessionError::CryptoError(
            "AES-GCM encryption requires cuda-runtime feature".into()
        ))
    }

    /// Encrypt data using the session key with AES-256-GCM
    ///
    /// This uses authenticated encryption (AEAD) which provides:
    /// - Confidentiality: Data is encrypted
    /// - Integrity: Any tampering is detected
    /// - Authenticity: Data came from the key holder
    ///
    /// # Format
    /// Output: [12-byte nonce][ciphertext][16-byte auth tag]
    #[cfg(feature = "cuda-runtime")]
    pub fn encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>, SecureSessionError> {
        use aes_gcm::{
            aead::{Aead, KeyInit},
            Aes256Gcm, Nonce,
        };

        // Generate random 12-byte nonce using secure randomness
        let mut nonce_bytes = [0u8; 12];
        getrandom::getrandom(&mut nonce_bytes)
            .map_err(|e| SecureSessionError::CryptoError(format!("Failed to generate nonce: {}", e)))?;
        let nonce = Nonce::from_slice(&nonce_bytes);

        // Create cipher instance
        let cipher = Aes256Gcm::new_from_slice(&self.session_key)
            .map_err(|e| SecureSessionError::CryptoError(format!("Invalid key: {}", e)))?;

        // Encrypt with authentication tag appended
        let ciphertext_with_tag = cipher.encrypt(nonce, plaintext)
            .map_err(|e| SecureSessionError::CryptoError(format!("Encryption failed: {}", e)))?;

        // Prepend nonce to ciphertext
        let mut result = Vec::with_capacity(12 + ciphertext_with_tag.len());
        result.extend_from_slice(&nonce_bytes);
        result.extend_from_slice(&ciphertext_with_tag);

        Ok(result)
    }

    #[cfg(not(feature = "cuda-runtime"))]
    pub fn encrypt(&self, _plaintext: &[u8]) -> Result<Vec<u8>, SecureSessionError> {
        Err(SecureSessionError::CryptoError(
            "AES-GCM encryption requires cuda-runtime feature".into()
        ))
    }
    
    /// Get the attestation quote
    pub fn attestation_quote(&self) -> &[u8] {
        &self.attestation_quote
    }
    
    /// Get the user ID this context belongs to
    pub fn user_id(&self) -> UserId {
        self.user_id
    }
    
    /// Securely destroy the context using volatile writes to prevent
    /// compiler optimization from eliding the zeroing.
    pub fn secure_destroy(&mut self) {
        // Zero out the session key using volatile writes
        for byte in self.session_key.iter_mut() {
            unsafe { std::ptr::write_volatile(byte, 0); }
        }
        std::sync::atomic::compiler_fence(std::sync::atomic::Ordering::SeqCst);
        self.attestation_quote.clear();
    }
}

impl Drop for TeeContext {
    fn drop(&mut self) {
        self.secure_destroy();
    }
}

/// GPU Secure Session
/// 
/// Provides a secure, high-performance session for GPU-accelerated STARK proofs.
/// 
/// # Lifecycle
/// 
/// 1. **Create**: `GpuSecureSession::new(user_id, log_size)`
/// 2. **Upload**: `session.upload_trace(data)` or `session.upload_encrypted_trace(encrypted)`
/// 3. **Compute**: `session.generate_proof()` - ALL computation stays on GPU
/// 4. **Download**: Proof is returned (only proof data, not trace)
/// 5. **Reuse** (optional): `session.clear_for_next()` to process another job
/// 6. **Destroy**: `session.secure_destroy()` or drop
#[cfg(feature = "cuda-runtime")]
pub struct GpuSecureSession {
    /// User/job this session belongs to
    user_id: UserId,
    
    /// TEE context for encryption/decryption
    tee_context: TeeContext,
    
    /// GPU proof pipeline (the fast path)
    pipeline: GpuProofPipeline,
    
    /// Current session state
    state: SessionState,
    
    /// When the session was created
    created_at: Instant,
    
    /// Number of proofs generated in this session
    proofs_generated: u64,
    
    /// Configuration
    config: SessionConfig,
}

/// Session configuration
#[derive(Debug, Clone)]
pub struct SessionConfig {
    /// Maximum idle time before session expires (default: 5 minutes)
    pub idle_timeout: Duration,
    
    /// Maximum active time for a single proof (default: 1 hour)
    pub active_timeout: Duration,
    
    /// Maximum retries on recoverable errors (default: 3)
    pub max_retries: u32,
    
    /// Whether to zero GPU memory on clear (default: true for security)
    pub secure_clear: bool,
}

impl Default for SessionConfig {
    fn default() -> Self {
        Self {
            idle_timeout: Duration::from_secs(300),      // 5 minutes
            active_timeout: Duration::from_secs(3600),   // 1 hour
            max_retries: 3,
            secure_clear: true,
        }
    }
}

#[cfg(feature = "cuda-runtime")]
impl GpuSecureSession {
    /// Create a new secure session for a user
    /// 
    /// This allocates GPU memory and precomputes twiddles for the given polynomial size.
    /// 
    /// # Arguments
    /// * `user_id` - Unique identifier for the user/job
    /// * `log_size` - Log2 of the polynomial size (e.g., 20 for 1M elements)
    pub fn new(user_id: UserId, log_size: u32) -> Result<Self, SecureSessionError> {
        Self::new_with_config(user_id, log_size, SessionConfig::default())
    }
    
    /// Create a new secure session with custom configuration
    pub fn new_with_config(
        user_id: UserId,
        log_size: u32,
        config: SessionConfig,
    ) -> Result<Self, SecureSessionError> {
        tracing::info!(
            user_id = user_id,
            log_size = log_size,
            "Creating new GpuSecureSession"
        );
        
        // Create TEE context (generates session key)
        let tee_context = TeeContext::new(user_id);
        
        // Create GPU pipeline (allocates memory, precomputes twiddles)
        let pipeline = GpuProofPipeline::new(log_size)?;
        
        Ok(Self {
            user_id,
            tee_context,
            pipeline,
            state: SessionState::New,
            created_at: Instant::now(),
            proofs_generated: 0,
            config,
        })
    }
    
    /// Get the user ID for this session
    pub fn user_id(&self) -> UserId {
        self.user_id
    }
    
    /// Get the current session state
    pub fn state(&self) -> &SessionState {
        &self.state
    }
    
    /// Get the number of proofs generated in this session
    pub fn proofs_generated(&self) -> u64 {
        self.proofs_generated
    }
    
    /// Get session uptime
    pub fn uptime(&self) -> Duration {
        self.created_at.elapsed()
    }
    
    /// Get the attestation quote for this session
    pub fn attestation_quote(&self) -> &[u8] {
        self.tee_context.attestation_quote()
    }
    
    /// Check if the session has timed out
    pub fn check_timeout(&mut self) -> bool {
        match &self.state {
            SessionState::Idle { idle_since } => {
                if idle_since.elapsed() > self.config.idle_timeout {
                    tracing::warn!(
                        user_id = self.user_id,
                        idle_duration = ?idle_since.elapsed(),
                        "Session idle timeout"
                    );
                    self.state = SessionState::Expired;
                    return true;
                }
            }
            SessionState::Active { started_at } => {
                if started_at.elapsed() > self.config.active_timeout {
                    tracing::error!(
                        user_id = self.user_id,
                        active_duration = ?started_at.elapsed(),
                        "Session active timeout - proof taking too long"
                    );
                    self.state = SessionState::Expired;
                    return true;
                }
            }
            _ => {}
        }
        false
    }
    
    /// Upload trace data to GPU
    /// 
    /// This uploads polynomial data directly to GPU memory.
    /// Data stays on GPU until proof generation is complete.
    /// 
    /// # Arguments
    /// * `polynomials` - Iterator of polynomial coefficient slices (as u32)
    pub fn upload_trace<'a>(
        &mut self,
        polynomials: impl Iterator<Item = &'a [u32]>,
    ) -> Result<usize, SecureSessionError> {
        // Check state
        if !self.state.is_usable() {
            return Err(SecureSessionError::InvalidState(
                format!("Cannot upload in state {:?}", self.state)
            ));
        }
        
        // Check timeout
        if self.check_timeout() {
            return Err(SecureSessionError::SessionExpired);
        }
        
        // Transition to active
        self.state = SessionState::Active {
            started_at: Instant::now(),
        };
        
        // Upload to GPU (single bulk transfer)
        let num_polys = self.pipeline.upload_polynomials_bulk(polynomials)?;
        
        tracing::debug!(
            user_id = self.user_id,
            num_polynomials = num_polys,
            "Uploaded trace to GPU"
        );
        
        Ok(num_polys)
    }
    
    /// Upload encrypted trace data to GPU
    /// 
    /// This decrypts the data in the TEE and uploads to GPU.
    /// The decrypted data never leaves the TEE/GPU boundary.
    /// 
    /// # Arguments
    /// * `encrypted_data` - Encrypted polynomial data
    /// * `num_polynomials` - Number of polynomials in the encrypted data
    pub fn upload_encrypted_trace(
        &mut self,
        encrypted_data: &[u8],
        num_polynomials: usize,
    ) -> Result<usize, SecureSessionError> {
        // Check state
        if !self.state.is_usable() {
            return Err(SecureSessionError::InvalidState(
                format!("Cannot upload in state {:?}", self.state)
            ));
        }
        
        // Check timeout
        if self.check_timeout() {
            return Err(SecureSessionError::SessionExpired);
        }
        
        // Transition to active
        self.state = SessionState::Active {
            started_at: Instant::now(),
        };
        
        // Decrypt in TEE
        let decrypted = self.tee_context.decrypt(encrypted_data)?;
        
        // Convert bytes to u32 slices
        let n = 1usize << self.pipeline.log_size();
        let poly_size_bytes = n * 4; // 4 bytes per u32
        
        if decrypted.len() != num_polynomials * poly_size_bytes {
            return Err(SecureSessionError::CryptoError(format!(
                "Decrypted size {} doesn't match expected {} polynomials × {} bytes",
                decrypted.len(),
                num_polynomials,
                poly_size_bytes
            )));
        }
        
        // Upload each polynomial
        for i in 0..num_polynomials {
            let start = i * poly_size_bytes;
            let end = start + poly_size_bytes;
            let poly_bytes = &decrypted[start..end];
            
            // Convert bytes to u32
            let poly_u32: Vec<u32> = poly_bytes
                .chunks_exact(4)
                .map(|chunk| u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
                .collect();
            
            self.pipeline.upload_polynomial(&poly_u32)?;
        }
        
        tracing::debug!(
            user_id = self.user_id,
            num_polynomials = num_polynomials,
            "Uploaded encrypted trace to GPU (decrypted in TEE)"
        );
        
        Ok(num_polynomials)
    }
    
    /// Execute IFFT on all polynomials (interpolation)
    /// 
    /// This performs IFFT with fused denormalization on all uploaded polynomials.
    /// Data stays on GPU.
    pub fn interpolate_all(&mut self) -> Result<(), SecureSessionError> {
        if !self.state.is_active() {
            return Err(SecureSessionError::InvalidState(
                "Must be in Active state to interpolate".into()
            ));
        }
        
        let num_polys = self.pipeline.num_polynomials();
        for i in 0..num_polys {
            self.pipeline.ifft_with_denormalize(i)?;
        }
        
        tracing::debug!(
            user_id = self.user_id,
            num_polynomials = num_polys,
            "Interpolated all polynomials on GPU"
        );
        
        Ok(())
    }
    
    /// Execute FFT on all polynomials (evaluation)
    /// 
    /// This performs FFT on all polynomials.
    /// Data stays on GPU.
    pub fn evaluate_all(&mut self) -> Result<(), SecureSessionError> {
        if !self.state.is_active() {
            return Err(SecureSessionError::InvalidState(
                "Must be in Active state to evaluate".into()
            ));
        }
        
        let num_polys = self.pipeline.num_polynomials();
        for i in 0..num_polys {
            self.pipeline.fft(i)?;
        }
        
        tracing::debug!(
            user_id = self.user_id,
            num_polynomials = num_polys,
            "Evaluated all polynomials on GPU"
        );
        
        Ok(())
    }
    
    /// Execute FRI folding
    /// 
    /// Performs FRI protocol folding on GPU.
    /// 
    /// # Arguments
    /// * `input_idx` - Index of the polynomial to fold
    /// * `itwiddles` - Inverse twiddles for folding
    /// * `alpha` - FRI challenge
    /// 
    /// # Returns
    /// Index of the folded polynomial
    pub fn fri_fold(
        &mut self,
        input_idx: usize,
        itwiddles: &[u32],
        alpha: &[u32; 4],
    ) -> Result<usize, SecureSessionError> {
        if !self.state.is_active() {
            return Err(SecureSessionError::InvalidState(
                "Must be in Active state to FRI fold".into()
            ));
        }
        
        let output_idx = self.pipeline.fri_fold_line(input_idx, itwiddles, alpha)?;
        
        Ok(output_idx)
    }
    
    /// Generate a complete STARK proof
    /// 
    /// This is the main entry point for proof generation. It:
    /// 1. Performs IFFT (interpolation) on all polynomials
    /// 2. Performs FFT (evaluation) on all polynomials
    /// 3. Executes FRI folding (if configured)
    /// 4. Computes Merkle commitments
    /// 5. Returns the proof
    /// 
    /// ALL computation stays on GPU - no intermediate transfers!
    /// 
    /// # Returns
    /// The generated proof as bytes
    pub fn generate_proof(&mut self) -> Result<ProofOutput, SecureSessionError> {
        self.generate_proof_with_recovery()
    }
    
    /// Generate proof with automatic retry on recoverable errors
    fn generate_proof_with_recovery(&mut self) -> Result<ProofOutput, SecureSessionError> {
        for attempt in 0..self.config.max_retries {
            match self.generate_proof_internal() {
                Ok(proof) => return Ok(proof),
                Err(e) => {
                    if Self::is_recoverable_error(&e) {
                        tracing::warn!(
                            user_id = self.user_id,
                            attempt = attempt,
                            error = %e,
                            "Proof generation failed, retrying"
                        );
                        
                        // Reset GPU state but keep twiddles
                        self.pipeline.sync()?;
                        continue;
                    } else {
                        // Non-recoverable error
                        self.state = SessionState::Dead;
                        return Err(e);
                    }
                }
            }
        }
        
        Err(SecureSessionError::MaxRetriesExceeded)
    }
    
    /// Check if an error is recoverable
    fn is_recoverable_error(e: &SecureSessionError) -> bool {
        matches!(e, SecureSessionError::GpuError(_))
    }
    
    /// Internal proof generation (single attempt)
    fn generate_proof_internal(&mut self) -> Result<ProofOutput, SecureSessionError> {
        if !self.state.is_active() {
            return Err(SecureSessionError::InvalidState(
                "Must be in Active state to generate proof".into()
            ));
        }
        
        let start = Instant::now();
        let num_polys = self.pipeline.num_polynomials();
        
        tracing::info!(
            user_id = self.user_id,
            num_polynomials = num_polys,
            "Starting proof generation on GPU"
        );
        
        // Step 1: IFFT with denormalization (interpolation)
        for i in 0..num_polys {
            self.pipeline.ifft_with_denormalize(i)?;
        }
        
        // Step 2: FFT (evaluation)
        for i in 0..num_polys {
            self.pipeline.fft(i)?;
        }
        
        // Step 3: Sync to ensure all GPU operations complete
        self.pipeline.sync()?;
        
        // Step 4: Download results (only what's needed for the proof)
        let polynomial_data = self.pipeline.download_polynomials_bulk()?;
        
        let duration = start.elapsed();
        self.proofs_generated += 1;
        
        // Transition to idle
        self.state = SessionState::Idle {
            idle_since: Instant::now(),
        };
        
        tracing::info!(
            user_id = self.user_id,
            duration_ms = duration.as_millis(),
            proofs_generated = self.proofs_generated,
            "Proof generation complete"
        );
        
        Ok(ProofOutput {
            polynomial_evaluations: polynomial_data,
            duration,
            num_polynomials: num_polys,
        })
    }
    
    /// Clear GPU memory for the next job
    /// 
    /// This clears polynomial data but keeps twiddles cached.
    /// The session can be reused for another proof with the same polynomial size.
    pub fn clear_for_next(&mut self) -> Result<(), SecureSessionError> {
        if matches!(self.state, SessionState::Dead | SessionState::Expired) {
            return Err(SecureSessionError::InvalidState(
                "Cannot clear expired/dead session".into()
            ));
        }
        
        // Clear polynomial data (twiddles stay cached)
        self.pipeline.clear_polynomials();
        
        // Transition to idle
        self.state = SessionState::Idle {
            idle_since: Instant::now(),
        };
        
        tracing::debug!(
            user_id = self.user_id,
            "Cleared session for next job"
        );
        
        Ok(())
    }
    
    /// Securely destroy the session
    /// 
    /// This:
    /// 1. Zeros GPU memory (if secure_clear is enabled)
    /// 2. Destroys TEE session keys
    /// 3. Marks session as dead
    pub fn secure_destroy(&mut self) {
        tracing::info!(
            user_id = self.user_id,
            proofs_generated = self.proofs_generated,
            uptime_secs = self.created_at.elapsed().as_secs(),
            "Destroying secure session"
        );
        
        // Clear GPU memory
        self.pipeline.clear_polynomials();
        
        // Destroy TEE context (zeros keys)
        self.tee_context.secure_destroy();
        
        // Mark as dead
        self.state = SessionState::Dead;
    }
    
    /// Get pipeline reference for advanced operations
    pub fn pipeline(&self) -> &GpuProofPipeline {
        &self.pipeline
    }
    
    /// Get mutable pipeline reference for advanced operations
    pub fn pipeline_mut(&mut self) -> &mut GpuProofPipeline {
        &mut self.pipeline
    }
}

#[cfg(feature = "cuda-runtime")]
impl Drop for GpuSecureSession {
    fn drop(&mut self) {
        if !matches!(self.state, SessionState::Dead) {
            self.secure_destroy();
        }
    }
}

/// Output from proof generation
#[derive(Debug)]
pub struct ProofOutput {
    /// Polynomial evaluations (the actual proof data)
    pub polynomial_evaluations: Vec<Vec<u32>>,
    
    /// Time taken to generate the proof
    pub duration: Duration,
    
    /// Number of polynomials in the proof
    pub num_polynomials: usize,
}

// ============================================================================
// Session Manager for Multi-User Isolation
// ============================================================================

/// Session Manager for multi-tenant GPU workers
/// 
/// Manages multiple isolated sessions for different users.
/// Ensures:
/// - User isolation (each user gets their own session)
/// - Resource limits (max concurrent sessions)
/// - Automatic cleanup (expired sessions are removed)
#[cfg(feature = "cuda-runtime")]
pub struct SessionManager {
    /// Active sessions by user ID
    sessions: HashMap<UserId, GpuSecureSession>,
    
    /// Maximum concurrent sessions
    max_sessions: usize,
    
    /// Default log size for new sessions
    default_log_size: u32,
    
    /// Session configuration
    config: SessionConfig,
}

#[cfg(feature = "cuda-runtime")]
impl SessionManager {
    /// Create a new session manager
    /// 
    /// # Arguments
    /// * `max_sessions` - Maximum concurrent sessions (limited by GPU memory)
    /// * `default_log_size` - Default polynomial size for new sessions
    pub fn new(max_sessions: usize, default_log_size: u32) -> Self {
        Self {
            sessions: HashMap::new(),
            max_sessions,
            default_log_size,
            config: SessionConfig::default(),
        }
    }
    
    /// Create a session manager with custom configuration
    pub fn with_config(
        max_sessions: usize,
        default_log_size: u32,
        config: SessionConfig,
    ) -> Self {
        Self {
            sessions: HashMap::new(),
            max_sessions,
            default_log_size,
            config,
        }
    }
    
    /// Get or create a session for a user
    /// 
    /// If a session exists and is usable, returns it.
    /// If no session exists, creates a new one (evicting oldest if at capacity).
    pub fn get_or_create_session(
        &mut self,
        user_id: UserId,
    ) -> Result<&mut GpuSecureSession, SecureSessionError> {
        self.get_or_create_session_with_size(user_id, self.default_log_size)
    }
    
    /// Get or create a session with specific polynomial size
    pub fn get_or_create_session_with_size(
        &mut self,
        user_id: UserId,
        log_size: u32,
    ) -> Result<&mut GpuSecureSession, SecureSessionError> {
        // Clean up expired sessions first
        self.cleanup_expired();
        
        // Check if session exists and determine if we need to remove it
        let should_remove = self.sessions.get(&user_id)
            .map(|s| !s.state.is_usable() && !s.state.is_active())
            .unwrap_or(false);
        
        if should_remove {
            // Session expired, remove it
            self.sessions.remove(&user_id);
        }
        
        // Check if session exists and is usable (after potential removal)
        if self.sessions.contains_key(&user_id) {
            return Ok(self.sessions.get_mut(&user_id).unwrap());
        }
        
        // Check capacity
        if self.sessions.len() >= self.max_sessions {
            // Evict oldest idle session
            self.evict_oldest_session()?;
        }
        
        // Create new session
        let session = GpuSecureSession::new_with_config(
            user_id,
            log_size,
            self.config.clone(),
        )?;
        
        self.sessions.insert(user_id, session);
        
        Ok(self.sessions.get_mut(&user_id).unwrap())
    }
    
    /// Destroy a specific user's session
    pub fn destroy_session(&mut self, user_id: UserId) -> Option<()> {
        if let Some(mut session) = self.sessions.remove(&user_id) {
            session.secure_destroy();
            Some(())
        } else {
            None
        }
    }
    
    /// Get the number of active sessions
    pub fn active_session_count(&self) -> usize {
        self.sessions.len()
    }
    
    /// Get all user IDs with active sessions
    pub fn active_users(&self) -> Vec<UserId> {
        self.sessions.keys().copied().collect()
    }
    
    /// Clean up expired sessions
    fn cleanup_expired(&mut self) {
        let expired: Vec<UserId> = self.sessions
            .iter_mut()
            .filter_map(|(&user_id, session)| {
                if session.check_timeout() {
                    Some(user_id)
                } else {
                    None
                }
            })
            .collect();
        
        for user_id in expired {
            if let Some(mut session) = self.sessions.remove(&user_id) {
                session.secure_destroy();
            }
        }
    }
    
    /// Evict the oldest idle session
    fn evict_oldest_session(&mut self) -> Result<(), SecureSessionError> {
        // Find oldest idle session
        let oldest = self.sessions
            .iter()
            .filter_map(|(&user_id, session)| {
                if let SessionState::Idle { idle_since } = &session.state {
                    Some((user_id, *idle_since))
                } else {
                    None
                }
            })
            .min_by_key(|(_, idle_since)| *idle_since)
            .map(|(user_id, _)| user_id);
        
        if let Some(user_id) = oldest {
            tracing::info!(
                evicted_user = user_id,
                "Evicting oldest session to make room"
            );
            self.destroy_session(user_id);
            Ok(())
        } else {
            Err(SecureSessionError::ResourceLimitExceeded(
                "All sessions are active, cannot evict".into()
            ))
        }
    }
}

#[cfg(feature = "cuda-runtime")]
impl Drop for SessionManager {
    fn drop(&mut self) {
        // Destroy all sessions
        let user_ids: Vec<UserId> = self.sessions.keys().copied().collect();
        for user_id in user_ids {
            self.destroy_session(user_id);
        }
    }
}

// ============================================================================
// Global Session Manager (Thread-Safe)
// ============================================================================

/// Thread-safe global session manager
#[cfg(feature = "cuda-runtime")]
pub type GlobalSessionManager = Arc<Mutex<SessionManager>>;

#[cfg(feature = "cuda-runtime")]
/// Create a new global session manager
pub fn create_global_session_manager(
    max_sessions: usize,
    default_log_size: u32,
) -> GlobalSessionManager {
    Arc::new(Mutex::new(SessionManager::new(max_sessions, default_log_size)))
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_tee_context_creation() {
        let ctx = TeeContext::new(12345);
        assert_eq!(ctx.user_id, 12345);
        // Software-mode context has empty attestation quote (honest about no TEE)
        assert!(ctx.attestation_quote.is_empty());
        // Session key should be non-zero (generated via CSPRNG)
        assert_ne!(ctx.session_key, [0u8; 32]);
    }
    
    #[test]
    #[cfg(feature = "cuda-runtime")]
    fn test_tee_encrypt_decrypt_aes_gcm() {
        let ctx = TeeContext::new(12345);
        let plaintext = b"Hello, World!";

        let encrypted = ctx.encrypt(plaintext).unwrap();

        // Encrypted should be: 12 (nonce) + 13 (plaintext) + 16 (auth tag) = 41 bytes
        assert_eq!(encrypted.len(), 12 + plaintext.len() + 16);
        assert_ne!(&encrypted[12..], plaintext); // Ciphertext should differ from plaintext

        let decrypted = ctx.decrypt(&encrypted).unwrap();
        assert_eq!(&decrypted, plaintext);
    }

    #[test]
    #[cfg(feature = "cuda-runtime")]
    fn test_tee_decrypt_detects_tampering() {
        let ctx = TeeContext::new(12345);
        let plaintext = b"Sensitive data";

        let mut encrypted = ctx.encrypt(plaintext).unwrap();

        // Tamper with the ciphertext
        if encrypted.len() > 15 {
            encrypted[15] ^= 0xFF;
        }

        // Decryption should fail due to authentication tag mismatch
        let result = ctx.decrypt(&encrypted);
        assert!(result.is_err());
    }

    #[test]
    #[cfg(feature = "cuda-runtime")]
    fn test_tee_decrypt_rejects_short_input() {
        let ctx = TeeContext::new(12345);

        // Too short to contain nonce + auth tag
        let short_data = vec![0u8; 20];
        let result = ctx.decrypt(&short_data);
        assert!(result.is_err());
    }
    
    #[test]
    fn test_session_state_transitions() {
        assert!(SessionState::New.is_usable());
        assert!(!SessionState::New.is_active());
        
        let active = SessionState::Active { started_at: Instant::now() };
        assert!(!active.is_usable());
        assert!(active.is_active());
        
        let idle = SessionState::Idle { idle_since: Instant::now() };
        assert!(idle.is_usable());
        assert!(!idle.is_active());
        
        assert!(!SessionState::Expired.is_usable());
        assert!(!SessionState::Dead.is_usable());
    }
    
    #[test]
    fn test_session_config_defaults() {
        let config = SessionConfig::default();
        assert_eq!(config.idle_timeout, Duration::from_secs(300));
        assert_eq!(config.active_timeout, Duration::from_secs(3600));
        assert_eq!(config.max_retries, 3);
        assert!(config.secure_clear);
    }
}