pub mod config;
mod estimate;
pub mod heartbeat;
pub mod metrics;
pub mod reader_server;
pub mod searcher;

use std::{
    error::Error,
    fmt,
    sync::{Arc, OnceLock},
};

use log::{error, info};
use prometheus::{Encoder, TextEncoder};
use tici_common::{
    codec::{codec_v2::add_keyspace_prefix, parse_table_id_from_ranges},
    index::IndexKey as CommonIndexKey,
};
use tici_shard::{
    reader::{
        query_info::QueryInfo,
        shard_mirror_manager::ShardMirrorManager,
        shard_reader::{ShardReader, ShardSearchStream},
    },
    utils::{key::ShardingKeyRange, shard_tracing_id},
};

use crate::{
    ffi::{
        CountResult, EstimateCountResult, Range, SearchParam, SearchResult, SearchStreamBatch,
        Shard,
    },
    searcher::search_shard_reader,
};

#[cxx::bridge]
pub mod ffi {
    #[derive(Debug, Default)]
    pub struct I64ColumnData {
        col_name: String,
        values: Vec<i64>,
        null_map: Vec<u8>,
    }

    #[derive(Debug, Default)]
    pub struct U64ColumnData {
        col_name: String,
        values: Vec<u64>,
        null_map: Vec<u8>,
    }

    #[derive(Debug, Default)]
    pub struct F64ColumnData {
        col_name: String,
        values: Vec<f64>,
        null_map: Vec<u8>,
    }

    #[derive(Debug, Default)]
    pub struct BytesColumnData {
        col_name: String,
        chars: Vec<u8>,
        offsets: Vec<u64>,
        null_map: Vec<u8>,
    }

    #[derive(Default)]
    pub struct SearchResult {
        row_count: u64,
        i64_columns: Vec<I64ColumnData>,
        u64_columns: Vec<U64ColumnData>,
        f64_columns: Vec<F64ColumnData>,
        bytes_columns: Vec<BytesColumnData>,
    }

    #[derive(Default)]
    pub struct SearchStreamBatch {
        result: SearchResult,
        is_eof: bool,
    }

    #[derive(Default)]
    pub struct CountResult {
        count: u64,
    }

    #[derive(Default, Debug)]
    pub struct EstimateCountResult {
        // The estimated total count of this query on the available shards.
        estimated_total_count: u64,
        available_shards: u64,
        sampled_shards: u64,
        available_frags: u64,
        sampled_frags: u64,
        population_live_docs: u64,
        sampled_live_docs: u64,
    }

    #[derive(Debug, Default)]
    pub struct SearchParam {
        limit: usize,
        sort_field_names: Vec<String>,
        is_asc: Vec<bool>,
    }

    #[derive(Clone)]
    struct Range<'a> {
        start: &'a [u8],
        end: &'a [u8],
    }

    struct ShardWithRange<'a> {
        shard_id: u64,
        ranges: Vec<Range<'a>>,
    }

    #[derive(Debug)]
    struct Shard {
        keyspace_id: u32,
        index_id: i64,
        shard_id: u64,
        shard_epoch: u64,
    }

    #[derive(Debug)]
    pub struct Expr {
        tp: i32,
        val: Vec<u8>,
        children: Vec<Expr>,
        rpn_args_len: usize,
        sig: i32,
        field_type: i32,
    }

    // definition of Rust interface
    extern "Rust" {
        type ShardsSnapshot;
        type SearchStream;

        fn search(
            snapshot: &ShardsSnapshot,
            shard: &Shard,
            ranges: &Vec<Range>,
            return_fields: &Vec<String>,
            query: &Expr,
            search_param: &SearchParam,
            read_ts: &u64,
        ) -> Result<SearchResult>;

        fn open_search_stream(
            snapshot: &ShardsSnapshot,
            shard: &Shard,
            ranges: &Vec<Range>,
            return_fields: &Vec<String>,
            query: &Expr,
            search_param: &SearchParam,
            read_ts: &u64,
        ) -> Result<Box<SearchStream>>;

        fn next_search_stream_batch(
            stream: &mut SearchStream,
            batch_size: usize,
        ) -> Result<SearchStreamBatch>;

        fn count(
            snapshot: &ShardsSnapshot,
            shard: &Shard,
            ranges: &Vec<Range>,
            query: &Expr,
            read_ts: &u64,
        ) -> Result<CountResult>;

        fn estimate_count(
            keyspace_id: u32,
            shards: &Vec<ShardWithRange>,
            query: &Expr,
        ) -> Result<EstimateCountResult>;

        fn check_shards_and_acquire_snapshot(
            shards: &Vec<Shard>,
            result: &mut Vec<bool>,
        ) -> Result<Box<ShardsSnapshot>>;

        fn start_reader_server(
            conf_path: &str,
            read_node_addr: &str,
            service_addr: &str,
            pd_addr: &str,
        ) -> Result<()>;
        fn stop_reader_server();

        fn gather_prometheus_metrics() -> Vec<u8>;
    }
}

