// Copyright 2025 PingCAP, Inc.

use std::{fs, path::Path, sync::Arc, time::Duration};

use futures::executor::block_on;
use kvengine::dfs::{Dfs, new_dfs_from_config};
use kvproto::{metapb::Region, raft_serverpb::RegionLocalState};
use protobuf::Message;
use raft_proto::eraftpb::{Entry, EntryType};
use rand::RngCore;
use rfengine::{
    MAX_EPOCH_BACKWARD, MAX_EPOCH_ROTATE_LEN, MIN_EPOCH_ROTATE_LEN, WriteBatch,
    parse_delayed_to_epoch_from_snapshot_key, region_state_key, snapshot_rlog_key_prefix,
};
use test_cloud_server::{must_wait, oss::prepare_dfs};
use tikv_util::config::{ReadableDuration, ReadableSize};

#[test]
fn test_rfengine_dfs_worker() {
    let (temp_dir, mut oss, mut dfs_conf) = prepare_dfs("rfengine_dfs_worker");
    // By modifying the dispatch_timeout and write_timeout_per_mb, we can make
    // dfs_worker fail.
    dfs_conf.conn_options.dispatch_timeout = ReadableDuration::millis(50);
    dfs_conf.conn_options.write_timeout_per_mb = ReadableDuration::millis(50);

    dfs_conf.conn_options.max_retry_count = 1;

    let mut rf_cfg = rfengine::RfEngineConfig::default();
    rf_cfg.wal_chunk_target_file_size = ReadableSize::kb(128);
    rf_cfg.target_file_size = ReadableSize::kb(512);
    rf_cfg.epoch_rotate_len = MAX_EPOCH_ROTATE_LEN;
    rf_cfg.lightweight_backup = true;
    rf_cfg.wal_sync_dir = format!("{}/wal", temp_dir.path().to_str().unwrap());
    fs::create_dir_all(&rf_cfg.wal_sync_dir).unwrap();
    let dir = temp_dir.path().join("raftdb");
    let dfs = new_dfs_from_config(dfs_conf);
    let mut raft = open_engine(dir.as_path(), &rf_cfg, dfs.clone());
    init_raft_data(&raft);

    // The dfs worker starts healthy after init.
    wait_health(&raft, true, "wait init");

    // Make DFS fail all so the worker becomes unhealthy.
    oss.set_fail_all(true);
    let mut entry_index = 1;
    let entry_size = 64 * 1024;
    write_wal_to_epoch(&raft, &mut entry_index, entry_size, 6);
    wait_health(&raft, false, "wait unhealthy");
    raft.stop_worker(false);

    // Restart the engine while the worker is unhealthy; it should remain unhealthy.
    raft = open_engine(dir.as_path(), &rf_cfg, dfs.clone());
    block_on(raft.upload_wal_chunk());
    assert!(!raft.is_dfs_worker_healthy());

    // Restore OSS so WAL uploads can succeed again, but the worker should
    // recover healthy after a snapshot is persisted.
    oss.set_fail_all(false);
    oss.set_put_delay("snapshots", Duration::from_millis(150));
    write_wal_to_epoch(&raft, &mut entry_index, entry_size, 9);
    assert!(
        !raft.is_dfs_worker_healthy(),
        "dfs worker should remain unhealthy until snapshot upload finishes"
    );
    oss.set_put_delay("snapshots", Duration::from_millis(0));
    write_wal_to_epoch(&raft, &mut entry_index, entry_size, 10);
    wait_health(&raft, true, "wait recover after snapshot");
    raft.stop_worker(false);

    // Restart the engine while healthy; it should remain healthy after restart.
    raft = open_engine(dir.as_path(), &rf_cfg, dfs.clone());
    wait_health(&raft, true, "stay healthy after restart");
    oss.set_fail_all(true);
    write_wal_to_epoch(&raft, &mut entry_index, entry_size, 11);
    wait_health(&raft, false, "became unhealthy after oss fail");
    oss.set_fail_all(false);
    write_wal_to_epoch(&raft, &mut entry_index, entry_size, 17);
    wait_health(&raft, true, "recover healthy after snapshot");

    write_wal_to_epoch(&raft, &mut entry_index, entry_size, 56);
    let snapshot_prefix = snapshot_rlog_key_prefix(100);
    let (snapshots, _) = dfs
        .list_objects("", Some(&snapshot_prefix), Some(64))
        .unwrap();
    let snapshot_epoch = snapshots
        .iter()
        .filter_map(|o| parse_delayed_to_epoch_from_snapshot_key(Some(&o.key)))
        .max()
        .unwrap();
    assert!(raft.is_dfs_worker_healthy());
    // tolerate snapshot fail, still healthy
    oss.set_put_delay("snapshots", Duration::from_millis(150));

    // A failed snapshot within the lag limit keeps the worker healthy.
    write_wal_to_epoch(
        &raft,
        &mut entry_index,
        entry_size,
        snapshot_epoch + MAX_EPOCH_BACKWARD - 1,
    );
    assert!(raft.is_dfs_worker_healthy());

    // Exceeding the snapshot lag limit makes the worker unhealthy.
    write_wal_to_epoch(
        &raft,
        &mut entry_index,
        entry_size,
        snapshot_epoch + MAX_EPOCH_BACKWARD + 1,
    );
    block_on(raft.upload_wal_chunk());
    wait_health(&raft, false, "wait snap lag too much");

    oss.set_put_delay("snapshots", Duration::from_secs(0));
    write_wal_to_epoch(
        &raft,
        &mut entry_index,
        entry_size,
        snapshot_epoch + MAX_EPOCH_BACKWARD + 10,
    );
    wait_health(&raft, true, "recover by snapshot");

    raft.stop_worker(false);
    oss.shutdown();
}

