// Copyright 2026 PingCAP, Inc.

use std::{iter::FromIterator, time::Duration};

use bytes::{BufMut, BytesMut};
use cloud_server::status_server::SstMeta;
use http::{Method, Request, StatusCode};
use hyper::{Body, Client};
use kvengine::dfs::DFSConfig;
use serde_json::{Value as JsonValue, json};
use test_cloud_server::{
    ServerClusterBuilder, TikvWorkerOptions, alloc_node_id_vec, oss::ObjectStorageService,
};
use tikv::config::TikvConfig;
use tikv_util::codec::bytes::encode_bytes;

use crate::{alloc_node_id, i_to_key_with_keyspace, i_to_val};

const DATA_SIZE: usize = 50;
const NODES_SIZE: usize = 1;

#[test]
fn test_write_sst() {
    let base_dir = tempfile::Builder::new()
        .prefix("test_write_sst")
        .tempdir()
        .unwrap();

    // Start mock S3 service
    let oss_dir = base_dir.path().join("oss");
    let mut oss = ObjectStorageService::new(oss_dir.clone());
    oss.start_server();

    // Configure DFS (used by both TiKV and Worker)
    let dfs_config = DFSConfig {
        prefix: "test_write_sst".to_string(),
        s3_endpoint: format!("http://127.0.0.1:{}", oss.port()),
        s3_key_id: "admin".to_string(),
        s3_secret_key: "admin".to_string(),
        s3_bucket: "test_write_sst".to_string(),
        s3_region: "local".to_string(),
        zstd_compression_level: "3".to_string(),
        ..Default::default()
    };

    // Start TiKV cluster (provides PD)
    let nodes = Vec::from_iter((0..NODES_SIZE).map(|_| alloc_node_id()));
    let mut cluster = ServerClusterBuilder::new(nodes.clone(), |_, conf: &mut TikvConfig| {
        conf.dfs = dfs_config.clone(); // Use the same DFS config
    })
    .pd_server_cnt(1)
    .tikv_worker_cnt(1)
    .build();
    let dfs = cluster.get_dfs().unwrap();
    let runtime = dfs.get_runtime();
    cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default());
    runtime.block_on(cluster.tikv_workers_must_healthy(Duration::from_secs(10)));
    cluster.wait_region_replicated(&[], NODES_SIZE);
    let mut tikv_client = cluster.new_client(); // Get client for later use if needed
    let pd_client = tikv_client.pd_client().clone();
    let cluster_id = pd_client.get_cluster_id().unwrap();

    let keyspace_id = 42;
    let i_to_tidb_key = i_to_key_with_keyspace(keyspace_id);
    let cfg = cloud_encryption::KeyspaceEncryptionConfig {
        enabled: true,
        cmek_id: Some("random".into()),
        vendor: Some("test".into()),
        ..Default::default()
    };
    pd_client.set_keyspace_encryption(keyspace_id, cfg).unwrap();
    tikv_client.split_keyspace(keyspace_id);
    cluster.wait_region_replicated(&i_to_tidb_key(0), NODES_SIZE);

    // Generate data [0..DATA_SIZE)
    let kvs: Vec<_> = (0..DATA_SIZE)
        .map(|i| (i_to_tidb_key(i), i_to_val(i)))
        .collect();

    // Encode data
    let mut body_buf = BytesMut::new();
    for (key, val) in &kvs {
        body_buf.put_u16_le(key.len() as u16);
        body_buf.put_slice(key);
        body_buf.put_u32_le(val.len() as u32);
        body_buf.put_slice(val);
    }
    let body_bytes = body_buf.freeze();

    // Send PUT /write_sst request
    // Use the cluster's PD client to get timestamp
    let commit_ts = runtime.block_on(pd_client.get_tso()).unwrap().into_inner();
    let write_url = format!(
        "http://{}/write_sst?cluster_id={}&commit_ts={}",
        cluster.tikv_worker_endpoints()[0],
        cluster_id,
        commit_ts
    );
    let http_client = Client::new();

    let req = Request::builder()
        .method(Method::PUT)
        .uri(&write_url)
        .header("Content-Type", "application/octet-stream")
        .body(Body::from(body_bytes.clone()))
        .unwrap();

    let resp = runtime.block_on(http_client.request(req)).unwrap();
    assert_eq!(resp.status(), StatusCode::OK,);

    // Read response body into bytes for JSON parsing
    let bytes = runtime
        .block_on(hyper::body::to_bytes(resp.into_body()))
        .unwrap();

    // Parse SstMeta from response
    let resp_body: JsonValue = serde_json::from_slice(&bytes).unwrap();
    let sst_meta_json = resp_body.get("sst_meta").unwrap();
    let sst_meta: SstMeta = serde_json::from_value(sst_meta_json.clone()).unwrap();

    // Get region info for ingestion
    let target_key = kvs[0].0.clone();
    // Use the PD client obtained earlier
    let region = pd_client.get_region(&target_key).unwrap();
    let region_id = region.get_id();
    let epoch = region.get_region_epoch().clone();

    // Send /ingest_s3 request to TiKV node
    let node_id = nodes[0];
    let status_addr = cluster.status_addr(node_id);
    let ingest_url = format!(
        "http://{}/ingest_s3?cluster_id={}&region_id={}&epoch_version={}",
        status_addr,
        cluster_id,
        region_id,
        epoch.get_version()
    );
    let req_body = json!(sst_meta).to_string();

    let ingest_req = Request::builder()
        .method(Method::POST)
        .uri(&ingest_url)
        .header("Content-Type", "application/json")
        .body(Body::from(req_body))
        .unwrap();

    let ingest_resp = runtime.block_on(http_client.request(ingest_req)).unwrap();
    assert_eq!(ingest_resp.status(), StatusCode::OK);

    // Verify data in TiKV
    std::thread::sleep(Duration::from_millis(200));
    for (key, expected_val) in &kvs {
        let (val, _) = tikv_client.must_get_key(key, tikv_util::time::Instant::now());
        assert_eq!(val, *expected_val,);
    }

    // Case: a new SST overlaps with existing data.
    // Generate KVs for an overlapping SST. If DATA_SIZE is 50, this will be keys
    // [25, 75).
    let overlap_kvs: Vec<_> = (DATA_SIZE / 2..DATA_SIZE + DATA_SIZE / 2)
        .map(|i| (i_to_tidb_key(i), i_to_val(i)))
        .collect();

    // Encode data for the overlapping SST
    let mut overlap_body_buf = BytesMut::new();
    for (key, val) in &overlap_kvs {
        overlap_body_buf.put_u16_le(key.len() as u16);
        overlap_body_buf.put_slice(key);
        overlap_body_buf.put_u32_le(val.len() as u32);
        overlap_body_buf.put_slice(val);
    }
    let overlap_body_bytes = overlap_body_buf.freeze();

    // Send PUT /write_sst request for the overlapping SST
    let overlap_commit_ts = runtime.block_on(pd_client.get_tso()).unwrap().into_inner();
    let overlap_write_url = format!(
        "http://{}/write_sst?cluster_id={}&commit_ts={}",
        cluster.tikv_worker_endpoints()[0],
        cluster_id,
        overlap_commit_ts
    );

    let overlap_req = Request::builder()
        .method(Method::PUT)
        .uri(&overlap_write_url)
        .header("Content-Type", "application/octet-stream")
        .body(Body::from(overlap_body_bytes.clone()))
        .unwrap();

    let overlap_resp = runtime.block_on(http_client.request(overlap_req)).unwrap();
    assert_eq!(overlap_resp.status(), StatusCode::OK,);

    // Parse SstMeta from the overlapping SST response
    let overlap_bytes_resp = runtime
        .block_on(hyper::body::to_bytes(overlap_resp.into_body()))
        .unwrap();
    let overlap_resp_body: JsonValue = serde_json::from_slice(&overlap_bytes_resp).unwrap();
    let overlap_sst_meta_json = overlap_resp_body.get("sst_meta").unwrap();
    let overlap_sst_meta: SstMeta = serde_json::from_value(overlap_sst_meta_json.clone()).unwrap();

    // Send /ingest_s3 request for the overlapping SST to the same TiKV node and
    // region
    let ingest_overlap_url = format!(
        "http://{}/ingest_s3?cluster_id={}&region_id={}&epoch_version={}",
        status_addr,
        cluster_id,
        region_id,
        epoch.get_version()
    );
    let overlap_req_body = json!(overlap_sst_meta).to_string();

    let ingest_overlap_req = Request::builder()
        .method(Method::POST)
        .uri(&ingest_overlap_url)
        .header("Content-Type", "application/json")
        .body(Body::from(overlap_req_body))
        .unwrap();

    let ingest_overlap_resp = runtime
        .block_on(http_client.request(ingest_overlap_req))
        .unwrap();
    assert_eq!(ingest_overlap_resp.status(), StatusCode::OK);

    // Verify data in TiKV after overlapping ingest
    std::thread::sleep(Duration::from_millis(200));
    for (key, expected_val) in &kvs {
        let (val, _) = tikv_client.must_get_key(key, tikv_util::time::Instant::now());
        assert_eq!(val, *expected_val);
    }
    for (key, expected_val) in &overlap_kvs {
        let (val, _) = tikv_client.must_get_key(key, tikv_util::time::Instant::now());
        assert_eq!(val, *expected_val);
    }

    // Case: Ingest SST with non-overlapping range should fail
    // Generate New Data [DATA_SIZE, DATA_SIZE * 2)
    let new_kvs: Vec<_> = (DATA_SIZE..DATA_SIZE * 2)
        .map(|i| (i_to_tidb_key(i), i_to_val(i)))
        .collect();

    let region_cnt_before = pd_client.get_regions_number();
    // Split Region at key(DATA_SIZE)
    let split_key = encode_bytes(&i_to_tidb_key(DATA_SIZE));
    runtime
        .block_on(pd_client.split_regions(vec![split_key]))
        .unwrap();
    cluster.wait_pd_region_count(region_cnt_before + 1);
    cluster.wait_region_replicated(&i_to_tidb_key(0), NODES_SIZE);
    cluster.wait_region_replicated(&i_to_tidb_key(DATA_SIZE), NODES_SIZE);

    // Get the regions after split
    let key_in_left = encode_bytes(&i_to_tidb_key(0));
    let key_in_right = encode_bytes(&i_to_tidb_key(DATA_SIZE));

    let left_region = pd_client.get_region(&key_in_left).unwrap();
    let right_region = pd_client.get_region(&key_in_right).unwrap();
    assert_ne!(left_region.get_id(), right_region.get_id(),);

    let left_region_id = left_region.get_id();
    let left_epoch = left_region.get_region_epoch().clone();

    // Create SST for new data
    let mut body_buf = BytesMut::new();
    for (key, val) in &new_kvs {
        body_buf.put_u16_le(key.len() as u16);
        body_buf.put_slice(key);
        body_buf.put_u32_le(val.len() as u32);
        body_buf.put_slice(val);
    }
    let new_body_bytes = body_buf.freeze();

    // Send PUT /write_sst for new data
    let new_commit_ts = runtime.block_on(pd_client.get_tso()).unwrap().into_inner();
    let new_write_url = format!(
        "http://{}/write_sst?cluster_id={}&commit_ts={}",
        cluster.tikv_worker_endpoints()[0],
        cluster_id,
        new_commit_ts
    );

    let new_req = Request::builder()
        .method(Method::PUT)
        .uri(&new_write_url)
        .header("Content-Type", "application/octet-stream")
        .body(Body::from(new_body_bytes.clone()))
        .unwrap();

    let new_resp = runtime.block_on(http_client.request(new_req)).unwrap();
    assert_eq!(new_resp.status(), StatusCode::OK,);

    // Parse SstMeta from new response
    let new_bytes = runtime
        .block_on(hyper::body::to_bytes(new_resp.into_body()))
        .unwrap();
    let new_resp_body: JsonValue = serde_json::from_slice(&new_bytes).unwrap();
    let new_sst_meta_json = new_resp_body.get("sst_meta").unwrap();
    let new_sst_meta: SstMeta = serde_json::from_value(new_sst_meta_json.clone()).unwrap();

    // Try to Ingest New SST to Left Region (Failure Case)
    let node_id = nodes[0];
    let status_addr = cluster.status_addr(node_id);
    let ingest_url = format!(
        "http://{}/ingest_s3?cluster_id={}&region_id={}&epoch_version={}",
        status_addr,
        cluster_id,
        left_region_id,
        left_epoch.get_version()
    );
    let req_body = json!(new_sst_meta).to_string();

    let ingest_req = Request::builder()
        .method(Method::POST)
        .uri(&ingest_url)
        .header("Content-Type", "application/json")
        .body(Body::from(req_body))
        .unwrap();

    let ingest_resp = runtime.block_on(http_client.request(ingest_req)).unwrap();
    let status = ingest_resp.status();
    // Expect failure because the SST range [key(DATA_SIZE), key(DATA_SIZE*2-1)]
    // does not overlap with the left region range [start_key, key(DATA_SIZE))
    assert_ne!(status, StatusCode::OK,);

    // Cleanup
    cluster.stop();
    oss.shutdown();
}