// Copyright 2025 PingCAP, Inc.

use std::{
    collections::{BTreeSet, HashMap, VecDeque},
    path::PathBuf,
    sync::{Arc, RwLock},
};

use arc_swap::ArcSwap;
use bytes::Bytes;
use kvproto::raft_serverpb;
use protobuf::Message;
use raft_proto::eraftpb::Entry;
use tikv_util::{debug, defer, info, time::Instant};

use crate::{
    AsyncFetchInfo, Error, PeerStats, RAFT_STATE_KEY_BYTE, REGION_META_KEY_BYTE,
    TRUNCATE_ALL_INDEX, WriteBatch,
    bg_service_worker::fetch_entries_from_rlog,
    engine::PeerData,
    log_batch::{RaftLogBlock, RaftLogs},
    manifest::{PeerFile, generate_rlog_read_plan, rlog_by_entry_index},
    metrics::{ENGINE_APPLY_DURATION_HISTOGRAM, ENGINE_FETCH_ENTRIES_DURATION_HISTOGRAM},
    region_state_key,
};

pub struct RaftPeers {
    /// The primary directory holding engine data (e.g. WAL, rlog files).
    pub(crate) dir: PathBuf,

    pub(crate) peers: papaya::HashMap<u64, RwLock<PeerData>>,

    /// Per-peer Rlog files, shared between RfEngineCore and CompactWorker.
    /// Updated by the CompactWorker after compaction; read by RfEngineCore to
    /// decide which raft entries to offload to disk. Uses `ArcSwap` for
    /// lock-free reads to avoid impacting RfEngineCore's write path.
    pub(crate) peer_rlog_files: Arc<ArcSwap<HashMap<u64, VecDeque<PeerFile>>>>,
}

impl RaftPeers {
    pub fn new(
        dir: PathBuf,
        peer_rlog_files: Arc<ArcSwap<HashMap<u64, VecDeque<PeerFile>>>>,
    ) -> RaftPeers {
        Self {
            dir,
            peers: Default::default(),
            peer_rlog_files,
        }
    }

    /// Applies the write batch to memory without persisting it to WAL.
    pub fn apply(&self, wb: &WriteBatch, offload_epoch: u32) -> Vec<Vec<RaftLogBlock>> {
        let timer = Instant::now_coarse();
        let mut truncated_logs = vec![];
        let peer_rlog_files = if offload_epoch > 0 {
            Some(self.peer_rlog_files.load())
        } else {
            None
        };
        for (&peer_id, batch_data) in &wb.peers {
            let region_id = batch_data.meta.region_id;
            let keyspace_id = batch_data.meta.keyspace_id;
            tikv_util::set_current_region_thread_local(region_id);
            let guard = self.peers.guard();
            let peer_data = self.get_or_init_peer_data(peer_id, region_id, keyspace_id, &guard);
            let mut peer_data = peer_data.write().unwrap();
            let mut truncated = peer_data.apply(batch_data);
            if let Some(peer_rlog_files) = peer_rlog_files.as_ref() {
                if let Some(offload_idx) = peer_rlog_files
                    .get(&peer_id)
                    .and_then(|r| find_offload_index(r, &peer_data.raft_logs, offload_epoch))
                {
                    truncated.extend(peer_data.raft_logs.truncate(offload_idx));
                }
            }
            drop(peer_data);
            if !truncated.is_empty() {
                truncated_logs.push(truncated);
            }
        }
        ENGINE_APPLY_DURATION_HISTOGRAM.observe(timer.saturating_elapsed_secs());
        truncated_logs
    }

    pub fn is_empty(&self) -> bool {
        self.peers.is_empty()
    }

    pub fn get_entry_from_rlog(
        &self,
        peer_id: u64,
        target_index: u64,
        peer_file: PeerFile,
    ) -> crate::Result<Entry> {
        fetch_entries_from_rlog(
            &self.dir,
            peer_id,
            peer_file.first_index,
            peer_file.last_index,
            target_index,
            target_index,
        )?
        .into_iter()
        .next()
        .map(|op| op.to_entry())
        .ok_or_else(|| {
            Error::Other(format!(
                "peer {peer_id} entry {target_index} not found in {:?}",
                peer_file
            ))
        })
    }

