use std::{
    path::Path,
    sync::{
        Arc, LazyLock,
        atomic::{AtomicU64, Ordering},
    },
    time::{Duration, UNIX_EPOCH},
};

use anyhow::{Context, Result, anyhow, bail};
use storage_backend::storage::BackendStorage;
use tantivy::Index;
use tici_common::{cdc_file_name::CDCFileName, config::S3Config};
use tokio::time::sleep;

use crate::{
    ShardWriterMeta,
    config::WorkerNodeConfig,
    metadata::metadata_store::MetadataStore,
    utils::{
        cmp_doc, frag_index_path,
        frag_path_generator::{FragIDGenerator, FragPathGenerator},
        key::ShardingKeyRange,
        meta_client::MetaClient,
        meta_client_pool::MetaClientPool,
        schema::{IndexMeta, SortField},
    },
    writer::{
        cdc_file_poller::CDCFilePoller, cdc_file_prefix::CDCPrefix,
        frag_batch_uploader::FragBatchUploader, shard_writer_manager::ShardWriterManager,
    },
};
pub fn test_get_inf_range() -> ShardingKeyRange {
    let start = [];
    let end = [255; 1024]; // Use a large enough end to represent infinity, just for testing
    ShardingKeyRange::half_open(start.to_vec(), end.to_vec()).unwrap()
}

pub struct LocalTSGenerator {
    last_id: AtomicU64,
}

impl LocalTSGenerator {
    pub fn new() -> Self {
        Self {
            last_id: AtomicU64::new(0),
        }
    }

    pub fn gen_local_id(&self) -> u64 {
        let now = std::time::SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos() as u64;

        loop {
            let last_id = self.last_id.load(Ordering::Relaxed);
            let next_id = now.max(last_id.saturating_add(1));
            if self
                .last_id
                .compare_exchange(last_id, next_id, Ordering::SeqCst, Ordering::Relaxed)
                .is_ok()
            {
                return next_id;
            }
        }
    }
}

impl Default for LocalTSGenerator {
    fn default() -> Self {
        Self::new()
    }
}

static TEST_FRAG_ID_GENERATOR: LazyLock<Arc<LocalTSGenerator>> =
    LazyLock::new(|| Arc::new(LocalTSGenerator::new()));

fn test_frag_id_generator() -> Arc<LocalTSGenerator> {
    Arc::clone(&TEST_FRAG_ID_GENERATOR)
}

#[async_trait::async_trait]
impl FragIDGenerator for LocalTSGenerator {
    async fn gen_id(&self) -> anyhow::Result<u64> {
        Ok(self.gen_local_id())
    }
}

#[async_trait::async_trait]
impl FragIDGenerator for Arc<LocalTSGenerator> {
    async fn gen_id(&self) -> anyhow::Result<u64> {
        Ok(self.gen_local_id())
    }
}

pub fn test_new_frag_path_generator(frag_base_path: &str) -> Arc<FragPathGenerator> {
    Arc::new(FragPathGenerator::new(
        frag_base_path,
        Box::new(test_frag_id_generator()),
    ))
}

pub fn create_test_frag_path(base_dir: impl AsRef<Path>) -> Result<String> {
    let frag_id = test_frag_id_generator().gen_local_id();
    let frag_path = base_dir.as_ref().join(format!("f_{frag_id}"));
    Ok(frag_path
        .to_str()
        .with_context(|| format!("frag path to_str failed, path={}", frag_path.display()))?
        .to_string())
}

pub fn is_sorted_index(index: &Index, sorted_fields: &[SortField]) -> bool {
    let reader = index.reader().unwrap();
    let searcher = reader.searcher();
    let mut docs = Vec::new();
    let segments = searcher.segment_readers();
    for segment_reader in segments {
        let store_reader = segment_reader.get_store_reader(0).unwrap();
        let max_doc = segment_reader.max_doc();
        for doc_id in 0..max_doc {
            let doc = store_reader.get(doc_id).unwrap();
            docs.push(doc);
        }
    }
    cmp_doc::is_sorted(&docs, sorted_fields)
}

pub fn is_sorted_frag(frag_path: &str, index_meta: &IndexMeta) -> bool {
    let sorted_fields = index_meta.get_sort_fields();
    let idx_path = frag_index_path(frag_path);
    let index = Index::open_in_dir(&idx_path).unwrap();
    is_sorted_index(&index, sorted_fields)
}

async fn wait_shard_ready(
    mock_client: &mut Box<dyn MetaClient>,
    shard_id: u64,
    last_file: &str,
) -> Result<()> {
    let mut retries = 0;
    const MAX_RETRIES: u32 = 240; // 2 minutes
    loop {
        if let Some(progress) = mock_client.get_shard_progress(shard_id).await? {
            if *progress >= *last_file {
                break;
            }
        }
        sleep(Duration::from_millis(500)).await;
        retries += 1;
        if retries >= MAX_RETRIES {
            bail!("Shard {} not ready after {} retries", shard_id, retries);
        }
    }
    Ok(())
}

