//! Profiling Integration
//!
//! This module provides profiling and performance analysis tools:
//! - NVTX/rocTX range markers for timeline visualization
//! - Automatic kernel timing and statistics
//! - Memory usage tracking
//! - Communication profiling
//! - Flamegraph generation support
//!
//! Integration with NVIDIA Nsight Systems and AMD ROCm profiler.

use crate::error::{DistributedError, Result};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::{Duration, Instant};

// =============================================================================
// Configuration
// =============================================================================

/// Profiling level
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ProfilingLevel {
    /// No profiling
    Off,
    /// Basic profiling (major operations only)
    Basic,
    /// Detailed profiling (includes sub-operations)
    Detailed,
    /// Full profiling (all operations, high overhead)
    Full,
}

/// Profiler backend
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ProfilerBackend {
    /// NVIDIA NVTX (Nsight Systems)
    Nvtx,
    /// AMD rocTX (ROCm profiler)
    RocTx,
    /// Built-in timing (no external dependency)
    Builtin,
    /// Chrome trace format
    ChromeTrace,
}

/// Configuration for profiling
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProfilingConfig {
    /// Profiling level
    pub level: ProfilingLevel,
    /// Profiler backend
    pub backend: ProfilerBackend,
    /// Enable memory profiling
    pub profile_memory: bool,
    /// Enable communication profiling
    pub profile_comm: bool,
    /// Enable kernel profiling
    pub profile_kernels: bool,
    /// Sample rate for detailed profiling (1 = every op, 10 = every 10th)
    pub sample_rate: u32,
    /// Output directory for traces
    pub output_dir: Option<String>,
    /// Enable automatic range nesting
    pub auto_nest: bool,
    /// Maximum range depth
    pub max_depth: usize,
}

impl Default for ProfilingConfig {
    fn default() -> Self {
        Self {
            level: ProfilingLevel::Basic,
            backend: ProfilerBackend::Builtin,
            profile_memory: true,
            profile_comm: true,
            profile_kernels: true,
            sample_rate: 1,
            output_dir: None,
            auto_nest: true,
            max_depth: 32,
        }
    }
}

// =============================================================================
// NVTX Colors
// =============================================================================

/// NVTX color for range visualization
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct NvtxColor(pub u32);

impl NvtxColor {
    pub const RED: Self = Self(0xFF0000);
    pub const GREEN: Self = Self(0x00FF00);
    pub const BLUE: Self = Self(0x0000FF);
    pub const YELLOW: Self = Self(0xFFFF00);
    pub const CYAN: Self = Self(0x00FFFF);
    pub const MAGENTA: Self = Self(0xFF00FF);
    pub const ORANGE: Self = Self(0xFFA500);
    pub const PURPLE: Self = Self(0x800080);
    pub const WHITE: Self = Self(0xFFFFFF);
    pub const GRAY: Self = Self(0x808080);

    /// Color for forward pass
    pub const FORWARD: Self = Self::GREEN;
    /// Color for backward pass
    pub const BACKWARD: Self = Self::BLUE;
    /// Color for optimizer step
    pub const OPTIMIZER: Self = Self::ORANGE;
    /// Color for communication
    pub const COMM: Self = Self::YELLOW;
    /// Color for memory operations
    pub const MEMORY: Self = Self::CYAN;
    /// Color for data loading
    pub const DATA: Self = Self::PURPLE;
}

// =============================================================================
// Profiling Range
// =============================================================================

/// A profiling range (NVTX range equivalent)
#[derive(Debug, Clone)]
pub struct ProfilingRange {
    /// Range name
    pub name: String,
    /// Category
    pub category: String,
    /// Color for visualization
    pub color: NvtxColor,
    /// Start time
    pub start_time: Instant,
    /// End time (None if still active)
    pub end_time: Option<Instant>,
    /// Parent range ID (for nesting)
    pub parent_id: Option<u64>,
    /// Range ID
    pub id: u64,
    /// Thread ID
    pub thread_id: u64,
    /// Additional metadata
    pub metadata: HashMap<String, String>,
}

impl ProfilingRange {
    /// Create a new range
    pub fn new(id: u64, name: String, category: String, color: NvtxColor) -> Self {
        Self {
            name,
            category,
            color,
            start_time: Instant::now(),
            end_time: None,
            parent_id: None,
            id,
            thread_id: {
                use std::hash::{Hash, Hasher};
                let mut hasher = std::collections::hash_map::DefaultHasher::new();
                std::thread::current().id().hash(&mut hasher);
                hasher.finish()
            },
            metadata: HashMap::new(),
        }
    }