pub struct ShardsSnapshot {
    inner: tici_shard::reader::shard_mirror_manager::ShardsSnapshot,
}

pub struct SearchStream {
    inner: ShardSearchStream,
    shard_tracing_id: String,
    read_ts: u64,
}

pub(crate) fn build_read_ranges(
    keyspace_id: u32,
    ranges: &[Range],
) -> Result<Vec<ShardingKeyRange>, Box<dyn Error>> {
    ranges
        .iter()
        .map(|range| {
            ShardingKeyRange::half_open(
                add_keyspace_prefix(keyspace_id, range.start),
                add_keyspace_prefix(keyspace_id, range.end),
            )
        })
        .collect::<Result<Vec<_>, _>>()
        .map_err(|err| err.into())
}

pub(crate) fn build_query_info(
    shard_reader: &ShardReader,
    query: &crate::ffi::Expr,
    read_ranges: &[ShardingKeyRange],
    return_fields: &[String],
    search_param: &SearchParam,
) -> Result<QueryInfo, Box<dyn Error>> {
    Ok(shard_reader.parse_expr(
        &query.into(),
        return_fields,
        read_ranges,
        search_param.limit,
        search_param.sort_field_names.to_vec(),
        search_param.is_asc.to_vec(),
    )?)
}

impl<'a> fmt::Debug for Range<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "[{}, {})",
            hex::encode(self.start),
            hex::encode(self.end)
        )
    }
}

