// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2026 PingCAP, Inc.

use builtin_dfs::BuiltinDfs;
use bytes::Bytes;
use http::{Request, StatusCode};
use hyper::{Body, Client};
use kvengine::dfs::{Dfs, FileType, Options};
use pd_client::PdClient;
use test_cloud_server::{ServerClusterBuilder, alloc_node_id_vec};
use tikv::config::TikvConfig;
use tikv_util::sys::SysQuota;
use tokio::runtime::Runtime;

use super::cfg_notify_once;

#[test]
fn test_builtin_dfs_byte_admission_rejects_while_body_is_active() {
    let byte_limit = 4 * 1024 * 1024;
    let nodes = alloc_node_id_vec(1);
    let mut cluster = ServerClusterBuilder::new(nodes.clone(), move |_, conf: &mut TikvConfig| {
        conf.dfs.s3_endpoint = "local".to_string();
        conf.kvengine.dfs_load_memory_ratio =
            byte_limit as f64 / SysQuota::memory_limit_in_bytes().max(1) as f64;
    })
    .pd_server_cnt(1)
    .build();
    cluster.wait_region_replicated(&[], 1);

    let status_addr = cluster.status_addr(nodes[0]);
    let pd_client = cluster.get_pd_client();
    let builtin_dfs = BuiltinDfs::new(pd_client.clone());
    let region = pd_client.get_region(b"").unwrap();
    let opts =
        Options::default().with_shard(region.get_id(), region.get_region_epoch().get_version());
    let runtime = Runtime::new().unwrap();

    let active_file_id = 12350u64;
    let second_file_id = 12351u64;
    runtime
        .block_on(builtin_dfs.create(
            active_file_id,
            Bytes::from(vec![8; 3 * 1024 * 1024 + 37]),
            opts,
        ))
        .unwrap();
    runtime
        .block_on(builtin_dfs.create(second_file_id, Bytes::from(vec![9; 2 * 1024 * 1024]), opts))
        .unwrap();

    let client = Client::new();
    let (body_started, body_chunk_guard) = cfg_notify_once("before_dfs_file_read_response_chunk");
    let active_url = format!(
        "http://{status_addr}/dfs/{active_file_id}?file_type={}",
        FileType::Sst.suffix()
    );
    let active_request =
        runtime.spawn(client.request(Request::get(active_url).body(Body::empty()).unwrap()));

    // Wait until the first response body is blocked before its first chunk.
    // Its admission guard must keep owning most of the byte budget.
    let resume_body = runtime.block_on(body_started).unwrap();
    let second_url = format!(
        "http://{status_addr}/dfs/{second_file_id}?file_type={}",
        FileType::Sst.suffix()
    );
    let rejected_response = runtime
        .block_on(client.request(Request::get(second_url).body(Body::empty()).unwrap()))
        .unwrap();
    assert_eq!(rejected_response.status(), StatusCode::SERVICE_UNAVAILABLE);

    resume_body.send(()).unwrap();
    drop(body_chunk_guard);
    let active_response = runtime.block_on(active_request).unwrap().unwrap();
    assert_eq!(active_response.status(), StatusCode::OK);
    drop(rejected_response);
    drop(active_response);
    cluster.stop();
}