// Copyright 2025 TiKV Project Authors. Licensed under Apache-2.0.

use std::{
    collections::VecDeque,
    fs::File,
    io::Write,
    path::Path,
    sync::{Arc, RwLock, mpsc as std_mpsc},
};

use bytes::BufMut;
use file_system::{IoRateLimiter, IoType};
use rfenginepb::{PeerMeta, RaftLogFile};
use tikv_util::{
    sys::SysQuota,
    worker::{Builder, Worker},
};

use crate::{RaftLogOp, Result, RlogCache, RlogHeader, raft_log_file_name, write_batch::PeerBatch};

const WRITE_RLOG_BATCH_SIZE: usize = 256 * 1024;

/// Result of compacting one WAL epoch with intra-epoch parallel
/// sub-compactions.
#[derive(Debug)]
pub(crate) struct EpochCompactionResult {
    pub(crate) metas: Vec<PeerMeta>,
    pub(crate) generated_files: u64,
    pub(crate) cached_files: u64,
}

struct BatchCompactionResult {
    metas: Vec<PeerMeta>,
    generated_files: u64,
    cached_files: u64,
}

pub(crate) struct WrittenRlogFile {
    pub(crate) meta: RaftLogFile,
    pub(crate) is_cached: bool,
}

/// Worker pool used for parallel sub-compaction within a single epoch.
///
/// It deliberately does not own any cross-epoch state. The serialized epoch
/// control plane remains in `CompactWorker`; this pool only executes the
/// per-batch file generation work for the current epoch and returns aggregated
/// results synchronously to the caller.
pub(crate) struct SubCompactionPool {
    worker: Worker,
}

impl SubCompactionPool {
    pub(crate) fn new(thread_count: usize) -> Self {
        let worker = Builder::new("sub-compact-wkr")
            .thread_count(thread_count)
            .thread_count_limits(1, SysQuota::cpu_cores_quota() as usize)
            .create();
        Self { worker }
    }

    pub(crate) fn compact_epoch(
        &self,
        dir: &Path,
        rate_limiter: Option<Arc<IoRateLimiter>>,
        rlog_cache: Arc<RwLock<RlogCache>>,
        rlog_file_size: u32,
        epoch_id: u32,
        peer_batches: Vec<Vec<PeerBatch>>,
    ) -> Result<EpochCompactionResult> {
        if peer_batches.is_empty() {
            return Ok(EpochCompactionResult {
                metas: vec![],
                generated_files: 0,
                cached_files: 0,
            });
        }

        let pool = self.worker.pool();
        // Use a channel to collect the results of the sub-compaction tasks.
        // Although the channel is unbounded, the number of spawned tasks is
        // limited by the number of peer batches, so the memory usage is expected to be
        // low.
        let (tx, rx) = std_mpsc::channel();
        let mut spawned = 0;
        let mut spawn_err = None;

        for batch in peer_batches {
            let tx = tx.clone();
            let dir = dir.to_path_buf();
            let rate_limiter = rate_limiter.clone();
            let rlog_cache = rlog_cache.clone();
            let task = async move {
                let res = compact_peer_batches(
                    &dir,
                    rate_limiter,
                    rlog_cache,
                    rlog_file_size,
                    epoch_id,
                    batch,
                );
                let _ = tx.send(res);
            };
            match pool.spawn(task) {
                Ok(()) => spawned += 1,
                Err(e) => {
                    spawn_err = Some(crate::Error::Other(format!(
                        "failed to spawn sub compaction task: {:?}",
                        e
                    )));
                    break;
                }
            }
        }
        drop(tx);

        let mut first_err = None;
        let mut metas = Vec::new();
        let mut generated_files = 0;
        let mut cached_files = 0;
        for _ in 0..spawned {
            match rx.recv() {
                Ok(Ok(batch_result)) => {
                    metas.extend(batch_result.metas);
                    generated_files += batch_result.generated_files;
                    cached_files += batch_result.cached_files;
                }
                Ok(Err(err)) => {
                    first_err.get_or_insert(err);
                }
                Err(err) => {
                    first_err.get_or_insert(crate::Error::Other(format!(
                        "sub compaction result channel closed unexpectedly: {:?}",
                        err
                    )));
                }
            }
        }

        // Do not return early on the first failed batch. Waiting for every
        // already-spawned batch keeps one epoch fully quiesced before the
        // caller retries from WAL, so we do not overlap retried writes with
        // stale outputs from the failed round.
        if let Some(err) = spawn_err.or(first_err) {
            return Err(err);
        }

        Ok(EpochCompactionResult {
            metas,
            generated_files,
            cached_files,
        })
    }
}

