// Copyright 2024 PingCAP, Inc.

use std::{
    collections::{HashMap, VecDeque},
    fmt, fs,
    io::{Read, Seek, SeekFrom},
    path::{Path, PathBuf},
    sync::{
        Arc, RwLock,
        atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering},
    },
    thread::JoinHandle,
};

use arc_swap::ArcSwap;
use bytes::Bytes;
use file_system::IoRateLimiter;
use kvengine::dfs::Dfs;
use rfenginepb::StoreBackupMeta;
use slog_global::{error, info};
use tikv_util::{
    DFS_WORKER_THREAD_NAME,
    mpsc::{Receiver, SendError, Sender},
    sys::thread::StdThreadBuildWrapper,
    time::{Instant, monotonic_raw_now_millis},
};

use crate::{
    BackupTask, CompressionType, DEFAULT_SNAPSHOT_RLOG_MULTIPART_UPLOAD_PART_SIZE, DfsMeter, Error,
    MAX_SNAPSHOT_RLOG_OBJECT_SIZE, NotifyOnDrop, Result, RfEngineConfig, RlogCache,
    bg_service_worker::{CompactTask, CompactWorker, WorkerHandle, backup_callback, wal_file_name},
    dfs_worker::{Healthy, LightweightBackupConfig, ObjectStorageTask, ObjectStorageWorker},
    log_batch::RaftLogBlock,
    manifest::{Manifest, PeerFile},
    write_batch::PeerBatch,
    writer::WalWriter,
};

pub(crate) struct ObjectStorageWorkerHandle {
    task_sender: Sender<ObjectStorageTask>,
    handle: JoinHandle<()>,
}

impl ObjectStorageWorkerHandle {
    fn try_send(&self, task: ObjectStorageTask) {
        // The DFS worker logs its exit; repeated send errors would cause a log storm.
        let _ = self.task_sender.send(task);
    }
}

#[derive(Serialize, Deserialize, Clone, Copy, Debug, PartialEq, Eq)]
#[serde(rename_all = "kebab-case")] // Be uniform with status server interfaces.
pub struct WalProgress {
    pub epoch: u32,
    pub offset: u64,
}

/// Runtime state of the DFS worker.
#[derive(Serialize, Deserialize, Clone, Copy, Debug, PartialEq, Eq)]
pub struct DfsWorkerStatus {
    /// Whether the worker receives new WAL and snapshot work.
    pub enabled: bool,
    /// Whether the configured worker thread is running.
    pub running: bool,
    /// Whether DFS contains a continuous recovery chain.
    pub healthy: bool,
}

pub(crate) enum ServiceTask {
    Dump {
        epoch_id: u32,
        start_off: u64,
        end_off: u64,
        callback: Box<dyn FnOnce(Result<(Bytes, bool)>) + Send>,
    },
    GetProgress {
        callback: Box<dyn FnOnce(Result<WalProgress>) + Send>,
    },
    GetDfsWorkerStatus {
        callback: tokio::sync::oneshot::Sender<Result<DfsWorkerStatus>>,
    },
    Rotate {
        epoch_id: u32,
        // Whether to cache new WriteBatch in memory for compaction.
        cache_wb_for_compact: bool,
        // Whether forcing compact worker to take and upload snapshot.
        require_snapshot: bool,
    },
    Write {
        wb: Arc<Vec<PeerBatch>>,
    },
    Backup(BackupTask),
    Truncates(Vec<Vec<RaftLogBlock>>),
    Upload(NotifyOnDrop),
    SetLightweightBackup {
        enabled: bool,
        callback: tokio::sync::oneshot::Sender<Result<()>>,
    },
    Close {
        force: bool,
    },
}

