//! Timestamp Oracle (TSO) for FlareDB
//!
//! The TSO provides globally unique, monotonically increasing timestamps
//! used for transaction ordering and MVCC in the distributed KV store.
//!
//! # Timestamp Format
//!
//! Each timestamp is a 64-bit integer with the following structure:
//! - **Upper 48 bits**: Physical time in milliseconds since Unix epoch
//! - **Lower 16 bits**: Logical counter (0-65535)
//!
//! ```text
//! ┌──────────────────────────────────────────────────────────┬──────────────────┐
//! │              Physical Time (48 bits)                     │  Logical (16 bits)│
//! │              milliseconds since epoch                    │  counter 0-65535  │
//! └──────────────────────────────────────────────────────────┴──────────────────┘
//! ```
//!
//! This format allows:
//! - Up to 65536 timestamps per millisecond
//! - Time-based ordering (most significant bits are physical time)
//! - Easy extraction of physical time: `ts >> 16`
//! - Easy extraction of logical counter: `ts & 0xFFFF`
//!
//! # Thread Safety
//!
//! The TSO is fully thread-safe and uses atomic operations to ensure
//! monotonicity even under high concurrency.
//!
//! # Usage
//!
//! ```rust
//! use flaredb_pd::tso::TsoOracle;
//!
//! let tso = TsoOracle::new();
//!
//! // Get a single timestamp
//! let ts = tso.get_timestamp(1);
//!
//! // Extract components
//! let physical = TsoOracle::physical_time(ts);
//! let logical = TsoOracle::logical_counter(ts);
//!
//! // Get a batch of 100 timestamps (returns the first one)
//! let first_ts = tso.get_timestamp(100);
//! // The allocated range is [first_ts, first_ts + 99]
//! ```

use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};

/// Timestamp Oracle providing globally unique timestamps.
///
/// The oracle maintains monotonicity by tracking the last issued timestamp
/// and advancing the logical counter or waiting for the next millisecond
/// when the counter would overflow.
pub struct TsoOracle {
    /// Combined physical (upper 48 bits) and logical (lower 16 bits) timestamp.
    /// Uses atomic operations to ensure thread safety.
    last_tso: AtomicU64,
}

impl Default for TsoOracle {
    fn default() -> Self {
        Self::new()
    }
}

impl TsoOracle {
    /// The number of bits used for the logical counter.
    pub const LOGICAL_BITS: u32 = 16;

    /// Maximum value for the logical counter (65535).
    pub const MAX_LOGICAL: u64 = 0xFFFF;

    /// Create a new TSO oracle.
    ///
    /// The oracle starts with no previously issued timestamps.
    pub fn new() -> Self {
        Self {
            last_tso: AtomicU64::new(0),
        }
    }

    /// Allocate `count` timestamps and return the first one.
    ///
    /// The caller is responsible for using timestamps in the range
    /// `[returned_ts, returned_ts + count - 1]`.
    ///
    /// # Arguments
    ///
    /// * `count` - Number of timestamps to allocate (1-65535)
    ///
    /// # Returns
    ///
    /// The first timestamp in the allocated batch.
    ///
    /// # Panics
    ///
    /// This function will spin-wait if the logical counter would overflow,
    /// waiting for the next millisecond. In extreme cases (>65536 requests
    /// per millisecond), this may cause delays.
    pub fn get_timestamp(&self, count: u32) -> u64 {
        debug_assert!(count > 0 && count <= Self::MAX_LOGICAL as u32);

        loop {
            let now = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_millis() as u64;

            let current_tso = self.last_tso.load(Ordering::Acquire);
            let current_physical = current_tso >> Self::LOGICAL_BITS;
            let current_logical = current_tso & Self::MAX_LOGICAL;

            if now > current_physical {
                // Time has advanced, reset logical counter
                let start_logical = 0;
                let end_logical = count as u64 - 1;
                let new_val = (now << Self::LOGICAL_BITS) | end_logical;

                if self
                    .last_tso
                    .compare_exchange(current_tso, new_val, Ordering::AcqRel, Ordering::Acquire)
                    .is_ok()
                {
                    return (now << Self::LOGICAL_BITS) | start_logical;
                }
            } else {
                // Same millisecond, advance logical counter
                if current_logical + (count as u64) > Self::MAX_LOGICAL {
                    // Would overflow, wait for next millisecond
                    std::thread::yield_now();
                    continue;
                }

                let start_logical = current_logical + 1;
                let end_logical = current_logical + (count as u64);
                let new_val = (current_physical << Self::LOGICAL_BITS) | end_logical;

                if self
                    .last_tso
                    .compare_exchange(current_tso, new_val, Ordering::AcqRel, Ordering::Acquire)
                    .is_ok()
                {
                    return (current_physical << Self::LOGICAL_BITS) | start_logical;
                }
            }
        }
    }