fn compact_peer_batches(
    dir: &Path,
    rate_limiter: Option<Arc<IoRateLimiter>>,
    rlog_cache: Arc<RwLock<RlogCache>>,
    rlog_file_size: u32,
    epoch_id: u32,
    peer_batches: Vec<PeerBatch>,
) -> Result<BatchCompactionResult> {
    let mut generated_files = 0;
    let mut cached_files = 0;
    let mut metas = Vec::with_capacity(peer_batches.len());

    for peer_batch in peer_batches {
        let mut peer_meta_pb = PeerMeta::default();
        peer_meta_pb.set_peer_id(peer_batch.peer_id);
        peer_meta_pb.set_region_id(peer_batch.meta.region_id);
        peer_meta_pb.set_truncated_index(peer_batch.truncated_idx);
        if let Some(keyspace_id) = peer_batch.get_keyspace_id() {
            peer_meta_pb.set_keyspace_id(keyspace_id);
        }
        for (k, v) in &peer_batch.meta.states {
            let mut state_pb = rfenginepb::PeerState::new();
            state_pb.set_key(k.to_vec());
            state_pb.set_value(v.to_vec());
            peer_meta_pb.mut_states().push(state_pb);
        }
        if !peer_batch.raft_logs.is_empty() {
            let files = write_raft_log_files(
                dir,
                rate_limiter.clone(),
                rlog_cache.clone(),
                rlog_file_size,
                peer_batch,
                epoch_id,
            )?;
            for written in files {
                peer_meta_pb.mut_files().push(written.meta);
                generated_files += 1;
                cached_files += written.is_cached as u64;
            }
        }
        metas.push(peer_meta_pb);
    }

    Ok(BatchCompactionResult {
        metas,
        generated_files,
        cached_files,
    })
}

fn write_raft_log_file(
    dir: &Path,
    rate_limiter: Option<Arc<IoRateLimiter>>,
    rlog_cache: Arc<RwLock<RlogCache>>,
    peer_id: u64,
    epoch_id: u32,
    raft_logs: &VecDeque<RaftLogOp>,
    start_idx: usize,
    count: usize,
    encoded_payload_len: usize,
) -> Result<WrittenRlogFile> {
    let first = raft_logs[start_idx].index;
    let last = raft_logs[start_idx + count - 1].index;
    let last_term = raft_logs[start_idx + count - 1].term;
    let filename = raft_log_file_name(dir, peer_id, first, last);

    let mut buf = Vec::with_capacity(
        RlogHeader::len() + count * std::mem::size_of::<u32>() + encoded_payload_len,
    );
    let header = RlogHeader::new(count as u32);
    header.encode_to(&mut buf);

    // Write index first to make the file addressable.
    let mut log_end_off = 0;
    for rlog in raft_logs.iter().skip(start_idx).take(count) {
        log_end_off += rlog.encoded_len() as u32 + 4 /* checksum */;
        buf.put_u32_le(log_end_off);
    }

    // Payload + checksums
    for rlog in raft_logs.iter().skip(start_idx).take(count) {
        let origin_len = buf.len();
        rlog.encode_to(&mut buf);
        let checksum = crc32c::crc32c(&buf[origin_len..]);
        buf.put_u32_le(checksum);
    }
    let mut file = File::create(filename)?;
    if let Some(rate_limiter) = &rate_limiter {
        let synced_size = file_system::write_all_with_rate_limiter(
            &mut file,
            &buf,
            rate_limiter,
            IoType::Compaction,
            WRITE_RLOG_BATCH_SIZE,
        )?;
        let remained_size = buf.len() - synced_size;
        if remained_size > 0 {
            // `write_all_with_rate_limiter(...)` has already persisted every
            // full `WRITE_RLOG_BATCH_SIZE` chunk. Only the final short tail may
            // remain unsynced, so finish durability here and let the caller
            // keep a simple "write file -> return durable metadata" contract.
            file.sync_data()?;
        }
    } else {
        file.write_all(&buf)?;
        file.sync_data()?;
    }

    let meta = RaftLogFile {
        first_index: first,
        last_index: last,
        last_term,
        epoch_id,
        ..Default::default()
    };

    let is_cached = rlog_cache.write().unwrap().add_rlog(peer_id, &meta, &buf);
    Ok(WrittenRlogFile { meta, is_cached })
}