    /// End the range
    pub fn end(&mut self) {
        self.end_time = Some(Instant::now());
    }

    /// Get duration
    pub fn duration(&self) -> Option<Duration> {
        self.end_time.map(|end| end.duration_since(self.start_time))
    }

    /// Add metadata
    pub fn with_metadata(mut self, key: &str, value: &str) -> Self {
        self.metadata.insert(key.to_string(), value.to_string());
        self
    }
}

// =============================================================================
// Range Guard (RAII)
// =============================================================================

/// RAII guard for automatic range ending
pub struct RangeGuard {
    profiler: Arc<Profiler>,
    range_id: u64,
}

impl RangeGuard {
    /// Create a new guard
    pub fn new(profiler: Arc<Profiler>, range_id: u64) -> Self {
        Self { profiler, range_id }
    }

    /// Add metadata to the range
    pub fn add_metadata(&self, key: &str, value: &str) {
        self.profiler.add_range_metadata(self.range_id, key, value);
    }
}

impl Drop for RangeGuard {
    fn drop(&mut self) {
        self.profiler.end_range(self.range_id);
    }
}

// =============================================================================
// Operation Statistics
// =============================================================================

/// Statistics for a profiled operation
#[derive(Debug, Clone, Default)]
pub struct OperationStats {
    /// Operation name
    pub name: String,
    /// Number of calls
    pub call_count: u64,
    /// Total time spent
    pub total_time: Duration,
    /// Minimum time
    pub min_time: Duration,
    /// Maximum time
    pub max_time: Duration,
    /// Total bytes processed (if applicable)
    pub total_bytes: u64,
}

impl OperationStats {
    /// Create new stats for an operation
    pub fn new(name: String) -> Self {
        Self {
            name,
            min_time: Duration::MAX,
            ..Default::default()
        }
    }

    /// Record a timing
    pub fn record(&mut self, duration: Duration, bytes: u64) {
        self.call_count += 1;
        self.total_time += duration;
        self.min_time = self.min_time.min(duration);
        self.max_time = self.max_time.max(duration);
        self.total_bytes += bytes;
    }

    /// Get average time
    pub fn avg_time(&self) -> Duration {
        if self.call_count == 0 {
            Duration::ZERO
        } else {
            self.total_time / self.call_count as u32
        }
    }

    /// Get throughput in GB/s
    pub fn throughput_gbps(&self) -> f64 {
        if self.total_time.as_secs_f64() == 0.0 {
            0.0
        } else {
            (self.total_bytes as f64 / (1024.0 * 1024.0 * 1024.0)) / self.total_time.as_secs_f64()
        }
    }
}

// =============================================================================
// Memory Snapshot
// =============================================================================

/// Memory usage snapshot
#[derive(Debug, Clone)]
pub struct MemorySnapshot {
    /// Timestamp
    pub timestamp: Instant,
    /// Device ID
    pub device_id: i32,
    /// Allocated bytes
    pub allocated_bytes: u64,
    /// Reserved bytes
    pub reserved_bytes: u64,
    /// Peak allocated bytes
    pub peak_allocated_bytes: u64,
    /// Number of allocations
    pub num_allocations: u64,
}

// =============================================================================
// Communication Event
// =============================================================================

/// Communication event for profiling
#[derive(Debug, Clone)]
pub struct CommEvent {
    /// Event name
    pub name: String,
    /// Start time
    pub start_time: Instant,
    /// End time
    pub end_time: Option<Instant>,
    /// Bytes transferred
    pub bytes: u64,
    /// Source rank
    pub src_rank: Option<i32>,
    /// Destination rank
    pub dst_rank: Option<i32>,
    /// Operation type
    pub op_type: String,
}

// =============================================================================
// Profiler
// =============================================================================

/// Main profiler instance
pub struct Profiler {
    /// Configuration
    config: ProfilingConfig,
    /// Whether profiling is enabled
    enabled: AtomicBool,
    /// Next range ID
    next_range_id: AtomicU64,
    /// Active ranges
    active_ranges: RwLock<HashMap<u64, ProfilingRange>>,
    /// Completed ranges
    completed_ranges: RwLock<Vec<ProfilingRange>>,
    /// Operation statistics
    op_stats: RwLock<HashMap<String, OperationStats>>,
    /// Memory snapshots
    memory_snapshots: RwLock<Vec<MemorySnapshot>>,
    /// Communication events
    comm_events: RwLock<Vec<CommEvent>>,
    /// Current range stack (per thread)
    range_stack: RwLock<Vec<u64>>,
    /// Profile start time
    start_time: Instant,
}

