use std::{
    collections::HashMap,
    fmt::Debug,
    sync::{Arc, atomic::AtomicU64},
    time::SystemTime,
};

use derivative::Derivative;
use portable_atomic::AtomicU128;
use tici_common::cdc_file_name::{CDCFileName, older_cdc_file};

use crate::{ShardInfo, meta::writer::compaction::*};

#[derive(Clone, Derivative, serde::Serialize, serde::Deserialize, Default)]
#[derivative(PartialEq)]
pub struct ShardBuildProgress {
    // Reported from writer, current consume location.
    pub cdc_s3_last_file: String,
    // Locally tracked cdc file that has been persisted to meta store.
    #[serde(skip)]
    #[derivative(PartialEq = "ignore")]
    pub persisted_cdc_s3_last_file: String,
    // Multiple compactions can be in progress at the same time.
    // Stored to avoid repeatedly creating compaction work over the same range.
    // NOTE Persist if we want the compaction is recoverable.
    #[serde(skip)]
    #[derivative(PartialEq = "ignore")]
    pub compaction_state: HashMap<CompactionStateKey, CompactionState>,
}

impl ShardBuildProgress {
    pub fn collect_compaction_states(&self) -> Vec<CompactionState> {
        self.compaction_state.values().cloned().collect::<Vec<_>>()
    }

    pub fn mark_cdc_last_file_persisted(&mut self) {
        self.persisted_cdc_s3_last_file = self.cdc_s3_last_file.clone();
    }
}

pub fn format_cdc_last_file(f: &str) -> String {
    if let Ok(o) = CDCFileName::parse(f) {
        format!("{}|{}", o.table_version(), o.sequence())
    } else {
        f.to_string()
    }
}

impl Debug for ShardBuildProgress {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SBP")
            .field_with("cdc_last_file", |f| {
                f.write_fmt(format_args!(
                    "{}",
                    format_cdc_last_file(&self.cdc_s3_last_file)
                ))
            })
            .field_with("persisted_cdc_last_file", |f| {
                f.write_fmt(format_args!(
                    "{}",
                    format_cdc_last_file(&self.persisted_cdc_s3_last_file)
                ))
            })
            .field("compaction_state", &self.collect_compaction_states())
            .finish()
    }
}

impl heapsize::HeapSizeOf for ShardBuildProgress {
    fn heap_size_of_children(&self) -> usize {
        self.cdc_s3_last_file.len()
            + self.persisted_cdc_s3_last_file.len()
            + self
                .compaction_state
                .iter()
                .map(|(k, v)| k.heap_size_of_children() + v.heap_size_of_children())
                .sum::<usize>()
    }
}

impl ShardBuildProgress {
    pub fn merge(&mut self, other: &ShardBuildProgress) {
        if !self.cdc_s3_last_file.is_empty() && !other.cdc_s3_last_file.is_empty() {
            let cdc1 = &self.cdc_s3_last_file;
            let cdc2 = &other.cdc_s3_last_file;
            self.cdc_s3_last_file = older_cdc_file(cdc1.as_str(), cdc2.as_str());
        } else {
            self.cdc_s3_last_file = String::new();
        }
    }
}

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

pub struct RateCounter {
    window: Duration,
    records: parking_lot::Mutex<VecDeque<(Instant, u64)>>,
    total: AtomicU64,
}

impl RateCounter {
    pub fn new(window: Duration) -> Self {
        Self {
            window,
            records: parking_lot::Mutex::new(VecDeque::new()),
            total: AtomicU64::new(0),
        }
    }

    pub fn record(&self, bytes: u64) {
        let now = Instant::now();
        self.records.lock().push_back((now, bytes));
        self.total
            .fetch_add(bytes, std::sync::atomic::Ordering::SeqCst);

        let cutoff = now - self.window;
        loop {
            if let Some(&(t, _)) = self.records.lock().front() {
                if t >= cutoff {
                    break;
                }
            } else {
                break;
            }
            if let Some((_, b)) = self.records.lock().pop_front() {
                self.total.fetch_sub(b, std::sync::atomic::Ordering::SeqCst);
            }
        }
    }