async fn find_last_cdc_file_for_test(poller: &CDCFilePoller, prefix: &CDCPrefix) -> Result<String> {
    let mut last_file = CDCFileName::empty();

    loop {
        let unread_files = poller
            .list_unread_files_for_table(&last_file, prefix)
            .await?;
        if unread_files.is_empty() {
            break;
        }

        // The poller intentionally returns at most one unread directory per
        // call, so test helpers that need the table tail must walk directory by
        // directory until the poller is exhausted.
        assert!(unread_files.len() <= 1000);
        let batch_last_file = unread_files.last().ok_or_else(|| {
            anyhow!("CDC poller returned an empty batch unexpectedly, prefix={prefix:?}")
        })?;
        last_file = CDCFileName::parse(batch_last_file)?;
    }

    if last_file.is_empty() {
        bail!("No unread files found for prefix {prefix:?}");
    }
    Ok(last_file.to_string())
}

pub async fn wait_index_ready(
    backend_storage: Arc<dyn BackendStorage>,
    bucket: &str,
    prefix: &str,
    shard_ids: &[u64],
    mut mock_client: Box<dyn MetaClient>,
) -> Result<()> {
    let poller = CDCFilePoller::new(backend_storage, bucket);
    let prefix = CDCPrefix::new(prefix.to_string());
    let last_file = find_last_cdc_file_for_test(&poller, &prefix).await?;

    for shard_id in shard_ids {
        wait_shard_ready(&mut mock_client, *shard_id, &last_file).await?;
    }

    Ok(())
}

/// Wait for each shard to reach its table's last CDC file.
/// For partition tables, each physical table has its own CDC prefix, so we
/// must use the last file per table prefix instead of a single global last
/// file.
pub async fn wait_index_ready_for_partition_table(
    backend_storage: Arc<dyn BackendStorage>,
    bucket: &str,
    shard_metas: &[ShardWriterMeta],
    mut mock_client: Box<dyn MetaClient>,
) -> Result<()> {
    let poller = CDCFilePoller::new(backend_storage, bucket);

    for shard_meta in shard_metas {
        let prefix = CDCPrefix::new(shard_meta.cdc_checkpoint.table_prefix.clone());
        let last_file = find_last_cdc_file_for_test(&poller, &prefix).await?;
        wait_shard_ready(
            &mut mock_client,
            shard_meta.shard_header.shard_id,
            &last_file,
        )
        .await?;
    }

    Ok(())
}

pub async fn write_test_data(
    config: &WorkerNodeConfig,
    shard_metas: &[ShardWriterMeta],
    index_meta: &IndexMeta,
    backend_storage: Arc<dyn BackendStorage>,
    mock_client_pool: Arc<dyn MetaClientPool>,
) -> Result<()> {
    assert!(!shard_metas.is_empty(), "No shard metas provided");

    let mut config = config.clone();
    config.s3 = S3Config::default_for_test();
    let metadata_store = Arc::new(MetadataStore::new(&config.storage.metadata_path())?);
    let frag_path_gen = test_new_frag_path_generator(&config.storage.frag_base_path());
    let shard_writer_manager = ShardWriterManager::new(
        config.clone(),
        backend_storage.clone(),
        mock_client_pool.clone(),
        metadata_store,
        frag_path_gen,
        Arc::new(FragBatchUploader::new(
            backend_storage.clone(),
            config.frag_upload.clone(),
        )),
    );

    let cdc_bucket = shard_metas[0].cdc_checkpoint.bucket.clone();
    let cdc_prefix = shard_metas[0].cdc_checkpoint.table_prefix.clone();
    let mut shard_ids = Vec::new();
    for shard_meta in shard_metas {
        assert_eq!(cdc_bucket, shard_meta.cdc_checkpoint.bucket);
        assert_eq!(cdc_prefix, shard_meta.cdc_checkpoint.table_prefix);
        shard_ids.push(shard_meta.shard_header.shard_id);

        shard_writer_manager
            .add_shard(shard_meta, index_meta)
            .await?;
    }

    let mock_client = mock_client_pool.get_meta_client().await?;
    wait_index_ready(
        backend_storage,
        &cdc_bucket,
        &cdc_prefix,
        &shard_ids,
        mock_client,
    )
    .await?;

    shard_writer_manager.shutdown().await;

    Ok(())
}

pub fn test_new_metadata_store(path: &str) -> Result<Arc<MetadataStore>> {
    let metadata_path = format!("{}/metadata", path);
    let store = MetadataStore::new(&metadata_path)?;
    Ok(Arc::new(store))
}