use std::{fs, sync::Arc};

use anyhow::{Context, Result};
use tantivy::{DocAddress, Index};
use tempfile::TempDir;

use crate::{
    reader::{
        frag_reader::{FragReader, FragReaderMeta, FragRowsStream},
        query_info::QueryInfo,
        query_result::SearchRows,
        shard_reader::ShardReaderStat,
    },
    test_utils::helpers::create_test_frag_path,
    utils::{
        self, create_index_writer,
        dedup_bitset::{FragAliveBitSetBytes, FragDocBitSet},
        key::ShardingKey,
        schema::{IndexMeta, VERSION_STORE_NAME},
        tantivy_doc::{RawDoc, TantivyDoc},
    },
    writer::cdc_log::OperationType,
};

/// A single document used when building multi-segment test fragments.
#[derive(Clone)]
pub struct TestDoc {
    pub handle: Vec<u8>,
    pub version: u64,
    pub cols: Vec<(i64, Option<String>)>,
}

impl TestDoc {
    pub fn normal(handle: &[u8], version: u64, name: &str) -> Self {
        Self {
            handle: handle.to_vec(),
            version,
            cols: vec![
                (1, Some(name.to_string())),
                (2, Some(format!("desc_{version}"))),
                (3, Some(format!("pk1_{version}"))),
                (4, Some(format!("pk2_{version}"))),
            ],
        }
    }

    pub fn numeric(handle: &[u8], version: u64, id: i64, tenant_id: i64) -> Self {
        Self {
            handle: handle.to_vec(),
            version,
            cols: vec![
                (1, Some(format!("severity_{id}"))),
                (2, Some(format!("body_{tenant_id}"))),
                (3, Some(id.to_string())),
                (4, Some(tenant_id.to_string())),
            ],
        }
    }
}

/// Builds a fragment with one tantivy commit per inner `Vec` (segment).
pub fn build_test_frag(
    index_meta: &IndexMeta,
    segments: Vec<Vec<TestDoc>>,
) -> Result<(TempDir, FragReaderMeta)> {
    let temp_dir = TempDir::new()?;
    let frag_path = create_test_frag_path(temp_dir.path())?;
    let index_path = utils::frag_index_path(&frag_path);
    fs::create_dir_all(&index_path)?;

    let schema = index_meta.get_tantivy_schema().clone();
    let mut index = Index::create_in_dir(&index_path, schema)?;
    index.set_tokenizers(index_meta.get_analyzer_manager().clone());
    let mut writer = create_index_writer(&index, 1, 15_000_000, &frag_path)?;

    let mut first_handle = None;
    let mut last_handle = None;
    for segment in segments {
        for test_doc in segment {
            first_handle.get_or_insert_with(|| test_doc.handle.clone());
            last_handle = Some(test_doc.handle.clone());
            let raw_doc = RawDoc {
                handle: ShardingKey::from_raw(test_doc.handle),
                version: test_doc.version,
                op_type: OperationType::Insert,
                cols: test_doc.cols,
                is_local_timezone: true,
            };
            writer.add_document(TantivyDoc::new(&raw_doc, index_meta)?.doc)?;
        }
        writer.commit()?;
    }
    writer.wait_merging_threads()?;

    let start = first_handle.context("test frag must contain at least one document")?;
    let end = ShardingKey::next_lexicographic(
        &last_handle.context("test frag must contain at least one document")?,
    );
    Ok((temp_dir, FragReaderMeta::new(&frag_path, start, end)?))
}

/// Loads a fragment reader with the given deleted doc addresses (no doc cache).
pub fn load_frag_reader(
    meta: FragReaderMeta,
    deleted_docs: &[DocAddress],
) -> Result<Arc<FragReader>> {
    let index = meta.load_index()?;
    let mut bitset = FragDocBitSet::empty_for_index(&index)?;
    for addr in deleted_docs {
        bitset.push_deleted_doc(*addr)?;
    }
    let valid_range = meta.frag_range()?;
    Ok(Arc::new(meta.load_reader_with_obsolete_doc_bitset(
        None,
        valid_range,
        FragAliveBitSetBytes::from_doc_bitset(bitset)?,
    )?))
}

/// Opens a [`FragRowsStream`] for tests and returns the initial reader stat.
pub fn open_frag_rows_stream(
    frag_reader: Arc<FragReader>,
    query_info: QueryInfo,
) -> Result<(FragRowsStream, ShardReaderStat)> {
    let mut stat = ShardReaderStat::default();
    let stream = frag_reader.open_rows_stream(Arc::new(query_info), &mut stat)?;
    Ok((stream, stat))
}

/// Extracts version column values from a [`SearchRows`] batch.
pub fn row_versions(rows: &SearchRows) -> Result<Vec<u64>> {
    let version_column = rows
        .u64_column(VERSION_STORE_NAME)
        .context("version column is missing from search result")?;
    (0..rows.len())
        .map(|row_idx| {
            version_column
                .value_at(row_idx)
                .context("version is missing")
        })
        .collect()
}

/// Drains all versions from a [`FragRowsStream`] using the given batch size.
pub fn collect_versions(
    stream: &mut FragRowsStream,
    batch_size: usize,
    stat: &mut ShardReaderStat,
) -> Result<Vec<u64>> {
    let mut versions = Vec::new();
    while let Some(rows) = stream.next_batch(batch_size, stat)? {
        versions.extend(row_versions(&rows)?);
    }
    Ok(versions)
}

/// Delete seg-0 doc-1 (v=20) and seg-1 doc-0 (v=40). Alive: v=10, v=30, v=50.
pub fn deleted_across_segments() -> Vec<DocAddress> {
    vec![DocAddress::new(0, 1), DocAddress::new(1, 0)]
}