fn build_shard_index_key(shard: &Shard, ranges: &[Range<'_>]) -> anyhow::Result<CommonIndexKey> {
    let table_id = parse_table_id_from_ranges(ranges, |range| range.start)?;
    Ok(CommonIndexKey::new(
        shard.keyspace_id,
        table_id,
        shard.index_id,
    ))
}

pub fn start_reader_server(
    conf_path: &str,
    read_node_addr: &str,
    service_addr: &str,
    pd_addr: &str,
) -> Result<(), Box<dyn Error>> {
    reader_server::start_server(conf_path, read_node_addr, service_addr, pd_addr)
        .inspect_err(|e| error!("Failed to start reader server, err={:?}", e))?;
    Ok(())
}

pub fn stop_reader_server() {
    reader_server::stop_server();
}

#[allow(clippy::ptr_arg)]
fn check_shards_and_acquire_snapshot(
    shards: &Vec<Shard>,
    result: &mut Vec<bool>,
) -> Result<Box<ShardsSnapshot>, Box<dyn Error>> {
    let shard_ids = shards
        .iter()
        .map(|shard| shard.shard_id)
        .collect::<Vec<_>>();
    let manager = get_shard_mirror_manager()?;
    let (inner, hits) = manager.check_shards_and_acquire_snapshot(shard_ids.as_slice());
    result.clear();
    result.reserve(hits.len());
    result.extend(hits);
    Ok(Box::new(ShardsSnapshot { inner }))
}

// Called by both FFI and tests.
#[allow(clippy::ptr_arg)]
pub fn search_by_reader(
    shard_reader: ShardReader,
    shard: &Shard,
    ranges: &Vec<Range>,
    return_fields: &Vec<String>,
    query: &crate::ffi::Expr,
    search_param: &SearchParam,
    read_ts: &u64,
) -> Result<SearchResult, Box<dyn Error>> {
    let start = std::time::Instant::now();
    let index_key = build_shard_index_key(shard, ranges)
        .inspect_err(|e| error!("Failed to parse index key from search ranges, err={:?}", e))?;
    let shard_tracing_id = shard_tracing_id(&index_key, shard.shard_id);
    let read_ranges = build_read_ranges(shard.keyspace_id, ranges)
        .inspect_err(|e| error!("Failed to create handle ranges, err={:?}", e))?;
    let query_info = Arc::new(
        build_query_info(
            &shard_reader,
            query,
            &read_ranges,
            return_fields,
            search_param,
        )
        .inspect_err(|e| error!("Failed to parse expr, err={:?}", e))?,
    );
    let (res, stat) = search_shard_reader(&shard_reader, query_info.clone()).inspect_err(|e| {
        error!(
            "Search {} failed, read_ts={}, query={:?}, err={:?}",
            shard_tracing_id, read_ts, query_info, e
        )
    })?;
    info!(
        "Search {} success, read_ts={}, query={:?}, plan={}, returned={}, elapsed={}us, stat={:?}",
        shard_tracing_id,
        read_ts,
        query_info,
        stat.plan.as_str(),
        stat.docs.returned_rows,
        start.elapsed().as_micros(),
        stat
    );
    Ok(res)
}

// Called by both FFI and tests.
#[allow(clippy::ptr_arg)]
pub fn open_search_stream_by_reader(
    shard_reader: ShardReader,
    shard: &Shard,
    ranges: &Vec<Range>,
    return_fields: &Vec<String>,
    query: &crate::ffi::Expr,
    search_param: &SearchParam,
    read_ts: &u64,
) -> Result<SearchStream, Box<dyn Error>> {
    let index_key = build_shard_index_key(shard, ranges)
        .inspect_err(|e| error!("Failed to parse index key from search ranges, err={:?}", e))?;
    let shard_tracing_id = shard_tracing_id(&index_key, shard.shard_id);
    let read_ranges = build_read_ranges(shard.keyspace_id, ranges)
        .inspect_err(|e| error!("Failed to create handle ranges, err={:?}", e))?;
    let query_info = Arc::new(
        build_query_info(
            &shard_reader,
            query,
            &read_ranges,
            return_fields,
            search_param,
        )
        .inspect_err(|e| error!("Failed to parse expr, err={:?}", e))?,
    );
    let inner = shard_reader
        .open_search_stream(query_info.clone())
        .inspect_err(|e| {
            error!(
                "Open search stream {} failed, read_ts={}, query={:?}, err={:?}",
                shard_tracing_id, read_ts, query_info, e
            )
        })?;
    Ok(SearchStream {
        inner,
        shard_tracing_id,
        read_ts: *read_ts,
    })
}

// Called by both FFI and tests.
pub fn next_search_stream_batch(
    stream: &mut SearchStream,
    batch_size: usize,
) -> Result<SearchStreamBatch, Box<dyn Error>> {
    let batch = stream.inner.next_batch(batch_size).inspect_err(|e| {
        error!(
            "Search stream {} next batch failed, read_ts={}, query={:?}, err={:?}",
            stream.shard_tracing_id,
            stream.read_ts,
            stream.inner.query_info(),
            e
        )
    })?;
    if batch.is_eof {
        let stat = stream.inner.stat_snapshot();
        info!(
            "Search stream {} success, read_ts={}, query={:?}, plan={}, returned={}, elapsed={}us, stat={:?}",
            stream.shard_tracing_id,
            stream.read_ts,
            stream.inner.query_info(),
            stat.plan.as_str(),
            stat.docs.returned_rows,
            stat.start.elapsed().as_micros(),
            stat
        );
    }
    Ok(SearchStreamBatch {
        result: batch.rows.into(),
        is_eof: batch.is_eof,
    })
}

// Called by both FFI and tests.
#[allow(clippy::ptr_arg)]
pub fn count_by_reader(
    shard_reader: ShardReader,
    shard: &Shard,
    ranges: &Vec<Range>,
    query: &crate::ffi::Expr,
    read_ts: &u64,
) -> Result<CountResult, Box<dyn Error>> {
    let start = std::time::Instant::now();
    let index_key = build_shard_index_key(shard, ranges)
        .inspect_err(|e| error!("Failed to parse index key from count ranges, err={:?}", e))?;
    let shard_tracing_id = shard_tracing_id(&index_key, shard.shard_id);
    let read_ranges = build_read_ranges(shard.keyspace_id, ranges)
        .inspect_err(|e| error!("Failed to create handle ranges, err={:?}", e))?;
    let return_fields = Vec::new();
    let search_param = SearchParam::default();
    let query_info = Arc::new(
        build_query_info(
            &shard_reader,
            query,
            &read_ranges,
            &return_fields,
            &search_param,
        )
        .inspect_err(|e| error!("Failed to parse expr, err={:?}", e))?,
    );
    let (count, stat) = shard_reader
        .count_with_query_info(query_info.clone())
        .inspect_err(|e| {
            error!(
                "Count {} failed, read_ts={}, query={:?}, err={:?}",
                shard_tracing_id, read_ts, query_info, e
            )
        })?;
    let res = CountResult { count };
    info!(
        "Count {} success, read_ts={}, query={:?}, plan={}, count={}, elapsed={}us, stat={:?}",
        shard_tracing_id,
        read_ts,
        query_info,
        stat.plan.as_str(),
        count,
        start.elapsed().as_micros(),
        stat
    );
    Ok(res)
}

#[allow(clippy::ptr_arg)]
pub fn estimate_count(
    keyspace_id: u32,
    shards: &Vec<ffi::ShardWithRange>,
    query: &crate::ffi::Expr,
) -> Result<EstimateCountResult, Box<dyn Error>> {
    let manager = get_shard_mirror_manager()?;
    estimate::estimate_count(manager, keyspace_id, shards, query)
}

static GLOBAL_SHARD_MIRROR_MANAGER: OnceLock<Arc<ShardMirrorManager>> = OnceLock::new();

fn set_global_shard_mirror_manager(
    shard_mirror_manager: Arc<ShardMirrorManager>,
) -> anyhow::Result<()> {
    GLOBAL_SHARD_MIRROR_MANAGER
        .set(shard_mirror_manager)
        .map_err(|_| anyhow::anyhow!("Failed to set global shard mirror manager"))
}

fn get_shard_mirror_manager() -> anyhow::Result<&'static ShardMirrorManager> {
    let manager = GLOBAL_SHARD_MIRROR_MANAGER
        .get()
        .ok_or_else(|| anyhow::anyhow!("Failed to get global shard mirror manager"))?;
    Ok(manager.as_ref())
}