    pub fn get_term(&self, peer_id: u64, index: u64) -> Option<u64> {
        let peers = self.peers.pin();
        let peer_data_ref = peers.get(&peer_id)?;
        let data = peer_data_ref.read().unwrap();
        if let Some(term_in_mem) = data.term(index) {
            return Some(term_in_mem);
        }
        // Check that the index is in the expected range.
        if index <= data.truncated_idx
            || (data.raft_logs.last_index() != 0 && index > data.raft_logs.last_index())
        {
            return None;
        }
        // Get the term from the rlog files.
        let file = self
            .peer_rlog_files
            .load()
            .get(&peer_id)
            .and_then(|files| rlog_by_entry_index(files, index))?;
        debug!(
            "Fetching term for peer_id {} index: {} from offloaded rlog file {:?}",
            peer_id, index, file
        );
        self.get_entry_from_rlog(peer_id, index, file)
            .ok()
            .map(|entry| entry.term)
    }

    pub fn get_truncated_index(&self, peer_id: u64) -> Option<u64> {
        let peers = self.peers.pin();
        let peer_data_ref = peers.get(&peer_id)?;
        let data = peer_data_ref.read().unwrap();
        Some(data.truncated_idx)
    }

    pub fn get_last_index(&self, peer_id: u64) -> Option<u64> {
        let peers = self.peers.pin();
        let peer_data_ref = peers.get(&peer_id)?;
        let data = peer_data_ref.read().unwrap();

        let last_index_in_mem = data.raft_logs.last_index();
        if last_index_in_mem != 0 {
            return Some(last_index_in_mem);
        }
        // Get the last index from offloaded entries.
        self.peer_rlog_files
            .load()
            .get(&peer_id)
            .and_then(|rlog_files| rlog_files.back())
            .map(|file| file.last_index)
            .filter(|&last_index| last_index > data.truncated_idx)
    }

    pub fn get_state(&self, peer_id: u64, key: &[u8]) -> Option<Bytes> {
        let peers = self.peers.pin();
        peers.get(&peer_id).and_then(|data| {
            data.read().unwrap().get_state(key).and_then(|val| {
                // TODO: seems it's impossible.
                if !val.is_empty() {
                    Some(val.clone())
                } else {
                    None
                }
            })
        })
    }

    /// Get the value of the last state key with the `prefix`. `prefix` must be
    /// non-empty.
    pub fn get_last_state_with_prefix(&self, peer_id: u64, prefix: &[u8]) -> Option<Bytes> {
        debug_assert!(!prefix.is_empty());
        let peers = self.peers.pin();
        let peer_data = peers.get(&peer_id)?;
        let peer_data = peer_data.read().unwrap();

        let mut end_prefix = prefix.to_vec();
        end_prefix[prefix.len() - 1] += 1;
        let range = Bytes::copy_from_slice(prefix)..Bytes::from(end_prefix);
        peer_data
            .meta
            .states
            .range(range)
            .next_back()
            .map(|(_, v)| v.clone())
    }

    /// Iterates states of the region in order or in desc order if `desc` is
    /// true until `f` returns error. The iterator will stop if the function
    /// returns false.
    pub fn iterate_peer_states<F>(&self, peer_id: u64, desc: bool, mut f: F)
    where
        F: FnMut(&Bytes, &Bytes) -> bool,
    {
        let peers = self.peers.pin();
        let peer_data = peers.get(&peer_id);
        let peer_data = match &peer_data {
            Some(data) => data.read().unwrap(),
            None => return,
        };

        let states = &peer_data.meta.states;
        if desc {
            for (k, v) in states.iter().rev() {
                if !f(k, v) {
                    break;
                }
            }
        } else {
            for (k, v) in states.iter() {
                if !f(k, v) {
                    break;
                }
            }
        }
    }

    pub fn get_peer_all_states(&self, peer_id: u64, desc: bool) -> Vec<(Bytes, Bytes)> {
        let mut states = vec![];
        self.iterate_peer_states(peer_id, desc, |k, v| {
            states.push((k.clone(), v.clone()));
            true
        });
        states
    }