    pub fn bps(&self) -> f64 {
        let now = Instant::now();
        let cutoff = now - self.window;
        let mut sum = 0u64;
        let mut earliest = None;
        let mut latest = None;
        for &(t, b) in self.records.lock().iter() {
            if t >= cutoff {
                sum += b;
                if earliest.is_none() {
                    earliest = Some(t);
                }
                latest = Some(t);
            }
        }
        if sum == 0 {
            // At least 2 points produce 1 estimation.
            return 0.0;
        }
        if let Some(latest) = latest {
            if let Some(earliest) = earliest {
                let elapsed = latest.duration_since(earliest).as_secs_f64();
                if elapsed == 0.0 {
                    // All records at the same instant, treat as window
                    sum as f64 / self.window.as_secs_f64()
                } else {
                    sum as f64 / elapsed
                }
            } else {
                0.0
            }
        } else {
            0.0
        }
    }
}
struct QpsCounter {
    window: Duration,
    times: parking_lot::Mutex<VecDeque<Instant>>,
}

impl QpsCounter {
    pub fn new(window: Duration) -> Self {
        Self {
            window,
            times: parking_lot::Mutex::new(VecDeque::new()),
        }
    }

    pub fn record(&self, now: Instant) {
        let mut times = self.times.lock();
        times.push_back(now);

        let cutoff = now - self.window;
        while let Some(&front) = times.front() {
            if front < cutoff {
                times.pop_front();
            } else {
                break;
            }
        }
    }

    pub fn qps(&self) -> f64 {
        let times = self.times.lock();

        let n = times.len();
        if n == 0 {
            0.0
        } else if n == 1 {
            match times.front() {
                Some(s) => {
                    let elapsed = Instant::now().duration_since(*s);
                    1.0 / elapsed.as_secs_f64()
                }
                None => 0.0,
            }
        } else {
            n as f64 / self.window.as_secs_f64()
        }
    }
}

#[derive(Clone)]
pub struct ShardStats {
    compaction_start: Arc<AtomicU64>,
    compaction_finish: Arc<AtomicU64>,
    compaction_abort: Arc<AtomicU64>,
    size: Arc<AtomicU64>,
    fragment_count: Arc<AtomicU64>,
    write_qps: Arc<QpsCounter>,
    pub last_heartbeat_ts: Arc<AtomicU128>,
    append_bps: Arc<RateCounter>,
}

impl Default for ShardStats {
    fn default() -> Self {
        let ops = QpsCounter::new(Duration::from_secs(60 * 5));
        ops.record(Instant::now());
        Self {
            compaction_start: Arc::new(AtomicU64::new(0)),
            compaction_finish: Arc::new(AtomicU64::new(0)),
            compaction_abort: Arc::new(AtomicU64::new(0)),
            size: Arc::new(AtomicU64::new(0)),
            fragment_count: Arc::new(AtomicU64::new(0)),
            write_qps: Arc::new(ops),
            last_heartbeat_ts: Arc::new(AtomicU128::new(
                SystemTime::now()
                    .duration_since(SystemTime::UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_millis(),
            )),
            append_bps: Arc::new(RateCounter::new(std::time::Duration::from_secs(30))),
        }
    }
}

impl Debug for ShardStats {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!(
            "compaction_(start={}|finish={}|abort={}),size={},frags={}",
            self.get_compaction_start(),
            self.get_compaction_finish(),
            self.get_compaction_abort(),
            self.size.load(std::sync::atomic::Ordering::Relaxed),
            self.fragment_count
                .load(std::sync::atomic::Ordering::Relaxed)
        ))
    }
}

impl ShardStats {
    pub fn new(size: u64, fragment_count: u64) -> Self {
        Self {
            size: Arc::new(AtomicU64::new(size)),
            fragment_count: Arc::new(AtomicU64::new(fragment_count)),
            ..Default::default()
        }
    }

    pub fn split(parent: &ShardStats, splits: u64) -> Self {
        Self {
            compaction_start: Arc::new(AtomicU64::new(parent.get_compaction_start())),
            compaction_finish: Arc::new(AtomicU64::new(parent.get_compaction_finish())),
            compaction_abort: Arc::new(AtomicU64::new(parent.get_compaction_abort())),
            size: Arc::new(AtomicU64::new(
                (parent.get_size() as f64 / splits as f64) as u64,
            )),
            fragment_count: Arc::new(AtomicU64::new(
                (parent.get_fragment_count() as f64 / splits as f64) as u64,
            )),
            ..Default::default()
        }
    }