fn get_shard_reader_from_snapshot(
    snapshot: &ShardsSnapshot,
    shard_id: u64,
) -> anyhow::Result<ShardReader> {
    Ok(snapshot
        .inner
        .get_shard(shard_id)
        .ok_or_else(|| anyhow::anyhow!("Shard {} not found", shard_id))?
        .get_reader())
}

// Called by FFI.
#[allow(clippy::ptr_arg)]
pub fn search(
    snapshot: &ShardsSnapshot,
    shard: &Shard,
    ranges: &Vec<Range>,
    return_fields: &Vec<String>,
    query: &crate::ffi::Expr,
    search_param: &SearchParam,
    read_ts: &u64,
) -> Result<SearchResult, Box<dyn Error>> {
    // Increment the request counter for this heartbeat cycle
    crate::metrics::increment_requests_in_cycle();
    let shard_reader = get_shard_reader_from_snapshot(snapshot, shard.shard_id)?;
    search_by_reader(
        shard_reader,
        shard,
        ranges,
        return_fields,
        query,
        search_param,
        read_ts,
    )
}

// Called by FFI.
#[allow(clippy::ptr_arg)]
pub fn open_search_stream(
    snapshot: &ShardsSnapshot,
    shard: &Shard,
    ranges: &Vec<Range>,
    return_fields: &Vec<String>,
    query: &crate::ffi::Expr,
    search_param: &SearchParam,
    read_ts: &u64,
) -> Result<Box<SearchStream>, Box<dyn Error>> {
    crate::metrics::increment_requests_in_cycle();
    let shard_reader = get_shard_reader_from_snapshot(snapshot, shard.shard_id)?;
    Ok(Box::new(open_search_stream_by_reader(
        shard_reader,
        shard,
        ranges,
        return_fields,
        query,
        search_param,
        read_ts,
    )?))
}

// Called by FFI.
#[allow(clippy::ptr_arg)]
pub fn count(
    snapshot: &ShardsSnapshot,
    shard: &Shard,
    ranges: &Vec<Range>,
    query: &crate::ffi::Expr,
    read_ts: &u64,
) -> Result<CountResult, Box<dyn Error>> {
    // Increment the request counter for this heartbeat cycle
    crate::metrics::increment_requests_in_cycle();
    let shard_reader = get_shard_reader_from_snapshot(snapshot, shard.shard_id)?;
    count_by_reader(shard_reader, shard, ranges, query, read_ts)
}

