// Copyright 2025 PingCAP, Inc.

use std::{sync::Arc, time::Duration};

use api_version::{
    ApiV2,
    api_v2::{DEFAULT_KEYSPACE_ID, KEYSPACE_PREFIX_LEN},
};
use bytes::{BufMut, Bytes, BytesMut};
use futures::StreamExt;
use http::{Request, StatusCode, Uri};
use hyper::Body;
use kvengine::{GLOBAL_SHARD_END_KEY, dfs::ReservableWriter, table::OwnedInnerKey};
use kvproto::{cdcpb::ChangeDataRequest, metapb::Region};
use log_wrappers::Value as LogValue;
use pd_client::RpcClient;
use security::{HttpClient, SecurityConfig, SecurityManager};
use tikv_util::{
    box_err,
    codec::bytes::{decode_bytes, encode_bytes},
    debug, error, info,
    time::{Instant, Limiter},
    warn,
};
use txn_types::TimeStamp;

use crate::{Error, ticdc_util, ticdc_util::TiCdcError};

pub(crate) const DISPATCH_CDC_TIMEOUT: Duration = Duration::from_secs(30);

#[derive(Clone)]
pub(crate) enum WriterWritePolicy {
    Inline,
    Blocking(Option<Limiter>),
}

const INLINE_WRITE_YIELD_THRESHOLD_BYTES: u64 = 4 * 1024 * 1024;
const BLOCKING_WRITE_CHANNEL_CAPACITY: usize = 16;

async fn read_hyper_body_to_writer_inline<W: ReservableWriter>(
    mut body: Body,
    timeout: Duration,
    mut writer: W,
    timeout_err: String,
) -> crate::Result<(W, u64 /* read_len */)> {
    let mut read_len: u64 = 0;
    let mut bytes_since_yield = 0;
    while let Some(res) = tokio::time::timeout(timeout, body.next())
        .await
        .map_err(|_| Error::StoreTimeout(timeout_err.clone()))?
    {
        let chunk = res?;
        let chunk_len = chunk.len() as u64;
        read_len += chunk_len;
        writer
            .write_all(&chunk)
            .map_err(|e| -> Error { box_err!("write response body failed: {:?}", e) })?;
        bytes_since_yield += chunk_len;
        if bytes_since_yield >= INLINE_WRITE_YIELD_THRESHOLD_BYTES {
            tokio::task::yield_now().await;
            bytes_since_yield = 0;
        }
    }
    writer
        .flush()
        .map_err(|e| -> Error { box_err!("flush response body failed: {:?}", e) })?;
    Ok((writer, read_len))
}

async fn read_hyper_body_to_writer_blocking<W>(
    mut body: Body,
    timeout: Duration,
    writer: W,
    speed_limiter: Option<Limiter>,
    timeout_err: String,
) -> crate::Result<(W, u64 /* read_len */)>
where
    W: ReservableWriter + Send + 'static,
{
    let (tx, rx) = tokio::sync::mpsc::channel::<Bytes>(BLOCKING_WRITE_CHANNEL_CAPACITY);
    let writer_handle = tokio::task::spawn_blocking(move || -> crate::Result<(W, u64)> {
        let mut writer = writer;
        let mut read_len: u64 = 0;
        let mut rx = rx;
        while let Some(chunk) = rx.blocking_recv() {
            read_len += chunk.len() as u64;
            writer
                .write_all(&chunk)
                .map_err(|e| -> Error { box_err!("write response body failed: {:?}", e) })?;
            if let Some(speed_limiter) = speed_limiter.as_ref() {
                speed_limiter.blocking_consume(chunk.len());
            }
        }
        writer
            .flush()
            .map_err(|e| -> Error { box_err!("flush response body failed: {:?}", e) })?;
        Ok((writer, read_len))
    });

    let mut read_err: Option<Error> = None;
    let mut send_failed = false;
    loop {
        let next = match tokio::time::timeout(timeout, body.next()).await {
            Ok(next) => next,
            Err(_) => {
                read_err = Some(Error::StoreTimeout(timeout_err.clone()));
                break;
            }
        };

        match next {
            None => break,
            Some(Ok(chunk)) => {
                if tx.send(chunk).await.is_err() {
                    // `send` fails only when the receiver is dropped, which indicates the
                    // blocking writer task has exited early (usually due to write error).
                    // Stop reading further to avoid unbounded buffering and await the writer
                    // task to return the actual error.
                    send_failed = true;
                    break;
                }
            }
            Some(Err(err)) => {
                read_err = Some(err.into());
                break;
            }
        }
    }

    drop(tx);
    let write_res = writer_handle
        .await
        .map_err(|e| -> Error { box_err!("write response body task join failed: {:?}", e) })?;

    match write_res {
        Ok((writer, read_len)) => {
            if let Some(err) = read_err {
                return Err(err);
            }
            if send_failed {
                return Err(box_err!(
                    "write response body channel unexpectedly closed without returning an error"
                ));
            }
            Ok((writer, read_len))
        }
        Err(err) => Err(err),
    }
}