fn init_raft_data(raft: &rfengine::RfEngine) {
    let mut wb = WriteBatch::new();
    let key = region_state_key(1);
    let mut region_local_state = RegionLocalState::new();
    let mut region = Region::new();
    region.set_id(1);
    region_local_state.set_region(region);
    let val = region_local_state.write_to_bytes().unwrap();
    wb.set_state_bytes(1, 1, 0, key, val.into());
    raft.write(wb).unwrap();
}

fn write_wal_to_epoch(
    engine: &rfengine::RfEngine,
    entry_index: &mut u64,
    payload: usize,
    target_epoch: u32,
) {
    let mut current_epoch = engine.get_epoch_offset().0;
    while current_epoch < target_epoch {
        append_logs(engine, entry_index, payload, 1);
        let new_epoch = engine.get_epoch_offset().0;
        if new_epoch != current_epoch {
            current_epoch = new_epoch;
            // Avoid the DFS worker falling behind and its WAL being overwritten.
            if current_epoch.is_multiple_of((MAX_EPOCH_ROTATE_LEN - MIN_EPOCH_ROTATE_LEN) as u32) {
                block_on(engine.upload_wal_chunk());
            }
        }
    }
}

fn open_engine(
    dir: &Path,
    cfg: &rfengine::RfEngineConfig,
    dfs: Arc<dyn Dfs>,
) -> rfengine::RfEngine {
    let raft = rfengine::RfEngine::open(dir, cfg, None, Some(dfs)).unwrap();
    raft.set_engine_id(100);
    raft
}

fn wait_health(raft: &rfengine::RfEngine, healthy: bool, label: &str) {
    must_wait(
        || raft.is_dfs_worker_healthy() == healthy,
        15,
        || label.to_string(),
    );
}

// Helper to append some raft logs to advance WAL progress.
fn append_logs(engine: &rfengine::RfEngine, entry_index: &mut u64, payload: usize, count: usize) {
    let mut rng = rand::thread_rng();
    for _ in 0..count {
        let mut entry = Entry::new();
        entry.set_index(*entry_index);
        entry.set_term(1);
        entry.set_entry_type(EntryType::EntryNormal);
        let mut data_buf = vec![0u8; payload];
        rng.fill_bytes(&mut data_buf);
        entry.set_data(data_buf.into());
        let mut wb = WriteBatch::new();
        wb.append_raft_log(1, 1, 0, &entry);
        engine.write(wb).unwrap();
        *entry_index += 1;
    }
}