use pd_proto_tso::{MAX_LOGICAL, Timestamp};

use crate::error::ServiceError;

/// Oracle runtime configuration.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct OracleConfig {
    /// Persisted save-window size in milliseconds.
    pub save_interval_ms: i64,
    /// Background physical-time update interval in milliseconds.
    pub update_physical_interval_ms: i64,
    /// Maximum tolerated backward jump before reset policy applies.
    pub max_reset_ts_gap_ms: i64,
}

impl Default for OracleConfig {
    fn default() -> Self {
        Self {
            save_interval_ms: 5_000,
            update_physical_interval_ms: 50,
            max_reset_ts_gap_ms: 24 * 60 * 60 * 1_000,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum UpdatePurpose {
    Interval,
    Overflow,
}

/// Timestamp oracle implementing monotonic physical/logical allocation.
#[derive(Debug, Clone)]
pub struct TimestampOracle {
    cfg: OracleConfig,
    initialized: bool,
    physical_ms: i64,
    logical: i64,
    last_saved_ms: i64,
}

impl TimestampOracle {
    /// Creates an uninitialized timestamp oracle with one runtime config snapshot.
    #[must_use]
    pub fn new(cfg: OracleConfig) -> Self {
        Self {
            cfg,
            initialized: false,
            physical_ms: 0,
            logical: 0,
            last_saved_ms: 0,
        }
    }

    /// Returns whether persisted state has been synced into this oracle.
    #[must_use]
    pub fn is_initialized(&self) -> bool {
        self.initialized
    }

    /// Returns current in-memory timestamp components.
    #[must_use]
    pub fn current_timestamp(&self) -> Timestamp {
        Timestamp {
            physical: self.physical_ms,
            logical: self.logical,
            suffix_bits: 0,
        }
    }

    /// Returns current persisted save-window upper bound in milliseconds.
    #[must_use]
    pub fn save_window_end_ms(&self) -> i64 {
        self.last_saved_ms
    }

    /// Applies updated runtime oracle configuration for hot-reload flows.
    pub fn apply_config(&mut self, cfg: OracleConfig) {
        self.cfg = cfg;
    }

    /// Clears in-memory oracle state after lease loss/step-down.
    pub fn reset(&mut self) {
        self.initialized = false;
        self.physical_ms = 0;
        self.logical = 0;
        self.last_saved_ms = 0;
    }

    /// Syncs in-memory oracle state from persisted window and writes a fresh save window.
    pub fn sync_timestamp<F>(
        &mut self,
        now_ms: i64,
        loaded_timestamp_nanos: Option<i64>,
        mut persist_save_window: F,
    ) -> Result<(), ServiceError>
    where
        F: FnMut(i64) -> Result<(), ServiceError>,
    {
        let loaded_ms = loaded_timestamp_nanos
            .map(|nanos| nanos / 1_000_000)
            .unwrap_or(0);

        if self.initialized
            && loaded_ms != 0
            && self.last_saved_ms != 0
            && loaded_ms == self.last_saved_ms
        {
            return Ok(());
        }

        let mut next_ms = now_ms;
        if loaded_ms != 0 && now_ms - loaded_ms < 1 {
            next_ms = loaded_ms + 1;
        }

        let save_ms = next_ms + self.cfg.save_interval_ms;
        persist_save_window(save_ms)?;

        self.initialized = true;
        self.physical_ms = next_ms;
        self.logical = 0;
        self.last_saved_ms = save_ms;
        Ok(())
    }

    /// Interval update path from allocator background loop.
    pub fn update_interval<F>(
        &mut self,
        now_ms: i64,
        persist_save_window: F,
    ) -> Result<(), ServiceError>
    where
        F: FnMut(i64) -> Result<(), ServiceError>,
    {
        let _ = self.update_timestamp(now_ms, UpdatePurpose::Interval, persist_save_window)?;
        Ok(())
    }

    /// Generates a timestamp range with `count` logical slots.
    pub fn generate_tso<F>(
        &mut self,
        now_ms: i64,
        count: u32,
        serving: bool,
        mut persist_save_window: F,
    ) -> Result<Timestamp, ServiceError>
    where
        F: FnMut(i64) -> Result<(), ServiceError>,
    {
        if count == 0 {
            return Err(ServiceError::InvalidCount);
        }

        for attempt in 0..10 {
            if !self.initialized {
                return Err(ServiceError::NotStarted);
            }
            if !serving {
                return Err(ServiceError::NotLeader);
            }

            self.logical += i64::from(count);
            if self.logical >= MAX_LOGICAL {
                let retry_now =
                    now_ms + (attempt as i64 + 1) * self.cfg.update_physical_interval_ms;
                let _ = self.update_timestamp(
                    retry_now,
                    UpdatePurpose::Overflow,
                    &mut persist_save_window,
                )?;
                continue;
            }

            return Ok(Timestamp {
                physical: self.physical_ms,
                logical: self.logical,
                suffix_bits: 0,
            });
        }

        Err(ServiceError::LogicalOverflow)
    }

    /// Resets in-memory timestamp to an explicitly provided value.
    pub fn set_tso<F>(
        &mut self,
        tso: u64,
        ignore_smaller: bool,
        skip_upper_bound_check: bool,
        mut persist_save_window: F,
    ) -> Result<(), ServiceError>
    where
        F: FnMut(i64) -> Result<(), ServiceError>,
    {
        if !self.initialized {
            return Err(ServiceError::NotStarted);
        }

        let parsed = Timestamp::parse(tso);
        let next_physical = parsed.physical;
        let next_logical = parsed.logical;

        let physical_diff = next_physical - self.physical_ms;
        let logical_diff = next_logical - self.logical;

        if physical_diff < 0 {
            if ignore_smaller {
                return Ok(());
            }
            return Err(ServiceError::Storage(
                "the specified ts is smaller than now".to_string(),
            ));
        }
        if physical_diff == 0 && logical_diff <= 0 {
            if ignore_smaller {
                return Ok(());
            }
            return Err(ServiceError::Storage(
                "the specified counter is smaller than now".to_string(),
            ));
        }
        if !skip_upper_bound_check && physical_diff >= self.cfg.max_reset_ts_gap_ms {
            return Err(ServiceError::Storage(
                "the specified ts is too larger than now".to_string(),
            ));
        }

        if self.last_saved_ms - next_physical <= 1 {
            let save_ms = next_physical + self.cfg.save_interval_ms;
            persist_save_window(save_ms)?;
            self.last_saved_ms = save_ms;
        }

        self.physical_ms = next_physical;
        self.logical = next_logical;
        Ok(())
    }

    fn update_timestamp<F>(
        &mut self,
        now_ms: i64,
        purpose: UpdatePurpose,
        mut persist_save_window: F,
    ) -> Result<bool, ServiceError>
    where
        F: FnMut(i64) -> Result<(), ServiceError>,
    {
        if !self.initialized {
            return Err(ServiceError::NotStarted);
        }

        let jet_lag_ms = now_ms - self.physical_ms;
        let next_ms = if jet_lag_ms > 1 {
            now_ms
        } else if self.logical > MAX_LOGICAL / 2 {
            self.physical_ms + 1
        } else {
            return Ok(false);
        };

        if self.last_saved_ms - next_ms <= 1 {
            if purpose == UpdatePurpose::Overflow {
                return Ok(true);
            }
            let save_ms = next_ms + self.cfg.save_interval_ms;
            persist_save_window(save_ms)?;
            self.last_saved_ms = save_ms;
        }

        if next_ms > self.physical_ms {
            self.physical_ms = next_ms;
            self.logical = 0;
        }

        Ok(self.logical >= MAX_LOGICAL)
    }
}

#[cfg(test)]
mod tests {
    use super::{OracleConfig, TimestampOracle};

    #[test]
    fn sync_respects_persisted_guard_and_saves_window() {
        let mut oracle = TimestampOracle::new(OracleConfig::default());
        let mut saved = Vec::new();

        oracle
            .sync_timestamp(1_000, Some(1_000_000_000), |save_ms| {
                saved.push(save_ms);
                Ok(())
            })
            .expect("sync");

        // persisted=1000ms, next must be >= persisted+1ms when now is not ahead.
        assert_eq!(oracle.current_timestamp().physical, 1_001);
        assert_eq!(saved, vec![6_001]);
    }

    #[test]
    fn overflow_path_advances_physical() {
        let mut oracle = TimestampOracle::new(OracleConfig {
            save_interval_ms: 5_000,
            update_physical_interval_ms: 1,
            max_reset_ts_gap_ms: 100_000,
        });
        oracle
            .sync_timestamp(10_000, None, |_save_ms| Ok(()))
            .expect("sync");

        // Push logical near overflow boundary.
        oracle
            .set_tso(
                ((10_000_u64) << 18) | ((1_u64 << 18) - 2),
                false,
                true,
                |_save_ms| Ok(()),
            )
            .expect("set tso");

        let ts = oracle
            .generate_tso(10_000, 2, true, |_save_ms| Ok(()))
            .expect("generate");
        assert!(ts.physical >= 10_001);
        assert_eq!(ts.logical, 2);
    }

    #[test]
    fn zero_count_is_rejected() {
        let mut oracle = TimestampOracle::new(OracleConfig::default());
        oracle
            .sync_timestamp(100, None, |_save_ms| Ok(()))
            .expect("sync");

        let err = oracle
            .generate_tso(100, 0, true, |_save_ms| Ok(()))
            .expect_err("count=0 should fail");
        assert_eq!(err.to_string(), "tso count should be positive");
    }
}