use protobuf::{ProtobufEnum, RepeatedField};
impl From<&crate::ffi::Expr> for tipb::Expr {
    fn from(expr: &crate::ffi::Expr) -> Self {
        let mut ret = tipb::Expr::default();
        ret.set_tp(tipb::ExprType::from_i32(expr.tp).unwrap_or(tipb::ExprType::Null));
        ret.set_val(expr.val.clone());
        let children = expr
            .children
            .iter()
            .map(|c| c.into())
            .collect::<Vec<tipb::Expr>>();
        let ret_children = RepeatedField::from_vec(children);
        ret.set_children(ret_children);
        ret.set_rpn_args_len(expr.rpn_args_len as u32);
        ret.set_sig(
            tipb::ScalarFuncSig::from_i32(expr.sig).unwrap_or(tipb::ScalarFuncSig::Unspecified),
        );
        ret.set_field_type(tipb::FieldType::new());
        ret
    }
}

pub fn gather_prometheus_metrics() -> Vec<u8> {
    let metric_families = prometheus::gather();
    let encoder = TextEncoder::new();
    let mut buffer = Vec::new();
    if let Err(e) = encoder.encode(&metric_families, &mut buffer) {
        error!("Failed to encode metrics, err={}", e);
        return Vec::new();
    }
    buffer
}

#[cfg(test)]
mod tests {
    use tici_common::codec::{table_end_key, table_start_key};

    use super::*;

    fn test_shard() -> Shard {
        Shard {
            keyspace_id: 7,
            index_id: 9,
            shard_id: 11,
            shard_epoch: 13,
        }
    }

    #[test]
    fn test_build_shard_index_key_uses_range_table_id() {
        let index_key = CommonIndexKey::new(7, 101, 9);
        let start = table_start_key(&index_key);
        let end = table_end_key(&index_key);
        let ranges = vec![Range {
            start: start.as_slice(),
            end: end.as_slice(),
        }];

        let actual = match build_shard_index_key(&test_shard(), &ranges) {
            Ok(index_key) => index_key,
            Err(err) => panic!("build_shard_index_key should succeed, err={err:?}"),
        };
        assert_eq!(actual, index_key);
    }

    #[test]
    fn test_build_shard_index_key_rejects_inconsistent_ranges() {
        let left = CommonIndexKey::new(7, 101, 9);
        let right = CommonIndexKey::new(7, 202, 9);
        let left_start = table_start_key(&left);
        let left_end = table_end_key(&left);
        let right_start = table_start_key(&right);
        let right_end = table_end_key(&right);
        let ranges = vec![
            Range {
                start: left_start.as_slice(),
                end: left_end.as_slice(),
            },
            Range {
                start: right_start.as_slice(),
                end: right_end.as_slice(),
            },
        ];

        let err = match build_shard_index_key(&test_shard(), &ranges) {
            Ok(index_key) => panic!("build_shard_index_key should fail, index_key={index_key:?}"),
            Err(err) => err.to_string(),
        };
        assert!(err.contains("Inconsistent table_id in ranges"));
    }

    #[test]
    fn test_build_shard_index_key_rejects_empty_ranges() {
        let err = match build_shard_index_key(&test_shard(), &[]) {
            Ok(index_key) => panic!("build_shard_index_key should fail, index_key={index_key:?}"),
            Err(err) => err.to_string(),
        };
        assert!(err.contains("Parse ranges failed"));
    }