async fn read_hyper_body_to_writer<W>(
    body: Body,
    timeout: Duration,
    writer: W,
    timeout_err: String,
    write_policy: WriterWritePolicy,
) -> crate::Result<(W, u64 /* read_len */)>
where
    W: ReservableWriter + Send + 'static,
{
    match write_policy {
        WriterWritePolicy::Inline => {
            read_hyper_body_to_writer_inline(body, timeout, writer, timeout_err).await
        }
        WriterWritePolicy::Blocking(speed_limiter) => {
            read_hyper_body_to_writer_blocking(body, timeout, writer, speed_limiter, timeout_err)
                .await
        }
    }
}

pub(crate) async fn send_request_to_store_with_writer<W, F>(
    req: Request<Body>,
    client: &HttpClient,
    timeout: Duration,
    write_policy: WriterWritePolicy,
    build_writer: F,
) -> crate::Result<(StatusCode, W, u64 /* read_len */)>
where
    W: ReservableWriter + Send + 'static,
    F: FnOnce() -> crate::Result<W>,
{
    debug_assert!(!timeout.is_zero());
    let uri_str = format!("{}", req.uri());
    let resp_res = tokio::time::timeout(timeout, client.request(req))
        .await
        .map_err(|_| Error::StoreTimeout(format!("send request to {uri_str}")))?;
    let resp = resp_res?;
    let status = resp.status();

    let mut writer = build_writer()?;
    if let Some(cap) = resp
        .headers()
        .get(http::header::CONTENT_LENGTH)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.parse::<u64>().ok())
    {
        writer.reserve_capacity(cap);
    }

    let (writer, read_len) = read_hyper_body_to_writer(
        resp.into_body(),
        timeout,
        writer,
        format!("read response from {uri_str}"),
        write_policy,
    )
    .await?;
    Ok((status, writer, read_len))
}

pub(crate) async fn send_request_to_store(
    req: Request<Body>,
    client: &HttpClient,
    timeout: Duration,
) -> crate::Result<(StatusCode, Bytes)> {
    let (status, writer, _) =
        send_request_to_store_with_writer(req, client, timeout, WriterWritePolicy::Inline, || {
            Ok(BytesMut::new().writer())
        })
        .await?;
    Ok((status, writer.into_inner().freeze()))
}

async fn dispatch_http_post_with_retry(
    client: &HttpClient,
    uri: &Uri,
    body: Bytes,
    timeout: Duration,
) -> crate::Result<(StatusCode, Bytes)> {
    let mut last_err: Option<Error> = None;
    let start_time = Instant::now_coarse();
    while start_time.saturating_elapsed() < timeout {
        match dispatch_http_post(client, uri, body.clone()).await {
            Ok(resp) => {
                return Ok(resp);
            }
            Err(err) => {
                last_err = Some(err);
                tokio::time::sleep(Duration::from_millis(500)).await;
            }
        }
    }
    Err(last_err.unwrap_or(box_err!("timeout")))
}

async fn dispatch_http_post(
    client: &HttpClient,
    uri: &Uri,
    req_body: Bytes,
) -> crate::Result<(StatusCode, Bytes)> {
    let req = http::Request::builder()
        .uri(uri)
        .method("POST")
        .body(req_body.into())
        .unwrap();
    let resp = client.request(req).await?;
    let status = resp.status();
    let body = hyper::body::to_bytes(resp.into_body()).await?;
    Ok((status, body))
}

pub(crate) async fn post_to_ticdc<F>(
    tag: &str,
    client: &HttpClient,
    uri: &Uri,
    body: Bytes,
    is_err_retryable: F,
) -> crate::Result<(StatusCode, Bytes)>
where
    F: Fn(&TiCdcError) -> bool,
{
    info!("{} post_to_ticdc", tag;
        "uri" => ?uri, "body" => String::from_utf8_lossy(&body).as_ref());
    let mut last_resp: Option<(StatusCode, Bytes)> = None;
    let start_time = Instant::now_coarse();
    while start_time.saturating_elapsed() < DISPATCH_CDC_TIMEOUT {
        let (status, resp) =
            dispatch_http_post_with_retry(client, uri, body.clone(), DISPATCH_CDC_TIMEOUT / 2)
                .await?;
        if !status.is_success() {
            let ticdc_err = ticdc_util::parse_ticdc_response(&resp);
            if is_err_retryable(&ticdc_err) {
                warn!("{} post_to_ticdc error", tag; "resp" => String::from_utf8_lossy(&resp).as_ref());
                last_resp = Some((status, resp));
                tokio::time::sleep(Duration::from_millis(500)).await;
                continue;
            }
        }
        info!("{} post_to_ticdc success", tag; "resp" => String::from_utf8_lossy(&resp).as_ref());
        return Ok((status, resp));
    }
    let (status, resp) = last_resp.ok_or(Error::TiCdcTimeout)?;
    error!("{} post_to_ticdc error", tag; "status" => ?status, "resp" => String::from_utf8_lossy(&resp).as_ref());
    Ok((status, resp))
}

