use std::{
    cmp::Ordering,
    sync::Arc,
    time::{Duration, Instant},
};

use anyhow::{Context, Result, bail, ensure};
use tantivy::{
    Index, Searcher, SegmentOrdinal, SegmentReader, index::SegmentId, query::AllQuery,
    schema::Schema,
};

use crate::{
    reader::{
        frag_reader::{build_custom_segment_readers, shared_searcher_segment_readers},
        merge_iter::{HeapEntry, MergeIter},
        seg_query_iter::{SegQueryIter, SegmentDocProvider},
    },
    utils::{
        cmp_doc::compare_docs, create_index_writer, dedup_bitset::FragAliveBitSetBytes,
        frag_index_path, schema::IndexMeta,
    },
};

/// Holds a Tantivy searcher together with the segment readers used for merge
/// iteration: either the index defaults or custom readers built from
/// obsolete-doc alive bitsets.
struct SearcherWithSegments {
    searcher: Searcher,
    segment_readers: Vec<Arc<SegmentReader>>,
}

fn create_searchers(
    frag_paths: &[String],
    obsolete_doc_bitsets: Vec<FragAliveBitSetBytes>,
) -> Result<(Option<Schema>, Vec<SearcherWithSegments>)> {
    ensure!(
        obsolete_doc_bitsets.is_empty() || obsolete_doc_bitsets.len() == frag_paths.len(),
        "obsolete_doc_bitsets count {} does not match input fragment count {}",
        obsolete_doc_bitsets.len(),
        frag_paths.len()
    );

    let mut searchers = Vec::with_capacity(frag_paths.len());
    let mut schema = None;
    let bitset_iter: Vec<_> = if obsolete_doc_bitsets.is_empty() {
        vec![None; frag_paths.len()]
    } else {
        obsolete_doc_bitsets.into_iter().map(Some).collect()
    };
    for (frag_path, obsolete_doc_bitset) in frag_paths.iter().zip(bitset_iter) {
        let index_path = frag_index_path(frag_path);
        let index = Index::open_in_dir(index_path.clone())
            .with_context(|| format!("Index::open_in_dir: {}", index_path))?;
        let index_reader = index
            .reader_builder()
            .reload_policy(tantivy::ReloadPolicy::Manual)
            .doc_store_cache_num_blocks(1)
            .try_into()
            .with_context(|| format!("Index::reader_builder: {}", index_path))?;

        if let Some(schema) = schema.as_ref() {
            if *schema != index.schema() {
                bail!("Schema mismatch for index_path={}", index_path);
            }
        } else {
            schema = Some(index.schema());
        }

        let searcher = index_reader.searcher();
        let segment_readers = if let Some(bitset) = obsolete_doc_bitset {
            build_custom_segment_readers(&index, &searcher, frag_path, bitset)?
        } else {
            shared_searcher_segment_readers(&searcher)
        };
        searchers.push(SearcherWithSegments {
            searcher,
            segment_readers,
        });
    }
    Ok((schema, searchers))
}

fn collect_seg_iter(
    searchers: &[SearcherWithSegments],
    query: &AllQuery,
) -> Result<Vec<SegQueryIter>> {
    let mut seg_iters = Vec::new();
    for searcher in searchers {
        let doc_provider: Arc<dyn SegmentDocProvider> = Arc::new(searcher.searcher.clone());
        for (i, seg_reader) in searcher.segment_readers.iter().enumerate() {
            let iter = SegQueryIter::new(
                i as SegmentOrdinal,
                doc_provider.clone(),
                seg_reader.clone(),
                query,
            )?;
            seg_iters.push(iter);
        }
    }
    Ok(seg_iters)
}

#[derive(Debug, Default)]
pub struct RebuildIndexStat {
    pub seg_count: u32,
    pub doc_count: u64,
    pub rebuild_dur: Duration,
}