    /// Extract the physical time (milliseconds since epoch) from a timestamp.
    ///
    /// # Example
    ///
    /// ```rust
    /// use flaredb_pd::tso::TsoOracle;
    ///
    /// let tso = TsoOracle::new();
    /// let ts = tso.get_timestamp(1);
    /// let physical = TsoOracle::physical_time(ts);
    /// // physical is approximately the current time in milliseconds
    /// ```
    #[inline]
    pub fn physical_time(ts: u64) -> u64 {
        ts >> Self::LOGICAL_BITS
    }

    /// Extract the logical counter from a timestamp.
    ///
    /// # Example
    ///
    /// ```rust
    /// use flaredb_pd::tso::TsoOracle;
    ///
    /// let tso = TsoOracle::new();
    /// let ts = tso.get_timestamp(1);
    /// let logical = TsoOracle::logical_counter(ts);
    /// // logical is the sequence number within the millisecond (0-65535)
    /// ```
    #[inline]
    pub fn logical_counter(ts: u64) -> u16 {
        (ts & Self::MAX_LOGICAL) as u16
    }

    /// Compose a timestamp from physical time and logical counter.
    ///
    /// # Arguments
    ///
    /// * `physical` - Physical time in milliseconds since epoch
    /// * `logical` - Logical counter (0-65535)
    ///
    /// # Returns
    ///
    /// The composed timestamp.
    #[inline]
    pub fn compose(physical: u64, logical: u16) -> u64 {
        (physical << Self::LOGICAL_BITS) | (logical as u64)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashSet;
    use std::sync::Arc;
    use std::thread;

    #[test]
    fn test_monotonic_tso() {
        let tso = Arc::new(TsoOracle::new());
        let mut threads = vec![];

        for _ in 0..10 {
            let tso = tso.clone();
            threads.push(thread::spawn(move || {
                let mut last = 0;
                for _ in 0..1000 {
                    let ts = tso.get_timestamp(1);
                    assert!(ts > last, "TSO not monotonic: {} vs {}", ts, last);
                    last = ts;
                }
            }));
        }

        for t in threads {
            t.join().unwrap();
        }
    }

    #[test]
    fn test_batch_allocation() {
        let tso = TsoOracle::new();

        // Allocate batch of 10
        let first = tso.get_timestamp(10);

        // Next allocation should be after the batch
        let next = tso.get_timestamp(1);
        assert!(
            next >= first + 10,
            "next={} should be >= first+10={}",
            next,
            first + 10
        );
    }

    #[test]
    fn test_batch_uniqueness() {
        let tso = Arc::new(TsoOracle::new());
        let all_timestamps = Arc::new(std::sync::Mutex::new(HashSet::new()));
        let mut handles = vec![];

        // Spawn threads each allocating batches
        for _ in 0..10 {
            let tso = tso.clone();
            let timestamps = all_timestamps.clone();
            handles.push(thread::spawn(move || {
                for _ in 0..100 {
                    let batch_size = 10u32;
                    let first = tso.get_timestamp(batch_size);
                    let mut ts_set = timestamps.lock().unwrap();
                    for i in 0..batch_size as u64 {
                        let ts = first + i;
                        assert!(ts_set.insert(ts), "Duplicate timestamp detected: {}", ts);
                    }
                }
            }));
        }

        for h in handles {
            h.join().unwrap();
        }

        // Should have 10 threads * 100 batches * 10 per batch = 10000 unique timestamps
        assert_eq!(all_timestamps.lock().unwrap().len(), 10000);
    }

    #[test]
    fn test_physical_logical_extraction() {
        let tso = TsoOracle::new();
        let ts = tso.get_timestamp(1);

        let physical = TsoOracle::physical_time(ts);
        let logical = TsoOracle::logical_counter(ts);

        // Reconstruct and verify
        assert_eq!(TsoOracle::compose(physical, logical), ts);

        // Physical time should be recent (within 1 second of now)
        let now_ms = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_millis() as u64;
        assert!(
            (now_ms - physical) < 1000,
            "physical={} should be within 1s of now={}",
            physical,
            now_ms
        );
    }

    #[test]
    fn test_compose_roundtrip() {
        // Test that compose and extract are inverses
        let physical = 1234567890123u64;
        let logical = 42u16;

        let ts = TsoOracle::compose(physical, logical);
        assert_eq!(TsoOracle::physical_time(ts), physical);
        assert_eq!(TsoOracle::logical_counter(ts), logical);
    }

    #[test]
    fn test_default_trait() {
        let tso1 = TsoOracle::default();
        let tso2 = TsoOracle::new();

        // Both should produce valid timestamps
        let ts1 = tso1.get_timestamp(1);
        let ts2 = tso2.get_timestamp(1);
        assert!(ts1 > 0);
        assert!(ts2 > 0);
    }

    #[test]
    fn test_constants() {
        assert_eq!(TsoOracle::LOGICAL_BITS, 16);
        assert_eq!(TsoOracle::MAX_LOGICAL, 0xFFFF);
        assert_eq!(TsoOracle::MAX_LOGICAL, 65535);
    }
}