    /// Collects at most `limit` distinct raft/region state versions older than
    /// `current_version` for a peer.
    pub fn collect_obsolete_state_versions(
        &self,
        peer_id: u64,
        current_version: u64,
        limit: usize,
    ) -> Vec<u64> {
        if limit == 0 {
            return Vec::new();
        }
        let peers = self.peers.pin();
        let Some(peer_data) = peers.get(&peer_id) else {
            return Vec::new();
        };
        let peer_data = peer_data.read().unwrap();
        let states = &peer_data.meta.states;
        let mut versions = BTreeSet::new();

        let version_state_ranges = [
            (
                Bytes::from_static(&[RAFT_STATE_KEY_BYTE]),
                Bytes::from_static(&[REGION_META_KEY_BYTE]),
            ),
            (
                Bytes::from_static(&[REGION_META_KEY_BYTE]),
                Bytes::from_static(&[REGION_META_KEY_BYTE + 1]),
            ),
        ];
        for (start, end) in version_state_ranges {
            for (key, _) in states.range(start..end) {
                if key.len() != 5 {
                    continue;
                }
                let version = u64::from(u32::from_be_bytes([key[1], key[2], key[3], key[4]]));
                if version >= current_version {
                    break;
                }
                versions.insert(version);
                if versions.len() == limit {
                    return versions.into_iter().collect();
                }
            }
        }
        versions.into_iter().collect()
    }

    /// Iterates stats of all regions in order or in desc order if `desc` is
    /// true and breaks one regions iteration if `f` returns false.
    pub fn iterate_all_states<F>(&self, desc: bool, mut f: F)
    where
        F: FnMut(u64, u64, &[u8], &[u8]) -> bool,
    {
        let peers = self.peers.pin();
        peers.iter().for_each(|(_, data)| {
            let data = data.read().unwrap();
            if data.truncated_idx == TRUNCATE_ALL_INDEX {
                return;
            }
            let peer_id = data.peer_id;
            let region_id = data.region_id;
            if desc {
                data.states
                    .iter()
                    .rev()
                    .take_while(|(k, v)| f(peer_id, region_id, k, v))
                    .count();
            } else {
                data.states
                    .iter()
                    .take_while(|(k, v)| f(peer_id, region_id, k, v))
                    .count();
            }
        });
    }