pub(crate) async fn read_from_ticdc<F>(
    tag: &str,
    http_client: &HttpClient,
    uri: Uri,
    timeout: Duration,
    retry_interval: Duration,
    is_err_retryable: F,
) -> crate::Result<Bytes>
where
    F: Fn(&TiCdcError) -> bool,
{
    let mut last_err = None;
    let start_time = Instant::now_coarse();
    while start_time.saturating_elapsed() < timeout {
        let resp = match tokio::time::timeout(timeout / 2, http_client.get(uri.clone())).await {
            Err(_) => {
                warn!("{} read_from_ticdc timeout", tag; "uri" => ?uri);
                last_err = Some(Error::TiCdcTimeout);
                tokio::time::sleep(retry_interval).await;
                continue;
            }
            Ok(Err(e)) => {
                warn!("{} read_from_ticdc error: {:?}", tag, e; "uri" => ?uri);
                last_err = Some(e.into());
                tokio::time::sleep(retry_interval).await;
                continue;
            }
            Ok(Ok(resp)) => resp,
        };

        let status = resp.status();
        if !status.is_success() {
            let body = hyper::body::to_bytes(resp.into_body())
                .await
                .unwrap_or_default();
            let ticdc_err = ticdc_util::parse_ticdc_response(&body);
            if is_err_retryable(&ticdc_err) {
                warn!("{} read_from_ticdc error: {:?}", tag, ticdc_err; "uri" => ?uri);
                tokio::time::sleep(retry_interval).await;
                continue;
            } else {
                return Err(ticdc_err.into());
            }
        }

        match tokio::time::timeout(timeout / 2, hyper::body::to_bytes(resp.into_body())).await {
            Err(_) => {
                warn!("{} read_from_ticdc body timeout", tag; "uri" => ?uri);
                last_err = Some(Error::TiCdcTimeout);
                tokio::time::sleep(retry_interval).await;
                continue;
            }
            Ok(Err(e)) => {
                warn!("{} read_from_ticdc body error: {:?}", tag, e; "uri" => ?uri);
                last_err = Some(e.into());
                tokio::time::sleep(retry_interval).await;
                continue;
            }
            Ok(Ok(body)) => {
                debug!("{} read_from_ticdc success", tag;
                    "body" => String::from_utf8_lossy(&body).as_ref(), "uri" => ?uri);
                return Ok(body);
            }
        }
    }

    let err = last_err.unwrap_or(Error::TiCdcTimeout);
    error!("{} read_from_ticdc error: {:?}", tag, err; "uri" => ?uri);
    Err(err)
}

pub(crate) fn build_request_range(request: &ChangeDataRequest) -> (Vec<u8>, Vec<u8>) {
    let mut encoded_start_key = request.get_start_key();
    let start_key = decode_bytes(&mut encoded_start_key, false).unwrap_or_default();
    let end_key = if request.get_end_key().is_empty() {
        GLOBAL_SHARD_END_KEY.to_vec()
    } else {
        let mut encoded_end_key = request.get_end_key();
        decode_bytes(&mut encoded_end_key, false).unwrap_or_default()
    };
    (start_key, end_key)
}

pub(crate) fn build_request_range_for_keyspace(
    keyspace_id: u32,
    request: &ChangeDataRequest,
) -> (Vec<u8>, Vec<u8>) {
    let (start_key, end_key) = build_request_range(request);
    let mut prepended_start_key = ApiV2::get_keyspace_prefix_by_id(keyspace_id);
    prepended_start_key.extend_from_slice(&start_key);
    let mut prepended_end_key = ApiV2::get_keyspace_prefix_by_id(keyspace_id);
    prepended_end_key.extend_from_slice(&end_key);
    (prepended_start_key, prepended_end_key)
}

pub(crate) fn get_inner_bounds_from_region(
    region: &Region,
) -> Option<(OwnedInnerKey, OwnedInnerKey)> {
    fn decode_outer_encoded_key(encoded_key: &[u8]) -> Option<Vec<u8>> {
        if encoded_key.is_empty() {
            return Some(vec![]);
        }
        let mut encoded = encoded_key;
        decode_bytes(&mut encoded, false).ok()
    }

    fn get_outer_keyspace_id(outer_key: &[u8]) -> u32 {
        ApiV2::get_u32_keyspace_id_by_key(outer_key).unwrap_or(DEFAULT_KEYSPACE_ID)
    }

    fn get_outer_end_keyspace_id(outer_end_key: &[u8]) -> Option<u32> {
        if outer_end_key.is_empty() {
            // Empty end key do not belong to any keyspace.
            return None;
        }
        let keyspace_id = get_outer_keyspace_id(outer_end_key);
        if outer_end_key.len() == KEYSPACE_PREFIX_LEN
            && ApiV2::is_txn_key(outer_end_key)
            && keyspace_id != DEFAULT_KEYSPACE_ID
        {
            keyspace_id.checked_sub(1)
        } else {
            Some(keyspace_id)
        }
    }

    let outer_start_key = decode_outer_encoded_key(region.get_start_key())?;
    let outer_end_key = decode_outer_encoded_key(region.get_end_key())?;
    let start_keyspace_id = get_outer_keyspace_id(&outer_start_key);
    let end_keyspace_id = get_outer_end_keyspace_id(&outer_end_key)?;
    if start_keyspace_id != end_keyspace_id {
        return None;
    }
    Some((
        OwnedInnerKey::new(outer_start_key.into()),
        OwnedInnerKey::new_end_key(outer_end_key.into()),
    ))
}