impl Profiler {
    /// Create a new profiler
    pub fn new(config: ProfilingConfig) -> Self {
        let enabled = config.level != ProfilingLevel::Off;

        Self {
            config,
            enabled: AtomicBool::new(enabled),
            next_range_id: AtomicU64::new(1),
            active_ranges: RwLock::new(HashMap::new()),
            completed_ranges: RwLock::new(Vec::new()),
            op_stats: RwLock::new(HashMap::new()),
            memory_snapshots: RwLock::new(Vec::new()),
            comm_events: RwLock::new(Vec::new()),
            range_stack: RwLock::new(Vec::new()),
            start_time: Instant::now(),
        }
    }

    /// Check if profiling is enabled
    pub fn is_enabled(&self) -> bool {
        self.enabled.load(Ordering::Relaxed)
    }

    /// Enable profiling
    pub fn enable(&self) {
        self.enabled.store(true, Ordering::Relaxed);
    }

    /// Disable profiling
    pub fn disable(&self) {
        self.enabled.store(false, Ordering::Relaxed);
    }

    /// Start a profiling range
    pub fn start_range(&self, name: &str, category: &str, color: NvtxColor) -> u64 {
        if !self.is_enabled() {
            return 0;
        }

        let id = self.next_range_id.fetch_add(1, Ordering::SeqCst);
        let mut range = ProfilingRange::new(id, name.to_string(), category.to_string(), color);

        // Set parent if auto-nesting
        if self.config.auto_nest {
            let stack = self.range_stack.read();
            range.parent_id = stack.last().copied();
        }

        // Add to active ranges
        {
            let mut active = self.active_ranges.write();
            active.insert(id, range);
        }

        // Push to stack
        {
            let mut stack = self.range_stack.write();
            if stack.len() < self.config.max_depth {
                stack.push(id);
            }
        }

        id
    }

    /// Start a range with RAII guard
    pub fn range_guard(
        self: &Arc<Self>,
        name: &str,
        category: &str,
        color: NvtxColor,
    ) -> RangeGuard {
        let id = self.start_range(name, category, color);
        RangeGuard::new(self.clone(), id)
    }

    /// End a profiling range
    pub fn end_range(&self, range_id: u64) {
        if !self.is_enabled() || range_id == 0 {
            return;
        }

        let range = {
            let mut active = self.active_ranges.write();
            active.remove(&range_id)
        };

        if let Some(mut range) = range {
            range.end();

            // Update operation stats
            if let Some(duration) = range.duration() {
                let mut stats = self.op_stats.write();
                let op_stats = stats
                    .entry(range.name.clone())
                    .or_insert_with(|| OperationStats::new(range.name.clone()));
                op_stats.record(duration, 0);
            }

            // Move to completed
            let mut completed = self.completed_ranges.write();
            completed.push(range);
        }

        // Pop from stack
        {
            let mut stack = self.range_stack.write();
            if stack.last() == Some(&range_id) {
                stack.pop();
            }
        }
    }

    /// Add metadata to an active range
    pub fn add_range_metadata(&self, range_id: u64, key: &str, value: &str) {
        if let Some(range) = self.active_ranges.write().get_mut(&range_id) {
            range.metadata.insert(key.to_string(), value.to_string());
        }
    }

    /// Record a memory snapshot
    pub fn record_memory(
        &self,
        device_id: i32,
        allocated: u64,
        reserved: u64,
        peak: u64,
        num_allocs: u64,
    ) {
        if !self.is_enabled() || !self.config.profile_memory {
            return;
        }

        let snapshot = MemorySnapshot {
            timestamp: Instant::now(),
            device_id,
            allocated_bytes: allocated,
            reserved_bytes: reserved,
            peak_allocated_bytes: peak,
            num_allocations: num_allocs,
        };

        self.memory_snapshots.write().push(snapshot);
    }

    /// Record a communication event
    pub fn record_comm_start(&self, name: &str, bytes: u64, op_type: &str) -> usize {
        if !self.is_enabled() || !self.config.profile_comm {
            return 0;
        }

        let event = CommEvent {
            name: name.to_string(),
            start_time: Instant::now(),
            end_time: None,
            bytes,
            src_rank: None,
            dst_rank: None,
            op_type: op_type.to_string(),
        };

        let mut events = self.comm_events.write();
        let idx = events.len();
        events.push(event);
        idx
    }