    pub(crate) fn get_or_init_peer_data<'a>(
        &self,
        peer_id: u64,
        region_id: u64,
        keyspace_id: u32,
        guard: &'a papaya::LocalGuard<'a>,
    ) -> &'a RwLock<PeerData> {
        self.peers.get_or_insert_with(
            peer_id,
            || RwLock::new(PeerData::new(peer_id, region_id, keyspace_id)),
            guard,
        )
    }

    /// Dumps the state of the region.
    pub fn get_peer_stats(&self, peer_id: u64) -> PeerStats {
        let peers = self.peers.pin();
        peers
            .get(&peer_id)
            .map(|data| data.read().unwrap().get_stats())
            .unwrap_or_default()
    }

    /// Returns the log index up to which logs should be truncated for the given
    /// peer so that its total log usage (including in-memory and offloaded
    /// entries) is reduced below the specified size limit.
    pub fn index_to_truncate_to_size(&self, peer_id: u64, size_limit: usize) -> u64 {
        let peers = self.peers.pin();
        peers
            .get(&peer_id)
            .map(|data| data.read().unwrap().index_to_truncate_to_size(size_limit))
            .unwrap_or_default()
    }

    pub fn get_region_peer_map(&self) -> HashMap<u64 /* region_id */, u64 /* peer_id */> {
        let mut region_to_peer = HashMap::with_capacity(self.peers.len());
        let mut id_pairs = Vec::with_capacity(self.peers.len());
        let peers = self.peers.pin();
        for (_, peer_ref) in peers.iter() {
            let peer_data = peer_ref.read().unwrap();
            let is_truncated = peer_data.truncated_idx == TRUNCATE_ALL_INDEX;
            id_pairs.push((peer_data.peer_id, peer_data.region_id, is_truncated));
        }
        // ensure the newer peer_id appear after the older peer_id, so it can replace
        // older.
        id_pairs.sort_by(|(peer_a, ..), (peer_b, ..)| peer_a.cmp(peer_b));
        for (peer_id, region_id, truncated) in id_pairs {
            if truncated {
                // The newer peer is already destroyed, the old peer is invalid too.
                region_to_peer.remove(&region_id);
            } else {
                // new peer_id replaces the older peer_id.
                region_to_peer.insert(region_id, peer_id);
            }
        }
        region_to_peer
    }

    pub fn get_raft_entry(&self, peer_id: u64, index: u64) -> Option<Entry> {
        let peers = self.peers.pin();
        let peer_data_ref = peers.get(&peer_id)?;
        let data = peer_data_ref.read().unwrap();
        if let Some(entry) = data.get(index) {
            return Some(entry);
        }
        // Check that the index is in the expected range.
        if index <= data.truncated_idx
            || (data.raft_logs.last_index() != 0 && index > data.raft_logs.last_index())
        {
            return None;
        }
        // Get the entry from the rlog files.
        let file = self
            .peer_rlog_files
            .load()
            .get(&peer_id)
            .and_then(|files| rlog_by_entry_index(files, index))?;
        self.get_entry_from_rlog(peer_id, index, file).ok()
    }

    pub fn fetch_raft_entries_to(
        &self,
        peer_id: u64,
        low: u64,
        high: u64,
        max_size: Option<usize>,
        buf: &mut Vec<Entry>,
        context: Option<raft::GetEntriesContext>,
    ) -> crate::Result<usize> {
        if high <= low {
            return Ok(0);
        }
        let old_len = buf.len();
        let peers = self.peers.pin();
        let peer_data = peers.get(&peer_id).ok_or(crate::Error::EntriesCompacted)?;
        let peer_data = peer_data.read().unwrap();
        if low <= peer_data.meta.truncated_idx {
            return Err(crate::Error::EntriesCompacted);
        }

        let timer = Instant::now_coarse();
        defer! {
            ENGINE_FETCH_ENTRIES_DURATION_HISTOGRAM.observe(timer.saturating_elapsed_secs())
        };
        let mut total_size = 0;
        // Some Raft logs may have been offloaded to disk. To fetch all entries
        // in the given range, we first check whether any part of the range
        // falls below the first in-memory index. If so, we need to read from
        // the offloaded log files (rlogs).
        //
        // Step 1: Load offloaded entries if necessary. `first_idx_in_mem == 0`
        // is a special case where all raft entries have been offloaded.
        let first_idx_in_mem = peer_data.raft_logs.first_index();
        let need_read_from_disk = first_idx_in_mem == 0 || low < first_idx_in_mem;
        if need_read_from_disk {
            info!(
                "peer {}: fetch_raft_entries_to entries [{},{}), first_idx_in_mem: {}",
                peer_id, low, high, first_idx_in_mem
            );
            let (disk_fetch_low, disk_fetch_high) = if first_idx_in_mem == 0 {
                (low, high)
            } else {
                (low, first_idx_in_mem.min(high))
            };
            let read_plan = generate_rlog_read_plan(
                self.peer_rlog_files
                    .load()
                    .get(&peer_id)
                    .ok_or(crate::Error::EntriesUnavailable)?,
                disk_fetch_low,
                disk_fetch_high,
            )?;

            if let Some(ctx) = &context {
                if ctx.can_async() {
                    return Err(crate::Error::AsyncFetch(AsyncFetchInfo {
                        peer_id,
                        low,
                        high,
                        max_size,
                        region_id: peer_data.region_id,
                    }));
                }
            }

            let mut expected_idx = disk_fetch_low;
            for (range, peer_file) in read_plan {
                debug!(
                    "peer {}: loading offloaded entries [{},{}] from file={:?}",
                    peer_id, range.start, range.end, peer_file
                );
                let entries = fetch_entries_from_rlog(
                    &self.dir,
                    peer_id,
                    peer_file.first_index,
                    peer_file.last_index,
                    range.start,
                    range.end,
                )
                .map_err(|_| crate::Error::EntriesUnavailable)?;
                for entry in entries {
                    let entry = entry.to_entry();
                    debug_assert!(
                        entry.index == expected_idx,
                        "peer {}, entry.index:{}, expected_idx:{}",
                        peer_id,
                        entry.index,
                        expected_idx
                    );
                    expected_idx += 1;
                    total_size += entry.compute_size() as usize;
                    buf.push(entry);
                    if max_size.is_some_and(|s| total_size >= s) {
                        // At least return one entry regardless of size limit.
                        return Ok(buf.len() - old_len);
                    }
                }
            }
            // Assert that the last entry index == disk_fetch_high - 1
            if let Some(last_entry) = buf.last() {
                debug_assert_eq!(
                    last_entry.index,
                    disk_fetch_high - 1,
                    "peer {}: last_entry.index = {}, expected {}",
                    peer_id,
                    last_entry.index,
                    disk_fetch_high - 1
                );
            }
        }

        // Step 2: Fetch in-memory entries to results.
        let need_read_from_mem = first_idx_in_mem != 0 && first_idx_in_mem < high;
        if need_read_from_mem {
            let start = low.max(first_idx_in_mem);
            for i in start..high {
                let entry = peer_data.get(i).ok_or(crate::Error::EntriesUnavailable)?;
                total_size += entry.compute_size() as usize;
                buf.push(entry);
                if max_size.is_some_and(|s| total_size >= s) {
                    // At least return one entry regardless of size limit.
                    break;
                }
            }
        }
        Ok(buf.len() - old_len)
    }

    pub fn load_region_state(
        &self,
        peer_id: u64,
        version: u64,
    ) -> Option<raft_serverpb::RegionLocalState> {
        let region_state_key = region_state_key(version);
        let region_state_val = self.get_state(peer_id, &region_state_key)?;
        let mut region_state = raft_serverpb::RegionLocalState::new();
        region_state.merge_from_bytes(&region_state_val).unwrap();
        Some(region_state)
    }
}

