use std::{collections::HashSet, sync::Arc};

use anyhow::{Result, bail};
use either::Either;
use log::info;
use tici::meta::reader::UnifiedShardInfo;
use tici_common::index::IndexKey;
use tici_grpc::{
    ErrorCode, GetShardLocalCacheRequest, GetShardLocalCacheResponse, GetShardManifestRequest,
    GetShardManifestResponse, KeyRange,
};
use tici_search_lib::reader_server::ReaderServer;
use tokio::sync::RwLock;

use crate::test_utils::{shard_cache_scaffold::MockReader, test_suite::E2ETestSuite};

impl E2ETestSuite {
    pub async fn send_get_shard_local_cache_request(
        &self,
        index_key: &IndexKey,
        key_ranges: Vec<KeyRange>,
    ) -> Result<tici_grpc::GetShardLocalCacheResponse> {
        let request = tici_grpc::GetShardLocalCacheRequest {
            keyspace_id: index_key.keyspace_id,
            table_id: index_key.table_id,
            index_id: index_key.index_id,
            key_ranges,
            limit: 0,
        };

        let client = self.get_meta_client();
        let response = client
            .clone()
            .get_shard_local_cache_info(request)
            .await
            .unwrap()
            .into_inner();

        Ok(response)
    }

    pub fn check_valid_shard_allocation(
        &self,
        response: &tici_grpc::GetShardLocalCacheResponse,
        expected_status: i32,
    ) -> Result<()> {
        if response.status != expected_status {
            bail!(
                "request should return status {}, got {}, response={:?}",
                expected_status,
                response.status,
                response
            );
        }

        if expected_status == 0 {
            if response.shard_local_cache_infos.is_empty() {
                bail!(
                    "shard_local_cache_infos should not be empty, response={:?}",
                    response
                );
            }
            for response_shard_info in response.shard_local_cache_infos.iter() {
                if response_shard_info.local_cache_addrs.is_empty() {
                    bail!(
                        "shard {:?} should have cache addresses",
                        response_shard_info
                    );
                }

                // Reported addrs in response_shard_info should exits.
                for cache_addr in response_shard_info.local_cache_addrs.iter() {
                    // TODO(P1) support real
                    if !(self
                        .mocked_readers
                        .iter()
                        .any(|x| &x.1.tiflash_addr == cache_addr)
                        || self.readers.iter().any(|x| &x.1.tiflash_addr == cache_addr))
                    {
                        bail!(
                            "want={}, found_mocked={:?}, found_readers={:?}, resp={:?}",
                            cache_addr,
                            self.mocked_readers.keys(),
                            self.readers.keys(),
                            response,
                        );
                    }
                }
            }
        }
        Ok(())
    }

    pub fn get_reader(&self, reader: &str) -> Either<MockReader, Arc<RwLock<ReaderServer>>> {
        if let Some(x) = self.readers.get(reader) {
            return Either::Right(x.reader.clone());
        }
        if let Some(x) = self.mocked_readers.get(reader) {
            return Either::Left(x.clone());
        }
        panic!("not found")
    }

    pub fn must_get_real_reader(&self, reader: &str) -> Arc<RwLock<ReaderServer>> {
        self.readers.get(reader).unwrap().reader.clone()
    }

    pub fn must_get_mock_reader(&self, reader: &str) -> MockReader {
        self.mocked_readers.get(reader).cloned().unwrap()
    }

    pub async fn check_shard_allocated_on_reader_node(
        &self,
        shard_id: &u64,
        expected_reader: &str,
    ) -> bool {
        let shard = self
            .get_meta_service()
            .shard_cache_manager
            .get_shard_by_id(shard_id)
            .await
            .unwrap();
        if !shard.has_ready_readers().await {
            return false;
        }
        let mut found = false;
        for reader in shard.get_reader_addrs().await {
            if reader.as_str() == expected_reader {
                found = true;
            }
        }
        if !found {
            info!(
                "expected_reader {} not found in shard's readers",
                expected_reader
            );
            return false;
        }
        self.get_cached_shard_ids(expected_reader)
            .await
            .contains(shard_id)
    }

    pub async fn must_shard_allocated_on_reader_node(&self, shard_id: &u64, expected_reader: &str) {
        assert!(
            self.check_shard_allocated_on_reader_node(shard_id, expected_reader)
                .await
        );
    }

    pub async fn get_uni_shard_by_id(&self, shard_id: &u64) -> Option<UnifiedShardInfo> {
        self.get_meta_service()
            .shard_cache_manager
            .get_shard_by_id(shard_id)
            .await
    }

    pub async fn check_get_local_cache_info(&self, index_key: &IndexKey) -> Result<()> {
        let mut client = self.get_meta_client();
        let req = GetShardLocalCacheRequest {
            keyspace_id: index_key.keyspace_id,
            table_id: index_key.table_id,
            index_id: index_key.index_id,
            key_ranges: vec![KeyRange {
                start_key: vec![],
                end_key: vec![],
            }],
            limit: 100,
        };
        let resp: GetShardLocalCacheResponse = client
            .get_shard_local_cache_info(req)
            .await
            .unwrap()
            .into_inner();
        if resp.shard_local_cache_infos.is_empty() {
            bail!(
                "shard_local_cache_infos should not be empty, resp={:?}",
                resp
            );
        }
        Ok(())
    }

    pub async fn check_get_shard_manifest(
        &self,
        index_key: &IndexKey,
        shard_id: u64,
        expected_seq: u64,
        expected_frag_paths: Vec<String>,
    ) -> bool {
        let mut client = self.get_meta_internal_client();
        let req = GetShardManifestRequest {
            table_id: index_key.table_id,
            index_id: index_key.index_id,
            keyspace_id: index_key.keyspace_id,
            shard_id,
        };
        let resp: GetShardManifestResponse =
            client.get_shard_manifest(req).await.unwrap().into_inner();
        if resp.status != ErrorCode::Success as i32 {
            return false;
        }
        if resp.seq != expected_seq {
            return false;
        }
        let mut expected = HashSet::new();
        for x in resp.frag_metas.iter() {
            expected.insert(x.frag_path.clone());
        }
        if expected.len() != expected_frag_paths.len() {
            return false;
        }
        for x in expected_frag_paths.iter() {
            if !expected.contains(x) {
                return false;
            }
        }
        true
    }

    pub async fn get_cached_shard_ids(&self, addr: &str) -> Vec<u64> {
        match self.get_reader(addr) {
            Either::Left(mock_reader) => mock_reader.get_cached_shard_ids().await,
            Either::Right(reader) => reader
                .read()
                .await
                .get_service()
                .shard_mirror_manager
                .get_cached_shard_ids(),
        }
    }
}