impl fmt::Debug for ServiceTask {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ServiceTask::Dump { .. } => write!(f, "ServiceTask::Dump"),
            ServiceTask::GetProgress { .. } => write!(f, "ServiceTask::GetProgress"),
            ServiceTask::GetDfsWorkerStatus { .. } => {
                write!(f, "ServiceTask::GetDfsWorkerStatus")
            }
            ServiceTask::Rotate { .. } => write!(f, "ServiceTask::Rotate"),
            ServiceTask::Write { .. } => write!(f, "ServiceTask::Write"),
            ServiceTask::Backup(_) => write!(f, "ServiceTask::Backup"),
            ServiceTask::Truncates(_) => write!(f, "ServiceTask::Truncates"),
            ServiceTask::Upload(_) => write!(f, "ServiceTask::Upload"),
            ServiceTask::SetLightweightBackup { enabled, .. } => {
                write!(f, "ServiceTask::SetLightweightBackup({enabled})")
            }
            ServiceTask::Close { force } => write!(f, "ServiceTask::Close({})", force),
        }
    }
}

const ASYNC_WRITER_SYNC_SIZE: u64 = 256 * 1024;

/// Service worker maintains the async WAL files, provide read service for HTTP
/// API, compaction and backup tasks.
pub(crate) struct ServiceWorker {
    engine_id: Arc<AtomicU64>,
    epoch_id: Arc<AtomicU32>,
    async_wal_writer: Option<WalWriter>,
    bytes_since_last_sync: u64,
    rx: Receiver<ServiceTask>,
    compact_worker_handle: WorkerHandle,
    dfs_worker_handle: Option<ObjectStorageWorkerHandle>,
    dfs_worker_healthy: Healthy,
    lightweight: Arc<AtomicBool>,
    // DFS upload statistics collected by the DFS worker.
    statistics: Arc<DfsMeter>,
    async_wal_sync_success_at_millis: Arc<AtomicU64>,
}