pub fn rebuild_index(
    input_frag_paths: &[String],
    index_meta: &IndexMeta,
    output_frag_path: &str,
    index_num_threads: usize,
    index_mem_budget: usize,
    obsolete_doc_bitsets: Vec<FragAliveBitSetBytes>,
) -> Result<(Vec<SegmentId>, RebuildIndexStat)> {
    let mut stat = RebuildIndexStat::default();

    let (input_schema, searchers) = create_searchers(input_frag_paths, obsolete_doc_bitsets)?;
    let input_schema = input_schema.context("No schema found from input fragments")?;
    let output_schema = index_meta.get_tantivy_schema();
    ensure!(
        input_schema == *output_schema,
        "Schema mismatch between input fragments and index_meta"
    );

    let query = AllQuery {};
    let seg_iters = collect_seg_iter(&searchers, &query)?;
    stat.seg_count = seg_iters.len() as u32;
    let sort_fields = index_meta.get_sort_fields();
    let cmp = |a: &HeapEntry, b: &HeapEntry| -> Ordering {
        compare_docs(b.doc.as_ref(), a.doc.as_ref(), sort_fields)
    };
    let mut merge_iter = MergeIter::new(seg_iters, cmp)?;

    let watcher = Instant::now();
    let index_path = frag_index_path(output_frag_path);
    let mut index = Index::create_in_dir(&index_path, output_schema.clone())?;
    index.set_tokenizers(index_meta.get_analyzer_manager().clone());
    let mut index_writer = create_index_writer(
        &index,
        index_num_threads,
        index_mem_budget,
        output_frag_path,
    )?;
    while let Some(doc) = merge_iter.next_doc()? {
        let mut doc = match Arc::try_unwrap(doc) {
            Ok(doc) => doc,
            Err(doc) => doc.as_ref().clone(),
        };
        index_meta.build_fulltext_fields(&mut doc)?;
        index_writer.add_document(doc).with_context(|| {
            format!("index add_document failed. frag_path={}", output_frag_path)
        })?;
        stat.doc_count += 1;
    }
    index_writer
        .commit()
        .with_context(|| format!("index commit failed. frag_path={}", output_frag_path))?;
    index_writer.wait_merging_threads().with_context(|| {
        format!(
            "wait_merging_threads failed. frag_path={}",
            output_frag_path
        )
    })?;
    stat.rebuild_dur = watcher.elapsed();

    let seg_ids = index.searchable_segment_ids().with_context(|| {
        format!(
            "searchable_segment_ids failed. frag_path={}",
            output_frag_path
        )
    })?;
    Ok((seg_ids, stat))
}

#[cfg(test)]
mod tests {
    use std::{fs, path::Path};

    use anyhow::Result;
    use tantivy::{TantivyDocument, collector::DocSetCollector, query::AllQuery, schema::Value};
    use tempfile::TempDir;

    use super::*;
    use crate::{
        test_utils::{TEST_FLOAT_INDEX_META, TEST_INDEX_META_OVERLAP},
        utils::{
            cmp_doc::sort_docs,
            dedup_bitset::FragDocBitSet,
            schema::{HANDLE_STORE_NAME, VERSION_STORE_NAME},
        },
    };

    struct MergeTestDoc {
        handle: String,
        name: String,
        description: String,
        version: u64,
    }

    impl MergeTestDoc {
        fn new(id: usize, description: String) -> Self {
            let handle = format!("handle_{id:06}");
            Self {
                name: format!("name_{id:06}"),
                description,
                handle,
                version: 1,
            }
        }
    }

    fn get_schema_field(schema: &Schema, field_name: &str) -> Result<tantivy::schema::Field> {
        schema
            .get_field(field_name)
            .with_context(|| format!("Field {} not found in schema", field_name))
    }

    fn build_test_doc(index_meta: &IndexMeta, test_doc: &MergeTestDoc) -> Result<TantivyDocument> {
        let schema = index_meta.get_tantivy_schema();
        let mut doc = TantivyDocument::default();
        doc.add_bytes(
            get_schema_field(schema, HANDLE_STORE_NAME)?,
            test_doc.handle.as_bytes(),
        );
        doc.add_u64(
            get_schema_field(schema, VERSION_STORE_NAME)?,
            test_doc.version,
        );
        doc.add_text(get_schema_field(schema, "column_1")?, &test_doc.name);
        doc.add_text(get_schema_field(schema, "column_2")?, &test_doc.description);
        doc.add_text(get_schema_field(schema, "column_3")?, &test_doc.handle);
        doc.add_text(get_schema_field(schema, "column_4")?, &test_doc.handle);
        Ok(doc)
    }

    fn write_input_frag(
        parent_dir: &Path,
        frag_name: &str,
        index_meta: &IndexMeta,
        docs: &[MergeTestDoc],
    ) -> Result<String> {
        let frag_path = parent_dir.join(frag_name);
        let frag_path = frag_path
            .to_str()
            .context("frag path to_str failed")?
            .to_string();
        let index_path = frag_index_path(&frag_path);
        fs::create_dir_all(&index_path)
            .with_context(|| format!("Failed to create index directory idx_path={}", index_path))?;

        let mut index = Index::create_in_dir(&index_path, index_meta.get_tantivy_schema().clone())
            .with_context(|| format!("Failed to create index idx_path={}", index_path))?;
        index.set_tokenizers(index_meta.get_analyzer_manager().clone());
        let mut index_writer = create_index_writer(&index, 1, 50_000_000, &frag_path)?;

        let tantivy_docs = docs
            .iter()
            .map(|doc| build_test_doc(index_meta, doc))
            .collect::<Result<Vec<_>>>()?;
        let tantivy_docs = sort_docs(tantivy_docs, index_meta.get_sort_fields());
        for doc in tantivy_docs {
            index_writer.add_document(doc)?;
        }
        index_writer.commit()?;
        index_writer.wait_merging_threads()?;
        Ok(frag_path)
    }