/// Returns whether the two half-open ranges `[a_start, a_end)` and
/// `[b_start, b_end)` intersect.
///
/// An empty end key represents an infinite upper bound.
pub(crate) fn ranges_intersect(a_start: &[u8], a_end: &[u8], b_start: &[u8], b_end: &[u8]) -> bool {
    if (!a_end.is_empty() && a_end <= a_start) || (!b_end.is_empty() && b_end <= b_start) {
        return false;
    }

    let a_start_lt_b_end = b_end.is_empty() || a_start < b_end;
    let b_start_lt_a_end = a_end.is_empty() || b_start < a_end;
    a_start_lt_b_end && b_start_lt_a_end
}

/// Clamps request range `[req_start, req_end)` to bounds `[bound_start,
/// bound_end)`.
///
/// Empty `req_end` and `bound_end` are normalized to `GLOBAL_SHARD_END_KEY`
/// before clamping. This keeps range arithmetic explicit and aligned with
/// kvengine's canonical global upper bound.
pub(crate) fn clamp_range_to_bounds<'a>(
    req_start: &'a [u8],
    req_end: &'a [u8],
    bound_start: &'a [u8],
    bound_end: &'a [u8],
) -> (&'a [u8], &'a [u8]) {
    let req_end = if req_end.is_empty() {
        GLOBAL_SHARD_END_KEY
    } else {
        req_end
    };
    let bound_end = if bound_end.is_empty() {
        GLOBAL_SHARD_END_KEY
    } else {
        bound_end
    };

    if bound_end <= bound_start {
        return (bound_start, bound_start);
    }

    let mut start = req_start;
    if start < bound_start {
        start = bound_start;
    } else if start > bound_end {
        start = bound_end;
    }

    let mut end = req_end;
    if end < bound_start {
        end = bound_start;
    } else if end > bound_end {
        end = bound_end;
    }

    if end < start {
        end = start;
    }
    (start, end)
}

pub(crate) async fn new_keyspace_pd_client(
    pd_url: String,
    sec_conf: &SecurityConfig,
) -> Arc<RpcClient> {
    let sec_mgr = Arc::new(SecurityManager::new(sec_conf).unwrap());
    Arc::new(
        RpcClient::new_async(&pd_client::Config::new(vec![pd_url]), None, sec_mgr)
            .await
            .unwrap(),
    )
}

pub(crate) fn keyspace_prefix_len(keyspace_id: u32) -> usize {
    if keyspace_id > 0 {
        KEYSPACE_PREFIX_LEN
    } else {
        0
    }
}

pub(crate) fn trim_keyspace_prefix_for_encoded_key(
    mut region_key: &[u8],
) -> crate::Result<(u32 /* keyspace_id */, Vec<u8>)> {
    let Some(keyspace_id) = ApiV2::get_u32_keyspace_id_by_key(region_key) else {
        return Ok((DEFAULT_KEYSPACE_ID, region_key.to_vec()));
    };

    let raw_key = decode_bytes(&mut region_key, false)
        .map_err(|err| -> Error { box_err!("decode failed: {:?}", err) })?;
    let encoded_key = if raw_key.len() == KEYSPACE_PREFIX_LEN {
        vec![]
    } else {
        encode_bytes(&raw_key[KEYSPACE_PREFIX_LEN..])
    };
    Ok((keyspace_id, encoded_key))
}

pub(crate) fn trim_keyspace_prefix_for_encoded_end_key(
    region_key: &[u8],
) -> crate::Result<(u32 /* keyspace_id */, Vec<u8>)> {
    let (keyspace_id, key) = trim_keyspace_prefix_for_encoded_key(region_key)?;
    Ok(if key.is_empty() {
        if keyspace_id == DEFAULT_KEYSPACE_ID {
            (keyspace_id, key)
        } else {
            // For a keyspace end boundary, ApiV2 encodes it as the next keyspace's prefix,
            // i.e. end(keyspace=13) has keyspace_id=14 with empty suffix. Adjust it back.
            let prev_keyspace_id = keyspace_id
                .checked_sub(1)
                .ok_or_else(|| -> Error { box_err!("invalid keyspace: {}", keyspace_id) })?;
            (prev_keyspace_id, key)
        }
    } else {
        (keyspace_id, key)
    })
}

pub(crate) fn trim_keyspace_prefix_for_encoded_range(
    start_key: &[u8],
    end_key: &[u8],
) -> crate::Result<(u32 /* keyspace_id */, Vec<u8>, Vec<u8>)> {
    let (start_keyspace_id, trim_start_key) = trim_keyspace_prefix_for_encoded_key(start_key)?;
    let (end_keyspace_id, trim_end_key) = trim_keyspace_prefix_for_encoded_end_key(end_key)?;
    if start_keyspace_id != end_keyspace_id {
        return Err(box_err!(
            "keyspace id mismatch between start & end key: {} vs {}",
            LogValue::key(start_key),
            LogValue::key(end_key)
        ));
    }
    Ok((start_keyspace_id, trim_start_key, trim_end_key))
}

#[derive(Default)]
pub(crate) struct ResolvedTsStats {
    pub(crate) resolved_regions: usize,
    /// Count of regions are still scanning locks.
    pub(crate) unresolved_regions: usize,

    pub(crate) min_ts: TimeStamp,
    pub(crate) min_ts_region_id: u64,
}

impl ResolvedTsStats {
    pub(crate) fn record_resolved_region(&mut self, region_id: u64, resolved_ts: TimeStamp) {
        self.resolved_regions += 1;
        if self.min_ts.is_zero() || self.min_ts > resolved_ts {
            self.min_ts = resolved_ts;
            self.min_ts_region_id = region_id;
        }
    }