impl ServiceWorker {
    pub(crate) fn new(
        dir: PathBuf,
        cfg: &RfEngineConfig,
        async_wal_writer: Option<WalWriter>,
        rx: Receiver<ServiceTask>,
        manifest: Manifest,
        compacted_epoch: Arc<AtomicU32>,
        lightweight_backup: Option<(LightweightBackupConfig, Arc<dyn Dfs>, bool)>,
        lightweight: Arc<AtomicBool>,
        healthy: Healthy,
        compact_rate_limiter: Option<Arc<IoRateLimiter>>,
        async_wal_sync_success_at_millis: Arc<AtomicU64>,
        #[cfg(any(test, feature = "testexport"))] compact_force_stop: Arc<AtomicBool>,
        peer_rlog_files: Arc<ArcSwap<HashMap<u64, VecDeque<PeerFile>>>>,
    ) -> Self {
        // Create a statistic object for the DFS worker.
        let statistic = Arc::new(DfsMeter::default());
        let engine_id = manifest.engine_id.clone();
        let epoch_id = manifest.epoch_id + 1;
        let (compact_worker_tx, compact_rx) = tikv_util::mpsc::unbounded();
        let (dfs_worker_tx, dfs_worker_rx) = tikv_util::mpsc::unbounded();
        let rlog_cache = Arc::new(RwLock::new(RlogCache::new(
            cfg.rlog_cache_capacity.0 as usize,
            cfg.rlog_cache_size_threshold.0 as usize,
        )));
        let rlog_compression_type = lightweight_backup
            .as_ref()
            .map(|(config, ..)| config.rlog_compression_type)
            .unwrap_or(CompressionType::NoCompression);
        let snapshot_rlog_multipart_upload_threshold = lightweight_backup
            .as_ref()
            .map(|(config, ..)| config.snapshot_rlog_multipart_upload_threshold)
            .unwrap_or_default();
        let snapshot_rlog_multipart_upload_part_size = lightweight_backup
            .as_ref()
            .map(|(config, ..)| config.snapshot_rlog_multipart_upload_part_size)
            .unwrap_or(DEFAULT_SNAPSHOT_RLOG_MULTIPART_UPLOAD_PART_SIZE);
        let snapshot_rlog_multipart_upload_concurrency = lightweight_backup
            .as_ref()
            .map(|(config, ..)| config.snapshot_rlog_multipart_upload_concurrency)
            .unwrap_or(1);
        let snapshot_rlog_max_object_size = lightweight_backup
            .as_ref()
            .map(|(config, ..)| config.snapshot_rlog_max_object_size)
            .unwrap_or(MAX_SNAPSHOT_RLOG_OBJECT_SIZE as usize);
        let snapshot_rlog_multipart_dfs = lightweight_backup
            .as_ref()
            .map(|(_, dfs, _)| Arc::clone(dfs));
        let mut compact_worker = CompactWorker::new(
            dir.clone(),
            cfg,
            compact_rx,
            dfs_worker_tx.clone(),
            manifest,
            compacted_epoch.clone(),
            Arc::clone(&rlog_cache),
            rlog_compression_type,
            snapshot_rlog_multipart_upload_threshold,
            snapshot_rlog_multipart_upload_part_size,
            snapshot_rlog_multipart_upload_concurrency,
            snapshot_rlog_max_object_size,
            snapshot_rlog_multipart_dfs,
            compact_rate_limiter,
            #[cfg(any(test, feature = "testexport"))]
            compact_force_stop,
            peer_rlog_files,
        );
        let handle = std::thread::Builder::new()
            .name("compact-wal-worker".to_string())
            .spawn_wrapper(move || compact_worker.run())
            .unwrap();
        let compact_worker_handle = WorkerHandle {
            task_sender: compact_worker_tx.clone(),
            handle: Some(handle),
        };
        let epoch_rotate_len = cfg.epoch_rotate_len;

        let service_worker_epoch = Arc::new(AtomicU32::new(epoch_id));
        let mut enabled = false;
        let dfs_worker_handle = if let Some((cfg, s3fs, initially_enabled)) = lightweight_backup {
            enabled = initially_enabled && async_wal_writer.is_some();
            let mut dfs_worker = ObjectStorageWorker::new(
                cfg,
                s3fs,
                enabled,
                epoch_id,
                epoch_rotate_len,
                engine_id.clone(),
                healthy.clone(),
                dfs_worker_rx,
                compact_worker_tx,
                service_worker_epoch.clone(),
                statistic.clone(),
            );
            let handle = std::thread::Builder::new()
                .name(DFS_WORKER_THREAD_NAME.to_string())
                .spawn_wrapper(move || dfs_worker.run())
                .unwrap();
            Some(ObjectStorageWorkerHandle {
                task_sender: dfs_worker_tx,
                handle,
            })
        } else {
            None
        };
        lightweight.store(enabled, Ordering::SeqCst);
        ServiceWorker {
            engine_id,
            epoch_id: service_worker_epoch,
            async_wal_writer,
            bytes_since_last_sync: 0,
            dfs_worker_handle,
            rx,
            compact_worker_handle,
            dfs_worker_healthy: healthy,
            lightweight,
            statistics: statistic,
            async_wal_sync_success_at_millis,
        }
    }

    pub(crate) fn dfs_statistic(&self) -> Option<Arc<DfsMeter>> {
        if self.dfs_worker_handle.is_some() {
            Some(self.statistics.clone())
        } else {
            None
        }
    }

    pub(crate) fn run(&mut self) {
        while let Ok(task) = self.rx.recv() {
            match task {
                ServiceTask::Write { wb } => {
                    self.handle_write(wb);
                }
                ServiceTask::Dump {
                    epoch_id,
                    start_off,
                    end_off,
                    callback,
                } => {
                    self.handle_dump(epoch_id, start_off, end_off, callback);
                }
                ServiceTask::GetProgress { callback } => {
                    self.handle_get_progress(callback);
                }
                ServiceTask::GetDfsWorkerStatus { callback } => {
                    let _ = callback.send(Ok(self.dfs_worker_status()));
                }
                ServiceTask::Backup(task) => {
                    self.handle_backup(task);
                }
                ServiceTask::Rotate {
                    epoch_id,
                    cache_wb_for_compact,
                    require_snapshot,
                } => {
                    self.handle_rotate(epoch_id, cache_wb_for_compact, require_snapshot);
                }
                ServiceTask::Truncates(truncates) => drop(truncates),
                ServiceTask::Upload(sender) => {
                    self.handle_flush(sender);
                }
                ServiceTask::SetLightweightBackup { enabled, callback } => {
                    self.handle_set_lightweight_backup(enabled, callback);
                }
                ServiceTask::Close { force } => {
                    self.handle_close(force);
                    break;
                }
            }
        }
    }