/// Finds the largest index eligible for offloading, given a target
/// `offload_epoch`.
///
/// Scans rlog files from newest to oldest and returns the last index of the
/// first file that satisfies:
/// 1. Has a valid epoch and `epoch_id <= offload_epoch`.
/// 2. Its last term matches the term of the corresponding in-memory entry. If
///    the terms don't match, the in-memory entry is newer and must not be
///    offloaded.
///
/// Returns:
/// - Some(index) if offloading can proceed up to that index
/// - None if nothing is eligible
pub(crate) fn find_offload_index(
    rlog_files: &VecDeque<PeerFile>,
    raft_logs: &RaftLogs,
    offload_epoch: u32,
) -> Option<u64> {
    if offload_epoch == 0 || raft_logs.first_index() == 0 {
        return None;
    }

    for f in rlog_files.iter().rev() {
        if f.epoch_id == 0 || f.epoch_id > offload_epoch {
            continue;
        }

        match raft_logs.get(f.last_index) {
            Some(entry) if entry.term == f.last_term as u64 => return Some(f.last_index),
            Some(_) => continue, // term mismatch - skip
            None => {
                if f.last_index > raft_logs.last_index() {
                    // In-memory entries with lower index have overwritten this
                    // rlog file. Skip it.
                    continue;
                } else {
                    // Entries in this rlog file are not in memory (i.e. already
                    // offloaded). No need to check older files.
                    return None;
                }
            }
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use std::{collections::VecDeque, path::PathBuf};

    use super::*;
    use crate::{RaftLogOp, WriteBatch, log_batch::RaftLogs, raft_state_key, region_state_key};

    #[test]
    fn test_collect_obsolete_state_versions() {
        let peers = RaftPeers::new(PathBuf::new(), Default::default());
        let peer_id = 10;
        let region_id = 20;
        let keyspace_id = 30;
        let mut wb = WriteBatch::new();
        for version in 1..=5001 {
            wb.set_state(
                peer_id,
                region_id,
                keyspace_id,
                &raft_state_key(version),
                b"raft",
            );
            wb.set_state(
                peer_id,
                region_id,
                keyspace_id,
                &region_state_key(version),
                b"region",
            );
        }
        peers.apply(&wb, 0);

        let first_batch = peers.collect_obsolete_state_versions(peer_id, 5001, 4096);
        assert_eq!(first_batch, (1..=4096).collect::<Vec<_>>());
        let mut gc_wb = WriteBatch::new();
        for version in first_batch {
            gc_wb.set_state(
                peer_id,
                region_id,
                keyspace_id,
                &raft_state_key(version),
                &[],
            );
            gc_wb.set_state(
                peer_id,
                region_id,
                keyspace_id,
                &region_state_key(version),
                &[],
            );
        }
        peers.apply(&gc_wb, 0);
        assert_eq!(
            peers.collect_obsolete_state_versions(peer_id, 5001, 4096),
            (4097..=5000).collect::<Vec<_>>()
        );
        assert!(peers.get_state(peer_id, &raft_state_key(5001)).is_some());
        assert!(peers.get_state(peer_id, &region_state_key(5001)).is_some());
        assert!(
            peers
                .collect_obsolete_state_versions(peer_id, 1, 4096)
                .is_empty()
        );
        assert!(
            peers
                .collect_obsolete_state_versions(peer_id, 5001, 0)
                .is_empty()
        );
        let orphan_peer_id = 11;
        let mut wb = WriteBatch::new();
        wb.set_state(
            orphan_peer_id,
            region_id,
            keyspace_id,
            &raft_state_key(1),
            b"raft",
        );
        wb.set_state(
            orphan_peer_id,
            region_id,
            keyspace_id,
            &region_state_key(2),
            b"region",
        );
        peers.apply(&wb, 0);
        assert_eq!(
            peers.collect_obsolete_state_versions(orphan_peer_id, 3, 4096),
            vec![1, 2]
        );

        let region_state_only_peer_id = 12;
        let mut wb = WriteBatch::new();
        for version in 1..=3 {
            wb.set_state(
                region_state_only_peer_id,
                region_id,
                keyspace_id,
                &region_state_key(version),
                b"region",
            );
        }
        peers.apply(&wb, 0);
        assert_eq!(
            peers.collect_obsolete_state_versions(region_state_only_peer_id, 4, 2),
            vec![1, 2]
        );
    }

    #[test]
    fn test_find_offload_index() {
        let mut logs = RaftLogs::default();
        let index_and_terms = vec![
            (20, 5),
            (21, 5),
            (22, 5),
            (23, 7),
            (24, 7),
            (25, 7),
            (26, 9),
            (27, 9),
            (28, 9),
        ];
        for (index, term) in index_and_terms {
            let mut e = Entry::default();
            e.index = index;
            e.term = term;
            logs.append(1, RaftLogOp::new(&e));
        }

        let files = VecDeque::from(vec![
            PeerFile::new(1, 15, 19, 4), // offloaded
            PeerFile::new(2, 20, 21, 5), // term matched
            PeerFile::new(0, 22, 22, 5), // rlog epoch == 0 (skipped)
            PeerFile::new(4, 23, 25, 6), // term mismatch
            PeerFile::new(5, 23, 24, 7), // term matched
            PeerFile::new(0, 25, 25, 7), // rlog epoch == 0 (skipped)
            PeerFile::new(7, 26, 30, 8), // term mismatch
            PeerFile::new(8, 26, 27, 9), // term matched
        ]);

        assert_eq!(find_offload_index(&files, &logs, 0), None);
        assert_eq!(find_offload_index(&files, &logs, 1), None);
        assert_eq!(find_offload_index(&files, &logs, 2), Some(21));
        assert_eq!(find_offload_index(&files, &logs, 3), Some(21));
        assert_eq!(find_offload_index(&files, &logs, 4), Some(21));
        assert_eq!(find_offload_index(&files, &logs, 5), Some(24));
        assert_eq!(find_offload_index(&files, &logs, 6), Some(24));
        assert_eq!(find_offload_index(&files, &logs, 7), Some(24));
        assert_eq!(find_offload_index(&files, &logs, 8), Some(27));

        // Edge cases: empty raft log or empty rlog files
        for offload_epoch in 0..=10 {
            assert_eq!(
                find_offload_index(&VecDeque::new(), &logs, offload_epoch),
                None
            );
            assert_eq!(
                find_offload_index(&files, &RaftLogs::default(), offload_epoch),
                None
            );
        }
    }
}