    pub(crate) fn record_unresolved_region(&mut self, _region_id: u64) {
        self.unresolved_regions += 1;
    }
}

#[cfg(test)]
mod tests {
    use std::{
        convert::Infallible,
        io,
        net::{SocketAddr, TcpListener},
        sync::Arc,
        time::Duration,
    };

    use api_version::{ApiV2, api_v2::DEFAULT_KEYSPACE_ID};
    use bytes::{BufMut, Bytes, BytesMut};
    use futures::stream;
    use hyper::{
        Body, Response, Server,
        service::{make_service_fn, service_fn},
    };
    use kvengine::dfs::ReservableWriter;
    use native_br::common::TempLocalObjectWriter;
    use tikv_util::codec::bytes::{decode_bytes, encode_bytes};

    use super::*;
    use crate::Error;

    fn start_test_server(
        handler: impl Fn(hyper::Request<Body>) -> Response<Body> + Send + Sync + 'static,
    ) -> (SocketAddr, tokio::sync::oneshot::Sender<()>) {
        let handler = Arc::new(handler);
        let make_svc = make_service_fn(move |_conn| {
            let handler = handler.clone();
            async move {
                Ok::<_, Infallible>(service_fn(move |req| {
                    let handler = handler.clone();
                    async move { Ok::<_, Infallible>((handler)(req)) }
                }))
            }
        });

        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
        let server = Server::from_tcp(listener).unwrap().serve(make_svc);
        let server = server.with_graceful_shutdown(async move {
            let _ = shutdown_rx.await;
        });
        tokio::spawn(server);
        (addr, shutdown_tx)
    }

    struct FailingWriter {
        fail_on_write: bool,
    }