    fn is_lightweight_enabled(&self) -> bool {
        self.lightweight.load(Ordering::SeqCst)
    }

    fn dfs_worker_status(&self) -> DfsWorkerStatus {
        let running = self.is_lightweight_enabled()
            && self
                .dfs_worker_handle
                .as_ref()
                .is_some_and(|handle| !handle.handle.is_finished());
        DfsWorkerStatus {
            enabled: self.is_lightweight_enabled(),
            running,
            healthy: self.dfs_worker_healthy.is_healthy(),
        }
    }

    fn get_engine_id(&self) -> u64 {
        self.engine_id.load(Ordering::SeqCst)
    }

    fn handle_write(&mut self, wb: Arc<Vec<PeerBatch>>) {
        let async_wal_sync_success_at_millis = self.async_wal_sync_success_at_millis.clone();
        if let Some(wal_writer) = &mut self.async_wal_writer {
            let old_file_off = wal_writer.file_off;
            let (epoch_id, file_off, rotated) = wal_writer.write_batch(&wb, false).unwrap();
            if rotated {
                // When rotated, the old file is already synced.
                self.bytes_since_last_sync = 0;
                async_wal_sync_success_at_millis
                    .fetch_max(monotonic_raw_now_millis(), Ordering::Relaxed);
            } else {
                self.bytes_since_last_sync += file_off - old_file_off;
            }
            if self.bytes_since_last_sync > ASYNC_WRITER_SYNC_SIZE {
                wal_writer.file().sync_data().unwrap();
                self.bytes_since_last_sync = 0;
                async_wal_sync_success_at_millis
                    .fetch_max(monotonic_raw_now_millis(), Ordering::Relaxed);
            }
            if self.is_lightweight_enabled() {
                // Send write task to object storage worker.
                let task = crate::dfs_worker::ObjectStorageTask::Sync { epoch_id, file_off };
                self.dfs_worker_handle.as_ref().unwrap().try_send(task);
            }
        }

        let _ = self
            .compact_worker_handle
            .task_sender
            .send(CompactTask::Write { wb });
    }