    /// End a communication event
    pub fn record_comm_end(&self, event_idx: usize) {
        let mut events = self.comm_events.write();
        if let Some(event) = events.get_mut(event_idx) {
            event.end_time = Some(Instant::now());
        }
    }

    /// Get operation statistics
    pub fn get_op_stats(&self) -> HashMap<String, OperationStats> {
        self.op_stats.read().clone()
    }

    /// Get memory snapshots
    pub fn get_memory_snapshots(&self) -> Vec<MemorySnapshot> {
        self.memory_snapshots.read().clone()
    }

    /// Get communication events
    pub fn get_comm_events(&self) -> Vec<CommEvent> {
        self.comm_events.read().clone()
    }

    /// Get completed ranges
    pub fn get_completed_ranges(&self) -> Vec<ProfilingRange> {
        self.completed_ranges.read().clone()
    }

    /// Clear all profiling data
    pub fn clear(&self) {
        self.active_ranges.write().clear();
        self.completed_ranges.write().clear();
        self.op_stats.write().clear();
        self.memory_snapshots.write().clear();
        self.comm_events.write().clear();
        self.range_stack.write().clear();
    }

    /// Generate summary report
    pub fn summary(&self) -> ProfileSummary {
        let stats = self.op_stats.read();
        let completed = self.completed_ranges.read();
        let memory = self.memory_snapshots.read();
        let comm = self.comm_events.read();

        let total_time = self.start_time.elapsed();

        let mut op_summaries: Vec<_> = stats
            .values()
            .map(|s| OpSummary {
                name: s.name.clone(),
                calls: s.call_count,
                total_ms: s.total_time.as_secs_f64() * 1000.0,
                avg_ms: s.avg_time().as_secs_f64() * 1000.0,
                percent: (s.total_time.as_secs_f64() / total_time.as_secs_f64()) * 100.0,
            })
            .collect();

        op_summaries.sort_by(|a, b| b.total_ms.partial_cmp(&a.total_ms).unwrap());

        let peak_memory = memory
            .iter()
            .map(|s| s.peak_allocated_bytes)
            .max()
            .unwrap_or(0);

        let total_comm_bytes: u64 = comm.iter().map(|e| e.bytes).sum();
        let total_comm_time: Duration = comm
            .iter()
            .filter_map(|e| e.end_time.map(|end| end.duration_since(e.start_time)))
            .sum();

        ProfileSummary {
            total_time,
            num_ranges: completed.len(),
            num_ops: stats.len(),
            op_summaries,
            peak_memory_bytes: peak_memory,
            total_comm_bytes,
            total_comm_time,
        }
    }

    /// Export to Chrome trace format
    pub fn export_chrome_trace(&self) -> String {
        let completed = self.completed_ranges.read();
        let base_time = self.start_time;

        let mut events = Vec::new();

        for range in completed.iter() {
            let start_us = range.start_time.duration_since(base_time).as_micros();
            let dur_us = range.duration().map_or(0, |d| d.as_micros());

            events.push(format!(
                r#"{{"name":"{}","cat":"{}","ph":"X","ts":{},"dur":{},"pid":1,"tid":{}}}"#,
                range.name, range.category, start_us, dur_us, range.thread_id
            ));
        }

        format!("[{}]", events.join(","))
    }
}

// =============================================================================
// Profile Summary
// =============================================================================

/// Summary of profiling results
#[derive(Debug, Clone)]
pub struct ProfileSummary {
    /// Total profiling time
    pub total_time: Duration,
    /// Number of ranges recorded
    pub num_ranges: usize,
    /// Number of unique operations
    pub num_ops: usize,
    /// Operation summaries (sorted by time)
    pub op_summaries: Vec<OpSummary>,
    /// Peak memory usage
    pub peak_memory_bytes: u64,
    /// Total communication bytes
    pub total_comm_bytes: u64,
    /// Total communication time
    pub total_comm_time: Duration,
}

/// Summary for a single operation type
#[derive(Debug, Clone)]
pub struct OpSummary {
    /// Operation name
    pub name: String,
    /// Number of calls
    pub calls: u64,
    /// Total time in milliseconds
    pub total_ms: f64,
    /// Average time in milliseconds
    pub avg_ms: f64,
    /// Percentage of total time
    pub percent: f64,
}