    fn prepare_output_frag(parent_dir: &Path, frag_name: &str) -> Result<String> {
        let frag_path = parent_dir.join(frag_name);
        let frag_path = frag_path
            .to_str()
            .context("output frag path to_str failed")?
            .to_string();
        fs::create_dir_all(frag_index_path(&frag_path))
            .with_context(|| format!("Failed to create output index dir for {}", frag_path))?;
        Ok(frag_path)
    }

    #[test]
    fn test_rebuild_index_rejects_schema_mismatch() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let input_frag = write_input_frag(
            temp_dir.path(),
            "input",
            &TEST_INDEX_META_OVERLAP,
            &[MergeTestDoc::new(1, "small text".to_string())],
        )?;
        let output_frag = prepare_output_frag(temp_dir.path(), "output")?;

        let err = rebuild_index(
            &[input_frag],
            &TEST_FLOAT_INDEX_META,
            &output_frag,
            1,
            50_000_000,
            Vec::new(),
        )
        .expect_err("schema mismatch should fail");
        assert!(
            err.to_string()
                .contains("Schema mismatch between input fragments and index_meta"),
            "err={err:?}"
        );
        Ok(())
    }

    #[test]
    fn test_rebuild_index_skips_docs_from_obsolete_doc_bitsets() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let input_frag = write_input_frag(
            temp_dir.path(),
            "input",
            &TEST_INDEX_META_OVERLAP,
            &[
                MergeTestDoc::new(1, "first".to_string()),
                MergeTestDoc::new(2, "second".to_string()),
            ],
        )?;
        let output_frag = prepare_output_frag(temp_dir.path(), "output")?;

        let input_index = Index::open_in_dir(frag_index_path(&input_frag))?;
        let input_searcher = input_index.reader()?.searcher();
        let mut input_doc_addrs = input_searcher
            .search(&AllQuery, &DocSetCollector)?
            .into_iter()
            .collect::<Vec<_>>();
        input_doc_addrs.sort();
        assert_eq!(input_doc_addrs.len(), 2);

        let input_schema = input_index.schema();
        let handle_field = get_schema_field(&input_schema, HANDLE_STORE_NAME)?;
        let deleted_handle = input_searcher
            .doc::<TantivyDocument>(input_doc_addrs[0])?
            .get_first(handle_field)
            .and_then(|value| value.as_bytes())
            .context("deleted handle is None")?
            .to_vec();
        let expected_handle = input_searcher
            .doc::<TantivyDocument>(input_doc_addrs[1])?
            .get_first(handle_field)
            .and_then(|value| value.as_bytes())
            .context("expected handle is None")?
            .to_vec();

        let mut obsolete_doc_bitset = FragDocBitSet::empty_for_index(&input_index)?;
        obsolete_doc_bitset.push_deleted_doc(input_doc_addrs[0])?;
        obsolete_doc_bitset.sort_and_dedup();
        let obsolete_doc_bitsets =
            vec![FragAliveBitSetBytes::from_doc_bitset(obsolete_doc_bitset)?];

        rebuild_index(
            &[input_frag],
            &TEST_INDEX_META_OVERLAP,
            &output_frag,
            1,
            50_000_000,
            obsolete_doc_bitsets,
        )?;

        let output_index = Index::open_in_dir(frag_index_path(&output_frag))?;
        let output_searcher = output_index.reader()?.searcher();
        let mut output_doc_addrs = output_searcher
            .search(&AllQuery, &DocSetCollector)?
            .into_iter()
            .collect::<Vec<_>>();
        output_doc_addrs.sort();
        assert_eq!(output_doc_addrs.len(), 1);

        let output_schema = output_index.schema();
        let output_handle_field = get_schema_field(&output_schema, HANDLE_STORE_NAME)?;
        let output_handle = output_searcher
            .doc::<TantivyDocument>(output_doc_addrs[0])?
            .get_first(output_handle_field)
            .and_then(|value| value.as_bytes())
            .context("output handle is None")?
            .to_vec();
        assert_eq!(output_handle, expected_handle);
        assert_ne!(output_handle, deleted_handle);
        Ok(())
    }
}