    /// `partial_content`:
    /// - When `end_off == 0`, `partial_content` is true means that there is
    ///   more epoches.
    /// - When `end_off > 0`, `partial_content` is true means that the epoch (of
    ///   `epoch_id`) has more content after `end_off`.
    fn handle_dump(
        &mut self,
        epoch_id: u32,
        start_off: u64,
        end_off: u64,
        callback: Box<dyn FnOnce(Result<(Bytes, bool /* partial_content */)>) + Send>,
    ) {
        let res = if let Some(writer) = self.async_wal_writer.as_ref() {
            let mut partial_content = false;
            let store_id = self.get_engine_id();
            // Check the WAL chunk meta is valid.
            if epoch_id > writer.epoch_id || (end_off > 0 && start_off >= end_off) {
                let msg = format!(
                    "{}: invalid dump wal chunk epoch {} start_off {} end_off {:?} writer epoch {}, file_off {}",
                    store_id, epoch_id, start_off, end_off, writer.epoch_id, writer.file_off
                );
                error!("{}", msg);
                callback(Err(Error::Other(msg)));
                return;
            } else if epoch_id + writer.epoch_rotate_len as u32 - 1 < writer.epoch_id {
                // The epoch_id is too old and has been overwritten.
                info!(
                    "{}: handle dump: epoch is overwritten: epoch {} writer epoch {}",
                    store_id, epoch_id, writer.epoch_id
                );
                callback(Err(Error::WalEpochOverwritten { epoch_id }));
                return;
            }
            // Dump the WAL chunk from offset start_off to end_off.
            #[allow(clippy::collapsible_else_if)]
            let effective_end_off = if end_off == 0 {
                if writer.epoch_id == epoch_id {
                    writer.file_off
                } else {
                    partial_content = true;
                    let wal_file_name =
                        wal_file_name(&writer.dir, epoch_id, writer.epoch_rotate_len);
                    let file_meta = fs::metadata(wal_file_name).unwrap();
                    file_meta.len()
                }
            } else {
                if writer.epoch_id == epoch_id {
                    debug_assert!(end_off <= writer.file_off);
                    partial_content = true;
                    end_off.min(writer.file_off)
                } else {
                    let wal_file_name =
                        wal_file_name(&writer.dir, epoch_id, writer.epoch_rotate_len);
                    let file_len = fs::metadata(wal_file_name).unwrap().len();
                    debug_assert!(end_off <= file_len);
                    partial_content = end_off < file_len;
                    end_off.min(file_len)
                }
            };
            info!(
                "{}: dump latest wal epoch {} start_off {} end_off {} effective_end_off {} writer epoch {}",
                store_id, epoch_id, start_off, end_off, effective_end_off, writer.epoch_id,
            );
            match dump_wal_chunk(
                &writer.dir,
                epoch_id,
                writer.epoch_rotate_len,
                start_off,
                effective_end_off,
            ) {
                Ok(chunk) => Ok((chunk, partial_content)),
                Err(err) => {
                    let msg = format!(
                        "{}: dump wal chunk epoch {} start_off {} end_off {} failed {:?}",
                        self.get_engine_id(),
                        epoch_id,
                        start_off,
                        effective_end_off,
                        err
                    );
                    error!("{}", msg);
                    Err(Error::Other(msg))
                }
            }
        } else {
            Err(Error::AsyncWriterDisabled)
        };
        callback(res);
    }

    fn handle_get_progress(&mut self, callback: Box<dyn FnOnce(Result<WalProgress>) + Send>) {
        let res = if let Some(writer) = &self.async_wal_writer {
            let progress = WalProgress {
                epoch: writer.epoch_id,
                offset: writer.file_off,
            };
            Ok(progress)
        } else {
            Err(Error::AsyncWriterDisabled)
        };
        callback(res);
    }

    fn handle_flush(&mut self, sender: NotifyOnDrop) {
        if self.is_lightweight_enabled() {
            // Send flush task to object storage worker.
            self.dfs_worker_handle
                .as_ref()
                .unwrap()
                .try_send(ObjectStorageTask::Flush(sender));
        }
    }

    fn handle_set_lightweight_backup(
        &mut self,
        enabled: bool,
        callback: tokio::sync::oneshot::Sender<Result<()>>,
    ) {
        if self.is_lightweight_enabled() == enabled {
            let _ = callback.send(Ok(()));
            return;
        }

        let Some(dfs_worker_handle) = self.dfs_worker_handle.as_ref() else {
            let _ = callback.send(Err(Error::Backup("dfs worker is not set".to_string())));
            return;
        };
        if enabled && self.async_wal_writer.is_none() {
            let _ = callback.send(Err(Error::Backup(
                "async wal writer is not set".to_string(),
            )));
            return;
        }

        self.lightweight.store(enabled, Ordering::SeqCst);
        let task = ObjectStorageTask::SetEnabled { enabled, callback };
        if let Err(SendError(ObjectStorageTask::SetEnabled { callback, .. })) =
            dfs_worker_handle.task_sender.send(task)
        {
            self.lightweight.store(!enabled, Ordering::SeqCst);
            let _ = callback.send(Err(Error::Backup("dfs worker is closed".to_string())));
            return;
        }

        // Enable is only a state transition. Reuse the normal Sync message to
        // catch the worker up even when no new write arrives afterwards.
        if enabled {
            let writer = self.async_wal_writer.as_ref().unwrap();
            dfs_worker_handle.try_send(ObjectStorageTask::Sync {
                epoch_id: writer.epoch_id,
                file_off: writer.file_off,
            });
        }
    }