// rfenginepb::RaftLogFile format: RlogHeader + [endoffset] +
// [RaftLogOp + checksum]
pub(crate) fn write_raft_log_files(
    dir: &Path,
    rate_limiter: Option<Arc<IoRateLimiter>>,
    rlog_cache: Arc<RwLock<RlogCache>>,
    rlog_file_size: u32, // Max size is limited to 4GiB
    peer_batch: PeerBatch,
    epoch_id: u32,
) -> Result<Vec<WrittenRlogFile>> {
    let PeerBatch {
        peer_id, raft_logs, ..
    } = peer_batch;
    let mut res = Vec::new();
    let mut i = 0;
    // Split raft logs into multiple files if the raft log offset
    // exceeds size limit.
    while i < raft_logs.len() {
        let mut count = 0;
        let mut log_end_off: u32 = 0;
        while i + count < raft_logs.len() {
            let op = &raft_logs[i + count];
            let size = (op.encoded_len() as u32) + 4 /* checksum */;
            if let Some(new_off) = log_end_off.checked_add(size) {
                if new_off <= rlog_file_size {
                    log_end_off = new_off;
                    count += 1;
                    continue;
                }
            }
            break;
        }
        if count == 0 {
            // A legal Raft entry may exceed the target `rlog_file_size`.
            // Write it alone in an oversized rlog file.
            count = 1;
        }
        res.push(write_raft_log_file(
            dir,
            rate_limiter.clone(),
            rlog_cache.clone(),
            peer_id,
            epoch_id,
            &raft_logs,
            i,
            count,
            log_end_off as usize,
        )?);
        i += count;
    }
    Ok(res)
}

#[cfg(test)]
mod tests {
    use raft_proto::eraftpb::EntryType;
    use tempfile::tempdir;

    use super::*;
    use crate::test_util::{init_logger, new_raft_entry};

    #[test]
    fn test_compact_epoch_with_empty_batches() {
        let dir = tempdir().unwrap();
        let rate_limiter = Arc::new(IoRateLimiter::new_for_test());
        let rlog_cache = Arc::new(RwLock::new(RlogCache::new(0, 0)));
        let pool = SubCompactionPool::new(1);

        let result = pool
            .compact_epoch(
                dir.path(),
                Some(rate_limiter),
                rlog_cache,
                1024 * 1024,
                42,
                vec![],
            )
            .unwrap();

        assert!(result.metas.is_empty());
        assert_eq!(result.generated_files, 0);
        assert_eq!(result.cached_files, 0);
    }

    #[test]
    fn test_compact_epoch_aggregates_parallel_batches() {
        init_logger();
        let dir = tempdir().unwrap();
        let rate_limiter = Arc::new(IoRateLimiter::new_for_test());
        let rlog_cache = Arc::new(RwLock::new(RlogCache::new(0, 0)));
        let pool = SubCompactionPool::new(2);

        let mut pb1 = PeerBatch::new(1, 11, 1);
        let mut pb2 = PeerBatch::new(2, 22, 2);
        for i in 1..=10 {
            let log = RaftLogOp::new(&new_raft_entry(EntryType::EntryNormal, 1, i, b"data", 0));
            pb1.append_raft_log(log.clone());
            pb2.append_raft_log(log);
        }

        let result = pool
            .compact_epoch(
                dir.path(),
                Some(rate_limiter),
                rlog_cache,
                1024 * 1024,
                100,
                vec![vec![pb1], vec![pb2]],
            )
            .unwrap();

        assert_eq!(result.metas.len(), 2);
        assert_eq!(result.generated_files, 2);
        assert_eq!(result.cached_files, 0);
    }

    #[test]
    fn test_compact_epoch_returns_error_when_output_dir_is_missing() {
        init_logger();
        let dir = tempdir().unwrap();
        let missing_dir = dir.path().join("missing");
        let rate_limiter = Arc::new(IoRateLimiter::new_for_test());
        let rlog_cache = Arc::new(RwLock::new(RlogCache::new(0, 0)));
        let pool = SubCompactionPool::new(2);

        let mut pb1 = PeerBatch::new(1, 11, 1);
        let mut pb2 = PeerBatch::new(2, 22, 2);
        for i in 1..=2 {
            let log = RaftLogOp::new(&new_raft_entry(EntryType::EntryNormal, 1, i, b"ok", 0));
            pb1.append_raft_log(log.clone());
            pb2.append_raft_log(log);
        }

        let err = pool
            .compact_epoch(
                &missing_dir,
                Some(rate_limiter),
                rlog_cache,
                16 * 1024,
                7,
                vec![vec![pb1], vec![pb2]],
            )
            .unwrap_err();

        assert!(matches!(err, crate::Error::Other(_) | crate::Error::Io(_)));
    }
}