impl ProfileSummary {
    /// Format as string
    pub fn to_string(&self) -> String {
        let mut s = String::new();
        s.push_str("=== Profile Summary ===\n");
        s.push_str(&format!(
            "Total time: {:.2}s\n",
            self.total_time.as_secs_f64()
        ));
        s.push_str(&format!("Ranges: {}\n", self.num_ranges));
        s.push_str(&format!("Unique ops: {}\n", self.num_ops));
        s.push_str(&format!(
            "Peak memory: {:.2} MB\n",
            self.peak_memory_bytes as f64 / 1024.0 / 1024.0
        ));
        s.push_str(&format!(
            "Total comm: {:.2} MB in {:.2}ms\n",
            self.total_comm_bytes as f64 / 1024.0 / 1024.0,
            self.total_comm_time.as_secs_f64() * 1000.0
        ));
        s.push_str("\nTop operations:\n");

        for (i, op) in self.op_summaries.iter().take(10).enumerate() {
            s.push_str(&format!(
                "  {}. {} - {} calls, {:.2}ms total ({:.1}%)\n",
                i + 1,
                op.name,
                op.calls,
                op.total_ms,
                op.percent
            ));
        }

        s
    }
}

// =============================================================================
// Thread-Safe Wrapper
// =============================================================================

/// Thread-safe shared profiler
pub type SharedProfiler = Arc<Profiler>;

/// Create a shared profiler
pub fn shared_profiler(config: ProfilingConfig) -> SharedProfiler {
    Arc::new(Profiler::new(config))
}

/// Global profiler instance
static GLOBAL_PROFILER: std::sync::OnceLock<SharedProfiler> = std::sync::OnceLock::new();

/// Get or initialize global profiler
pub fn global_profiler() -> &'static SharedProfiler {
    GLOBAL_PROFILER.get_or_init(|| shared_profiler(ProfilingConfig::default()))
}

/// Initialize global profiler with config
pub fn init_global_profiler(config: ProfilingConfig) -> Result<()> {
    GLOBAL_PROFILER
        .set(shared_profiler(config))
        .map_err(|_| DistributedError::runtime("Global profiler already initialized"))
}

// =============================================================================
// Convenience Macros (as functions)
// =============================================================================

/// Profile a forward pass operation
pub fn profile_forward<F, T>(profiler: &SharedProfiler, name: &str, f: F) -> T
where
    F: FnOnce() -> T,
{
    let _guard = profiler.range_guard(name, "forward", NvtxColor::FORWARD);
    f()
}

/// Profile a backward pass operation
pub fn profile_backward<F, T>(profiler: &SharedProfiler, name: &str, f: F) -> T
where
    F: FnOnce() -> T,
{
    let _guard = profiler.range_guard(name, "backward", NvtxColor::BACKWARD);
    f()
}

/// Profile an optimizer step
pub fn profile_optimizer<F, T>(profiler: &SharedProfiler, name: &str, f: F) -> T
where
    F: FnOnce() -> T,
{
    let _guard = profiler.range_guard(name, "optimizer", NvtxColor::OPTIMIZER);
    f()
}