    fn handle_rotate(&mut self, epoch_id: u32, cache_wb_for_compact: bool, require_snapshot: bool) {
        let prev_epoch_id = self.epoch_id.swap(epoch_id + 1, Ordering::SeqCst);
        assert_eq!(prev_epoch_id, epoch_id); // panic on unexpected to avoid data corruption.
        let mut dfs_file_off = None;
        if let Some(writer) = self.async_wal_writer.as_mut() {
            assert!(
                writer.epoch_id >= epoch_id,
                "invalid rotate epoch {}, async writer epoch {}",
                epoch_id,
                writer.epoch_id,
            );
            if writer.epoch_id == epoch_id {
                dfs_file_off = Some(writer.file_off);
                writer.rotate().unwrap();
            }
        }

        // Queue compaction before the DFS worker can request its snapshot.
        self.compact_worker_handle.try_send(CompactTask::Rotate {
            epoch_id,
            cache_wb_for_compact,
        });
        if let (Some(dfs_worker_handle), Some(file_off)) = (&self.dfs_worker_handle, dfs_file_off) {
            dfs_worker_handle.try_send(ObjectStorageTask::Rotate {
                epoch_id,
                file_off,
                require_snapshot,
            });
        }
    }

    fn handle_close(&mut self, force: bool) {
        // Close and join object storage thread.
        if let Some(dfs_worker_handle) = self.dfs_worker_handle.take() {
            // If force close, we skip flushing wal chunk and close task thread.
            if !force {
                let (tx, _) = tokio::sync::oneshot::channel();
                dfs_worker_handle.try_send(ObjectStorageTask::Flush(tx));
            }
            dfs_worker_handle.try_send(ObjectStorageTask::Close);
            dfs_worker_handle.handle.join().unwrap();
        }
        self.compact_worker_handle.try_send(CompactTask::Close);
        let join_handle = self.compact_worker_handle.handle.take().unwrap();
        join_handle.join().unwrap();
    }

    fn handle_backup(&mut self, mut task: BackupTask) {
        if let Some(async_writer) = self.async_wal_writer.as_ref() {
            task.file_off = async_writer.file_off;
        }
        if task.config.lightweight && self.is_lightweight_enabled() {
            self.lightweight_backup(task);
        } else {
            (task.callback)(Err(Error::Backup("backup not enabled".into())));
        }
    }

    fn lightweight_backup(&mut self, task: BackupTask) {
        if self.async_wal_writer.is_none() {
            return backup_callback(
                task,
                Err(Error::Backup("async wal writer is not set".to_string())),
                "light_fail",
                Instant::now(),
            );
        }
        let async_writer = self.async_wal_writer.as_ref().unwrap();

        if !self.dfs_worker_healthy.is_healthy() {
            return backup_callback(
                task,
                Err(Error::DfsWorkerUnhealthy {
                    store_id: self.get_engine_id(),
                    epoch_id: async_writer.epoch_id,
                }),
                "light_fail",
                Instant::now(),
            );
        }

        let mut backup_meta = StoreBackupMeta::default();
        backup_meta.set_store_id(self.engine_id.load(Ordering::SeqCst));
        backup_meta.set_epoch(async_writer.epoch_id);
        backup_meta.set_offset(task.file_off);

        backup_callback(task, Ok(backup_meta), "light_success", Instant::now());
    }
}

fn dump_wal_chunk(
    dir: &Path,
    epoch_id: u32,
    epoch_rotate_len: usize,
    start_off: u64,
    end_off: u64,
) -> crate::Result<Bytes> {
    // `epoch_id` already checked in the caller.
    let mut file = fs::File::open(wal_file_name(dir, epoch_id, epoch_rotate_len))?;
    let file_len = file.metadata()?.len();
    if end_off > file_len {
        return Err(Error::Eof);
    }

    // `start_off` < `end_off` already checked in the caller.
    file.seek(SeekFrom::Start(start_off))?;
    let dump_len = (end_off - start_off) as usize;
    let mut buf = vec![0; dump_len];
    file.read_exact(&mut buf)?;
    Ok(Bytes::from(buf))
}