    impl io::Write for FailingWriter {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            if self.fail_on_write {
                return Err(io::Error::other("write failed"));
            }
            Ok(buf.len())
        }

        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }

    impl ReservableWriter for FailingWriter {
        fn reserve_capacity(&mut self, _additional: u64) {}
    }

    #[tokio::test]
    async fn test_read_hyper_body_to_writer_blocking_temp_file_cleanup_on_timeout() {
        let dir = test_util::temp_dir(Some("read_hyper_body_to_writer_timeout"), false);
        let path = dir.path().join("tmp.wal");
        let writer = TempLocalObjectWriter::new(
            native_br::common::TempLocalObject::create(path.clone()).unwrap(),
        );
        let body = Body::wrap_stream(stream::unfold(0, |state| async move {
            match state {
                0 => Some((Ok::<Bytes, hyper::Error>(Bytes::from_static(b"abc")), 1)),
                1 => {
                    tokio::time::sleep(Duration::from_millis(200)).await;
                    Some((Ok::<Bytes, hyper::Error>(Bytes::from_static(b"def")), 2))
                }
                _ => None,
            }
        }));

        let res = read_hyper_body_to_writer(
            body,
            Duration::from_millis(50),
            writer,
            "read response body".to_string(),
            WriterWritePolicy::Blocking(None),
        )
        .await;
        assert!(matches!(res, Err(Error::StoreTimeout(_))));
        assert!(!path.exists());
    }

    #[rstest::rstest]
    #[case(WriterWritePolicy::Inline)]
    #[case(WriterWritePolicy::Blocking(None))]
    #[tokio::test]
    async fn test_send_request_to_store_with_writer_network_error(
        #[case] write_policy: WriterWritePolicy,
    ) {
        let client = security::HttpClient::Http(hyper::Client::new());
        let req = http::Request::get("http://127.0.0.1:0/")
            .body(Body::empty())
            .unwrap();
        let res = send_request_to_store_with_writer(
            req,
            &client,
            Duration::from_millis(200),
            write_policy,
            || Ok(BytesMut::new().writer()),
        )
        .await;
        assert!(matches!(
            res,
            Err(Error::SecurityHttpClientError(
                security::HttpClientError::Hyper(_)
            ))
        ));
    }

    #[rstest::rstest]
    #[case(WriterWritePolicy::Inline)]
    #[case(WriterWritePolicy::Blocking(None))]
    #[tokio::test]
    async fn test_send_request_to_store_with_writer_timeout(
        #[case] write_policy: WriterWritePolicy,
    ) {
        let (addr, shutdown_tx) = start_test_server(|_req| {
            let body = Body::wrap_stream(stream::unfold(0, |state| async move {
                match state {
                    0 => Some((Ok::<Bytes, hyper::Error>(Bytes::from_static(b"abc")), 1)),
                    1 => {
                        tokio::time::sleep(Duration::from_millis(200)).await;
                        Some((Ok::<Bytes, hyper::Error>(Bytes::from_static(b"def")), 2))
                    }
                    _ => None,
                }
            }));
            Response::new(body)
        });

        let client = security::HttpClient::Http(hyper::Client::new());
        let req = http::Request::get(format!("http://{addr}/"))
            .body(Body::empty())
            .unwrap();
        let res = send_request_to_store_with_writer(
            req,
            &client,
            Duration::from_millis(50),
            write_policy,
            || Ok(BytesMut::new().writer()),
        )
        .await;
        let _ = shutdown_tx.send(());

        assert!(matches!(res, Err(Error::StoreTimeout(_))));
    }

    #[rstest::rstest]
    #[case(WriterWritePolicy::Inline)]
    #[case(WriterWritePolicy::Blocking(None))]
    #[tokio::test]
    async fn test_send_request_to_store_with_writer_write_error(
        #[case] write_policy: WriterWritePolicy,
    ) {
        let (addr, shutdown_tx) =
            start_test_server(|_req| Response::new(Body::from(Bytes::from_static(b"abcdef"))));

        let client = security::HttpClient::Http(hyper::Client::new());
        let req = http::Request::get(format!("http://{addr}/"))
            .body(Body::empty())
            .unwrap();
        let res = send_request_to_store_with_writer(
            req,
            &client,
            Duration::from_secs(5),
            write_policy,
            || {
                Ok(FailingWriter {
                    fail_on_write: true,
                })
            },
        )
        .await;
        let _ = shutdown_tx.send(());

        assert!(matches!(res, Err(Error::OtherError(_))));
    }

    #[rstest::rstest]
    #[case(WriterWritePolicy::Inline)]
    #[case(WriterWritePolicy::Blocking(None))]
    #[tokio::test]
    async fn test_send_request_to_store_with_writer_success(
        #[case] write_policy: WriterWritePolicy,
    ) {
        let (addr, shutdown_tx) =
            start_test_server(|_req| Response::new(Body::from(Bytes::from_static(b"abcdef"))));

        let client = security::HttpClient::Http(hyper::Client::new());
        let req = http::Request::get(format!("http://{addr}/"))
            .body(Body::empty())
            .unwrap();
        let (status, writer, read_len) = send_request_to_store_with_writer(
            req,
            &client,
            Duration::from_secs(5),
            write_policy,
            || Ok(BytesMut::new().writer()),
        )
        .await
        .unwrap();
        let _ = shutdown_tx.send(());

        assert!(status.is_success());
        assert_eq!(read_len, 6);
        assert_eq!(writer.into_inner().freeze(), Bytes::from_static(b"abcdef"));
    }

    #[tokio::test]
    async fn test_send_request_to_store_with_writer_success_blocking_temp_file() {
        let (addr, shutdown_tx) =
            start_test_server(|_req| Response::new(Body::from(Bytes::from_static(b"abcdef"))));

        let client = security::HttpClient::Http(hyper::Client::new());
        let req = http::Request::get(format!("http://{addr}/"))
            .body(Body::empty())
            .unwrap();

        let dir = test_util::temp_dir(Some("send_request_to_store_with_writer_blocking"), false);
        let path = dir.path().join("tmp.wal");

        let (status, writer, read_len) = send_request_to_store_with_writer(
            req,
            &client,
            Duration::from_secs(5),
            WriterWritePolicy::Blocking(None),
            || {
                Ok(TempLocalObjectWriter::new(
                    native_br::common::TempLocalObject::create(path.clone()).unwrap(),
                ))
            },
        )
        .await
        .unwrap();
        let _ = shutdown_tx.send(());

        assert!(status.is_success());
        assert_eq!(read_len, 6);

        let obj = match writer.into_inner() {
            Ok(obj) => obj,
            Err(err) => panic!("extract TempLocalObject failed: {}", err.into_error()),
        };
        assert_eq!(std::fs::read(obj.path.as_ref()).unwrap(), b"abcdef");
        drop(obj);
        assert!(!path.exists());
    }

    fn encode_keyspace_key(keyspace_id: u32, suffix: &[u8]) -> Vec<u8> {
        let mut raw = ApiV2::get_keyspace_prefix_by_id(keyspace_id);
        raw.extend_from_slice(suffix);
        encode_bytes(&raw)
    }

    fn decode_encoded_key(encoded: &[u8]) -> Vec<u8> {
        let mut slice = encoded;
        decode_bytes(&mut slice, false).unwrap()
    }

    #[test]
    fn test_trim_keyspace_prefix_cases() {
        struct Case {
            name: &'static str,
            input: Vec<u8>,
            expected_keyspace_id: u32,
            expected_trimmed: Vec<u8>,
            expected_trimmed_raw: Vec<u8>,
        }

        let unprefixed = encode_bytes(b"t123");
        let cases = vec![
            Case {
                name: "non_default_keyspace_suffix",
                input: encode_keyspace_key(12, b"a"),
                expected_keyspace_id: 12,
                expected_trimmed: encode_bytes(b"a"),
                expected_trimmed_raw: b"a".to_vec(),
            },
            Case {
                name: "non_default_keyspace_suffix_multi_bytes",
                input: encode_keyspace_key(14, b"abc"),
                expected_keyspace_id: 14,
                expected_trimmed: encode_bytes(b"abc"),
                expected_trimmed_raw: b"abc".to_vec(),
            },
            Case {
                name: "non_default_keyspace_prefix_only",
                input: encode_keyspace_key(14, b""),
                expected_keyspace_id: 14,
                expected_trimmed: vec![],
                expected_trimmed_raw: vec![],
            },
            Case {
                name: "unprefixed_key_defaults_to_keyspace_0",
                input: unprefixed.clone(),
                expected_keyspace_id: DEFAULT_KEYSPACE_ID,
                expected_trimmed: unprefixed,
                expected_trimmed_raw: b"t123".to_vec(),
            },
        ];

        for case in cases {
            let (got_keyspace_id, got_trimmed) = trim_keyspace_prefix_for_encoded_key(&case.input)
                .unwrap_or_else(|err| panic!("case {}: {:?}", case.name, err));
            assert_eq!(
                got_keyspace_id, case.expected_keyspace_id,
                "case {}",
                case.name
            );
            assert_eq!(got_trimmed, case.expected_trimmed, "case {}", case.name);
            if case.expected_trimmed_raw.is_empty() {
                assert!(got_trimmed.is_empty(), "case {}", case.name);
            } else {
                assert_eq!(
                    decode_encoded_key(&got_trimmed),
                    case.expected_trimmed_raw,
                    "case {}",
                    case.name
                );
            }
        }
    }

    #[test]
    fn test_trim_keyspace_prefix_end_cases() {
        struct Case {
            name: &'static str,
            input: Vec<u8>,
            expected_keyspace_id: u32,
            expected_trimmed_raw: Vec<u8>,
        }

        let default_end = encode_bytes(&ApiV2::get_keyspace_end_by_id(DEFAULT_KEYSPACE_ID));
        let keyspace_id = 13;
        let non_default_end = encode_bytes(&ApiV2::get_keyspace_end_by_id(keyspace_id));
        let cases = vec![
            Case {
                name: "default_keyspace_end_boundary",
                input: default_end,
                expected_keyspace_id: DEFAULT_KEYSPACE_ID,
                expected_trimmed_raw: vec![],
            },
            Case {
                name: "non_default_keyspace_end_boundary",
                input: non_default_end,
                expected_keyspace_id: keyspace_id,
                expected_trimmed_raw: vec![],
            },
            Case {
                name: "regular_end_key",
                input: encode_keyspace_key(7, b"k"),
                expected_keyspace_id: 7,
                expected_trimmed_raw: b"k".to_vec(),
            },
        ];

        for case in cases {
            let (got_keyspace_id, got_trimmed) =
                trim_keyspace_prefix_for_encoded_end_key(&case.input)
                    .unwrap_or_else(|err| panic!("case {}: {:?}", case.name, err));
            assert_eq!(
                got_keyspace_id, case.expected_keyspace_id,
                "case {}",
                case.name
            );
            if case.expected_trimmed_raw.is_empty() {
                assert!(got_trimmed.is_empty(), "case {}", case.name);
            } else {
                assert_eq!(
                    decode_encoded_key(&got_trimmed),
                    case.expected_trimmed_raw,
                    "case {}",
                    case.name
                );
            }
        }
    }

    #[test]
    fn test_trim_keyspace_prefix_range_cases() {
        enum Expected {
            Ok {
                keyspace_id: u32,
                start_raw: Vec<u8>,
                end_raw: Vec<u8>,
            },
            ErrContains(&'static str),
        }

        struct Case {
            name: &'static str,
            start_key: Vec<u8>,
            end_key: Vec<u8>,
            expected: Expected,
        }

        fn expected_trimmed(raw: &[u8]) -> Vec<u8> {
            if raw.is_empty() {
                vec![]
            } else {
                encode_bytes(raw)
            }
        }

        let keyspace_id = 15;
        let cases = vec![
            Case {
                name: "accepts_keyspace_end_boundary",
                start_key: encode_keyspace_key(keyspace_id, b"a"),
                end_key: encode_bytes(&ApiV2::get_keyspace_end_by_id(keyspace_id)),
                expected: Expected::Ok {
                    keyspace_id,
                    start_raw: b"a".to_vec(),
                    end_raw: vec![],
                },
            },
            Case {
                name: "accepts_same_keyspace",
                start_key: encode_keyspace_key(16, b"a"),
                end_key: encode_keyspace_key(16, b"b"),
                expected: Expected::Ok {
                    keyspace_id: 16,
                    start_raw: b"a".to_vec(),
                    end_raw: b"b".to_vec(),
                },
            },
            Case {
                name: "rejects_keyspace_mismatch",
                start_key: encode_keyspace_key(20, b"a"),
                end_key: encode_keyspace_key(21, b"b"),
                expected: Expected::ErrContains("keyspace id mismatch"),
            },
        ];

        for case in cases {
            let got = trim_keyspace_prefix_for_encoded_range(&case.start_key, &case.end_key);
            match (got, case.expected) {
                (
                    Ok((got_keyspace_id, got_start, got_end)),
                    Expected::Ok {
                        keyspace_id,
                        start_raw,
                        end_raw,
                    },
                ) => {
                    assert_eq!(got_keyspace_id, keyspace_id, "case {}", case.name);
                    assert_eq!(
                        got_start,
                        expected_trimmed(&start_raw),
                        "case {}",
                        case.name
                    );
                    assert_eq!(got_end, expected_trimmed(&end_raw), "case {}", case.name);

                    assert_eq!(
                        decode_encoded_key(&got_start),
                        start_raw,
                        "case {}",
                        case.name
                    );
                    if end_raw.is_empty() {
                        assert!(got_end.is_empty(), "case {}", case.name);
                    } else {
                        assert_eq!(decode_encoded_key(&got_end), end_raw, "case {}", case.name);
                    }
                }
                (Err(err), Expected::ErrContains(expected)) => {
                    assert!(
                        err.to_string().contains(expected),
                        "case {}: err: {:?}",
                        case.name,
                        err
                    );
                }
                (Ok(got), _expected) => panic!("case {}: unexpected Ok({:?})", case.name, got),
                (Err(err), _expected) => panic!("case {}: unexpected Err({:?})", case.name, err),
            }
        }
    }

    #[test]
    fn test_get_inner_bounds_from_region_cases() {
        enum Expected {
            Some {
                start_inner: Vec<u8>,
                end_inner: Vec<u8>,
            },
            None,
        }

        struct Case {
            name: &'static str,
            start_key: Vec<u8>,
            end_key: Vec<u8>,
            expected: Expected,
        }

        let non_default_keyspace_id = 13;
        let cases = vec![
            Case {
                name: "default_keyspace",
                start_key: encode_bytes(b"a"),
                end_key: encode_bytes(b"z"),
                expected: Expected::Some {
                    start_inner: b"a".to_vec(),
                    end_inner: b"z".to_vec(),
                },
            },
            Case {
                name: "default_keyspace_end_boundary",
                start_key: encode_bytes(b"a"),
                end_key: encode_bytes(&ApiV2::get_keyspace_end_by_id(DEFAULT_KEYSPACE_ID)),
                expected: Expected::Some {
                    start_inner: b"a".to_vec(),
                    end_inner: GLOBAL_SHARD_END_KEY.to_vec(),
                },
            },
            Case {
                name: "default_keyspace_unbounded_end",
                start_key: encode_bytes(b"a"),
                end_key: vec![],
                expected: Expected::None,
            },
            Case {
                name: "non_default_keyspace",
                start_key: encode_keyspace_key(non_default_keyspace_id, b"a"),
                end_key: encode_keyspace_key(non_default_keyspace_id, b"z"),
                expected: Expected::Some {
                    start_inner: b"a".to_vec(),
                    end_inner: b"z".to_vec(),
                },
            },
            Case {
                name: "non_default_keyspace_end_boundary",
                start_key: encode_keyspace_key(non_default_keyspace_id, b"a"),
                end_key: encode_bytes(&ApiV2::get_keyspace_end_by_id(non_default_keyspace_id)),
                expected: Expected::Some {
                    start_inner: b"a".to_vec(),
                    end_inner: GLOBAL_SHARD_END_KEY.to_vec(),
                },
            },
            Case {
                name: "non_default_keyspace_unbounded_end",
                start_key: encode_keyspace_key(non_default_keyspace_id, b"a"),
                end_key: vec![],
                expected: Expected::None,
            },
            Case {
                name: "keyspace_mismatch",
                start_key: encode_keyspace_key(20, b"a"),
                end_key: encode_keyspace_key(21, b"z"),
                expected: Expected::None,
            },
            Case {
                name: "invalid_encoded_start_key",
                start_key: vec![0x1],
                end_key: encode_bytes(b"z"),
                expected: Expected::None,
            },
        ];

        for case in cases {
            let mut region = Region::default();
            region.set_start_key(case.start_key);
            region.set_end_key(case.end_key);
            let got = get_inner_bounds_from_region(&region);

            match (got, case.expected) {
                (
                    Some((start, end)),
                    Expected::Some {
                        start_inner,
                        end_inner,
                    },
                ) => {
                    assert_eq!(start.to_vec(), start_inner, "case {}", case.name);
                    assert_eq!(end.to_vec(), end_inner, "case {}", case.name);
                }
                (None, Expected::None) => {}
                (Some(got), Expected::None) => {
                    panic!("case {}: unexpected Some({:?})", case.name, got)
                }
                (None, Expected::Some { .. }) => panic!("case {}: unexpected None", case.name),
            }
        }
    }

    #[test]
    fn test_ranges_intersect_basic() {
        assert!(ranges_intersect(b"a", b"b", b"a", b"b"));
        assert!(ranges_intersect(b"a", b"c", b"b", b"d"));
        assert!(!ranges_intersect(b"a", b"b", b"b", b"c"));
        assert!(!ranges_intersect(b"b", b"c", b"a", b"b"));
    }

    #[test]
    fn test_ranges_intersect_infinite_end() {
        assert!(ranges_intersect(b"b", b"", b"a", b"c"));
        assert!(ranges_intersect(b"b", b"", b"b", b"c"));
        assert!(ranges_intersect(b"b", b"", b"c", b"d"));
        assert!(!ranges_intersect(b"b", b"c", b"c", b""));
    }

    #[test]
    fn test_clamp_range_to_bounds() {
        let cases: Vec<(&[u8], &[u8], &[u8], &[u8], &[u8], &[u8])> = vec![
            (b"", b"", b"m", b"n", b"m", b"n"),
            (b"a", b"z", b"m", b"n", b"m", b"n"),
            (b"a", b"c", b"m", b"n", b"m", b"m"),
            (b"z", b"zz", b"m", b"n", b"n", b"n"),
            (b"b", b"", b"a", b"", b"b", GLOBAL_SHARD_END_KEY),
        ];

        for (idx, (req_start, req_end, bound_start, bound_end, expected_start, expected_end)) in
            cases.into_iter().enumerate()
        {
            let (start, end) = clamp_range_to_bounds(req_start, req_end, bound_start, bound_end);
            assert_eq!(start, expected_start, "case {}", idx);
            assert_eq!(end, expected_end, "case {}", idx);
        }
    }
}