/// Profile a communication operation
pub fn profile_comm<F, T>(profiler: &SharedProfiler, name: &str, f: F) -> T
where
    F: FnOnce() -> T,
{
    let _guard = profiler.range_guard(name, "comm", NvtxColor::COMM);
    f()
}

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

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_profiling_config_default() {
        let config = ProfilingConfig::default();
        assert_eq!(config.level, ProfilingLevel::Basic);
        assert_eq!(config.backend, ProfilerBackend::Builtin);
        assert!(config.profile_memory);
    }

    #[test]
    fn test_nvtx_colors() {
        assert_eq!(NvtxColor::RED.0, 0xFF0000);
        assert_eq!(NvtxColor::GREEN.0, 0x00FF00);
        assert_eq!(NvtxColor::FORWARD.0, NvtxColor::GREEN.0);
    }

    #[test]
    fn test_profiler_creation() {
        let config = ProfilingConfig::default();
        let profiler = Profiler::new(config);

        assert!(profiler.is_enabled());
    }

    #[test]
    fn test_profiler_enable_disable() {
        let config = ProfilingConfig::default();
        let profiler = Profiler::new(config);

        profiler.disable();
        assert!(!profiler.is_enabled());

        profiler.enable();
        assert!(profiler.is_enabled());
    }

    #[test]
    fn test_range_start_end() {
        let config = ProfilingConfig::default();
        let profiler = Profiler::new(config);

        let id = profiler.start_range("test_op", "test", NvtxColor::GREEN);
        assert!(id > 0);

        std::thread::sleep(std::time::Duration::from_millis(10));
        profiler.end_range(id);

        let completed = profiler.get_completed_ranges();
        assert_eq!(completed.len(), 1);
        assert_eq!(completed[0].name, "test_op");
        assert!(completed[0].duration().unwrap() >= Duration::from_millis(10));
    }

    #[test]
    fn test_range_guard() {
        let config = ProfilingConfig::default();
        let profiler = Arc::new(Profiler::new(config));

        {
            let _guard = profiler.range_guard("scoped_op", "test", NvtxColor::BLUE);
            std::thread::sleep(std::time::Duration::from_millis(5));
        }

        let completed = profiler.get_completed_ranges();
        assert_eq!(completed.len(), 1);
    }

    #[test]
    fn test_operation_stats() {
        let mut stats = OperationStats::new("test".to_string());

        stats.record(Duration::from_millis(10), 1024);
        stats.record(Duration::from_millis(20), 2048);

        assert_eq!(stats.call_count, 2);
        assert_eq!(stats.total_time, Duration::from_millis(30));
        assert_eq!(stats.avg_time(), Duration::from_millis(15));
        assert_eq!(stats.total_bytes, 3072);
    }

    #[test]
    fn test_memory_snapshot() {
        let config = ProfilingConfig::default();
        let profiler = Profiler::new(config);

        profiler.record_memory(0, 1000, 2000, 1500, 10);

        let snapshots = profiler.get_memory_snapshots();
        assert_eq!(snapshots.len(), 1);
        assert_eq!(snapshots[0].allocated_bytes, 1000);
    }

    #[test]
    fn test_comm_event() {
        let config = ProfilingConfig::default();
        let profiler = Profiler::new(config);

        let idx = profiler.record_comm_start("allreduce", 1024, "AllReduce");
        std::thread::sleep(std::time::Duration::from_millis(5));
        profiler.record_comm_end(idx);

        let events = profiler.get_comm_events();
        assert_eq!(events.len(), 1);
        assert_eq!(events[0].bytes, 1024);
        assert!(events[0].end_time.is_some());
    }

    #[test]
    fn test_profile_summary() {
        let config = ProfilingConfig::default();
        let profiler = Profiler::new(config);

        let id = profiler.start_range("op1", "test", NvtxColor::GREEN);
        profiler.end_range(id);

        let summary = profiler.summary();
        assert_eq!(summary.num_ranges, 1);
        assert!(summary.num_ops >= 1);
    }

    #[test]
    fn test_chrome_trace_export() {
        let config = ProfilingConfig::default();
        let profiler = Profiler::new(config);

        let id = profiler.start_range("test_op", "category", NvtxColor::GREEN);
        profiler.end_range(id);

        let trace = profiler.export_chrome_trace();
        assert!(trace.starts_with('['));
        assert!(trace.ends_with(']'));
        assert!(trace.contains("test_op"));
    }

    #[test]
    fn test_profiler_clear() {
        let config = ProfilingConfig::default();
        let profiler = Profiler::new(config);

        let id = profiler.start_range("test", "test", NvtxColor::GREEN);
        profiler.end_range(id);

        profiler.clear();

        assert!(profiler.get_completed_ranges().is_empty());
        assert!(profiler.get_op_stats().is_empty());
    }

    #[test]
    fn test_shared_profiler() {
        let profiler = shared_profiler(ProfilingConfig::default());

        let id = profiler.start_range("test", "test", NvtxColor::GREEN);
        profiler.end_range(id);

        assert_eq!(profiler.get_completed_ranges().len(), 1);
    }

    #[test]
    fn test_disabled_profiler() {
        let config = ProfilingConfig {
            level: ProfilingLevel::Off,
            ..Default::default()
        };
        let profiler = Profiler::new(config);

        assert!(!profiler.is_enabled());

        let id = profiler.start_range("test", "test", NvtxColor::GREEN);
        assert_eq!(id, 0);
    }

    #[test]
    fn test_profile_helpers() {
        let profiler = shared_profiler(ProfilingConfig::default());

        let result = profile_forward(&profiler, "forward", || 42);

        assert_eq!(result, 42);
        assert_eq!(profiler.get_completed_ranges().len(), 1);
    }
}