    #[test]
    fn test_open_search_stream_rejects_count_query() {
        use std::sync::Arc;

        use anyhow::{Context, Result};
        use tici_shard::{
            reader::shard_reader::{EstimateSampleCounts, ShardReader},
            test_utils::{
                TEST_INDEX_KEY, TEST_INDEX_META_NORMAL, TestDoc, build_test_frag, load_frag_reader,
                test_get_inf_range,
            },
            utils::schema::VERSION_STORE_NAME,
        };
        use tipb::{ExprType, ScalarFuncSig};

        use crate::ffi::Expr;

        fn build_match_word_expr(word: &str, column: &str) -> Expr {
            Expr {
                tp: ExprType::ScalarFunc as i32,
                val: Vec::new(),
                children: vec![
                    Expr {
                        tp: ExprType::String as i32,
                        val: word.as_bytes().to_vec(),
                        children: Vec::new(),
                        rpn_args_len: 0,
                        sig: ScalarFuncSig::Unspecified as i32,
                        field_type: 0,
                    },
                    Expr {
                        tp: ExprType::String as i32,
                        val: column.as_bytes().to_vec(),
                        children: Vec::new(),
                        rpn_args_len: 0,
                        sig: ScalarFuncSig::Unspecified as i32,
                        field_type: 0,
                    },
                ],
                rpn_args_len: 0,
                sig: ScalarFuncSig::FtsMatchWord as i32,
                field_type: 0,
            }
        }

        fn build_count_query_expr() -> Expr {
            Expr {
                tp: ExprType::ScalarFunc as i32,
                val: Vec::new(),
                children: vec![
                    Expr {
                        tp: ExprType::String as i32,
                        val: b"testword".to_vec(),
                        children: Vec::new(),
                        rpn_args_len: 0,
                        sig: ScalarFuncSig::Unspecified as i32,
                        field_type: 0,
                    },
                    Expr {
                        tp: ExprType::String as i32,
                        val: b"column_1".to_vec(),
                        children: Vec::new(),
                        rpn_args_len: 0,
                        sig: ScalarFuncSig::Unspecified as i32,
                        field_type: 0,
                    },
                ],
                rpn_args_len: 0,
                sig: ScalarFuncSig::FtsMatchWord as i32,
                field_type: 0,
            }
        }

        fn run_test() -> Result<()> {
            let index_meta = TEST_INDEX_META_NORMAL.clone();
            let (temp_dir, meta) = build_test_frag(
                &index_meta,
                vec![vec![
                    TestDoc::normal(b"doc_00", 10, "alpha testword"),
                    TestDoc::normal(b"doc_01", 20, "beta testword"),
                ]],
            )
            .context("build test frag")?;
            let frag_reader = load_frag_reader(meta, &[]).context("load frag reader")?;
            let _guard = temp_dir;
            let make_shard_reader = || {
                ShardReader::new(
                    test_get_inf_range(),
                    index_meta.clone(),
                    1,
                    EstimateSampleCounts::default(),
                    Arc::new(vec![frag_reader.clone()]),
                )
            };
            let index_key = TEST_INDEX_KEY.clone();
            let shard = Shard {
                keyspace_id: index_key.keyspace_id,
                index_id: index_key.index_id,
                shard_id: 11,
                shard_epoch: 13,
            };
            let start = table_start_key(&index_key);
            let end = table_end_key(&index_key);
            let ranges = vec![Range {
                start: start.as_slice(),
                end: end.as_slice(),
            }];
            let return_fields = vec![VERSION_STORE_NAME.to_string()];
            let query = build_match_word_expr("testword", "column_1");
            let search_param = SearchParam::default();
            let read_ts = 0_u64;

            let mut stream = match open_search_stream_by_reader(
                make_shard_reader(),
                &shard,
                &ranges,
                &return_fields,
                &query,
                &search_param,
                &read_ts,
            ) {
                Ok(stream) => stream,
                Err(err) => anyhow::bail!("open rows search stream failed, err={err:?}"),
            };

            loop {
                let batch = match next_search_stream_batch(&mut stream, 2) {
                    Ok(batch) => batch,
                    Err(err) => anyhow::bail!("next search batch failed, err={err:?}"),
                };
                if batch.is_eof {
                    break;
                }
            }

            let one_shot = match search_by_reader(
                make_shard_reader(),
                &shard,
                &ranges,
                &return_fields,
                &query,
                &search_param,
                &read_ts,
            ) {
                Ok(result) => result,
                Err(err) => anyhow::bail!("one-shot search failed, err={err:?}"),
            };
            assert_eq!(
                stream.inner.stat_snapshot().docs.returned_rows,
                one_shot.row_count as u64
            );

            let count_query = build_count_query_expr();
            let err = match open_search_stream_by_reader(
                make_shard_reader(),
                &shard,
                &ranges,
                &Vec::new(),
                &count_query,
                &search_param,
                &read_ts,
            ) {
                Ok(_) => anyhow::bail!("count query unexpectedly opened a search stream"),
                Err(err) => err.to_string(),
            };
            assert!(
                err.contains(
                    "SearchExecutionPlan::FragScanCount is not supported in ShardSearchStream"
                ),
                "expected count stream error, got {err}"
            );
            Ok(())
        }

        match run_test() {
            Ok(()) => {}
            Err(err) => panic!("open search stream test failed, err={err:#}"),
        }
    }
}