    pub fn get_compaction_start(&self) -> u64 {
        self.compaction_start
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn get_compaction_finish(&self) -> u64 {
        self.compaction_finish
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn get_compaction_abort(&self) -> u64 {
        self.compaction_abort
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn observe_compaction_finish(&self) {
        self.compaction_finish
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    pub fn observe_compaction_abort(&self) {
        self.compaction_abort
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    pub fn observe_compaction_start(&self) {
        self.compaction_start
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    pub fn get_size(&self) -> u64 {
        self.size.load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn get_fragment_count(&self) -> u64 {
        self.fragment_count
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn get_append_bps(&self) -> f64 {
        self.append_bps.bps()
    }

    pub async fn on_write(&self, shard: &ShardInfo) {
        self.size.store(
            shard.get_total_estimated_bytes(),
            std::sync::atomic::Ordering::Relaxed,
        );
        self.fragment_count.store(
            shard.get_total_fragment_count() as u64,
            std::sync::atomic::Ordering::Relaxed,
        );
        self.ping_write();
    }

    pub async fn on_append(&self, shard: &ShardInfo, added: u64) {
        self.size
            .fetch_add(added, std::sync::atomic::Ordering::Relaxed);
        self.fragment_count.store(
            shard.get_total_fragment_count() as u64,
            std::sync::atomic::Ordering::Relaxed,
        );
        self.ping_write();
        self.append_bps.record(added);
    }
}

impl ShardStats {
    pub fn ping_hb(&self) {
        self.last_heartbeat_ts.store(
            SystemTime::now()
                .duration_since(SystemTime::UNIX_EPOCH)
                .unwrap_or_default()
                .as_millis(),
            std::sync::atomic::Ordering::Relaxed,
        );
    }

    pub fn ping_write(&self) {
        self.write_qps.record(Instant::now());
    }

    pub async fn get_write_qps(&self) -> f64 {
        self.write_qps.qps()
    }
}

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

    #[test]
    fn test_shard_build_progress_skip_persisted_cdc_serde() {
        let mut progress = ShardBuildProgress {
            cdc_s3_last_file:
                "tici_default_prefix/cdc/10313/461869219056451591/CDC00000000000000000002.json"
                    .to_string(),
            persisted_cdc_s3_last_file:
                "tici_default_prefix/cdc/10313/461869219056451591/CDC00000000000000000001.json"
                    .to_string(),
            ..Default::default()
        };
        let json = serde_json::to_string(&progress).unwrap();
        assert!(json.contains("cdc_s3_last_file"));
        assert!(!json.contains("persisted_cdc_s3_last_file"));

        let restored: ShardBuildProgress = serde_json::from_str(&json).unwrap();
        assert_eq!(restored.cdc_s3_last_file, progress.cdc_s3_last_file);
        assert!(restored.persisted_cdc_s3_last_file.is_empty());

        progress.mark_cdc_last_file_persisted();
        assert_eq!(
            progress.persisted_cdc_s3_last_file,
            progress.cdc_s3_last_file
        );
    }

    #[test]
    fn test_qps_counter() {
        let counter = QpsCounter::new(Duration::from_secs(10));
        let base = Instant::now();

        counter.record(base);
        counter.record(base + Duration::from_secs(3));
        counter.record(base + Duration::from_secs(7));

        // 3 req in the last 10 secs, QPS = 0.3
        assert!((counter.qps() - 0.3).abs() < 1e-6);

        counter.record(base + Duration::from_secs(15));
        // Only retain the req from 5s and 15s.
        assert!((counter.qps() - 0.2).abs() < 1e-6);

        let counter = QpsCounter::new(Duration::from_secs(10));
        assert_eq!(counter.qps(), 0.0);
    }

    #[test]
    fn test_rate_counter_basic() {
        let window = Duration::from_secs(1);
        let rc = RateCounter::new(window);

        assert_eq!(rc.bps(), 0.0);

        rc.record(100);
        let bps = rc.bps();
        assert!(bps > 0.0);

        std::thread::sleep(Duration::from_millis(100));
        rc.record(200);
        let bps2 = rc.bps();
        assert!(bps2 > bps);

        // Wait for window to expire
        std::thread::sleep(window);
        std::thread::sleep(Duration::from_millis(123));
        rc.record(50);
        let bps3 = rc.bps();
        assert!(bps3 > 0.0);

        std::thread::sleep(window);
        assert_eq!(rc.bps(), 0.0);
    }
}