use anyhow::{Context, Result, ensure};
use serde::{Deserialize, Serialize};
use tantivy::{
    DocAddress, DocId, Index,
    directory::OwnedBytes,
    fastfield::{AliveBitSet, write_alive_bitset},
    index::{Segment, SegmentId},
};
use tantivy_common::BitSet;

/// Dedup bitsets are modeled in two layers so the dedup pipeline can build
/// them efficiently and the reader pipeline can restore them safely.
///
/// `FragDocBitSet` is the mutable build-time form. Dedup discovers obsolete
/// documents as `DocAddress` values and records deleted doc ids per searchable
/// segment while the local Tantivy index is open.
///
/// `FragAliveBitSetBytes` is the persistable/runtime form. Tantivy's
/// `AliveBitSet` cannot be serialized directly, so this module stores the raw
/// alive-bitset bytes together with segment metadata and reopens an
/// `AliveBitSet` on demand. Keeping `segment_id` and `max_doc` alongside the
/// bytes lets restore code validate that the persisted bitset still matches the
/// current frag segments after restarts or frag replacement, instead of
/// silently applying deletions to the wrong segment layout.
///
/// Segments without deleted docs are encoded as `None` to avoid materializing
/// empty bitsets in memory or on disk.

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct FragBitSet<S> {
    pub segments: Vec<S>,
}

impl<S> Default for FragBitSet<S> {
    fn default() -> Self {
        Self {
            segments: Vec::new(),
        }
    }
}

pub type FragDocBitSet = FragBitSet<SegmentDocBitSet>;
pub type FragAliveBitSetBytes = FragBitSet<SegmentAliveBitSetBytes>;

trait AliveBitSetSegment {
    fn has_deleted_docs(&self) -> bool;

    fn open_alive_bitset(&self) -> Result<Option<AliveBitSet>>;
}

fn frag_has_deleted_docs<S: AliveBitSetSegment>(segments: &[S]) -> bool {
    segments.iter().any(AliveBitSetSegment::has_deleted_docs)
}

fn open_frag_alive_bitsets<S: AliveBitSetSegment>(
    segments: &[S],
) -> Result<Vec<Option<AliveBitSet>>> {
    segments
        .iter()
        .map(AliveBitSetSegment::open_alive_bitset)
        .collect()
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SegmentDocBitSet {
    segment_id: SegmentId,
    max_doc: DocId,
    deleted_doc_ids: Vec<DocId>,
}

impl SegmentDocBitSet {
    pub fn new(segment_id: SegmentId, max_doc: DocId, deleted_doc_ids: Vec<DocId>) -> Self {
        Self {
            segment_id,
            max_doc,
            deleted_doc_ids,
        }
    }

    pub fn encode_alive_bitset_bytes(&self) -> Result<Option<Vec<u8>>> {
        if self.deleted_doc_ids.is_empty() {
            return Ok(None);
        }

        let mut alive_bitset = BitSet::with_max_value_and_full(self.max_doc);
        for &doc_id in &self.deleted_doc_ids {
            ensure!(
                doc_id < self.max_doc,
                "deleted doc id {} exceeds max_doc {}",
                doc_id,
                self.max_doc
            );
            alive_bitset.remove(doc_id);
        }

        let mut bytes = Vec::new();
        write_alive_bitset(&alive_bitset, &mut bytes)
            .context("serialize custom alive bitset failed")?;
        Ok(Some(bytes))
    }

    pub fn open_alive_bitset(&self) -> Result<Option<AliveBitSet>> {
        Ok(self
            .encode_alive_bitset_bytes()?
            .map(|bytes| AliveBitSet::open(OwnedBytes::new(bytes))))
    }
}

impl AliveBitSetSegment for SegmentDocBitSet {
    fn has_deleted_docs(&self) -> bool {
        !self.deleted_doc_ids.is_empty()
    }

    fn open_alive_bitset(&self) -> Result<Option<AliveBitSet>> {
        SegmentDocBitSet::open_alive_bitset(self)
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SegmentAliveBitSetBytes {
    pub segment_id: SegmentId,
    pub max_doc: DocId,
    pub alive_bitset_bytes: Option<Vec<u8>>,
}

impl SegmentAliveBitSetBytes {
    const ENCODED_MAX_DOC_BYTES_LEN: usize = std::mem::size_of::<DocId>();
    const ENCODED_ALIVE_BLOCK_BYTES_LEN: usize = std::mem::size_of::<u64>();

    fn validate_alive_bitset_bytes(&self) -> Result<()> {
        let Some(bytes) = self.alive_bitset_bytes.as_ref() else {
            return Ok(());
        };
        ensure!(
            bytes.len() >= Self::ENCODED_MAX_DOC_BYTES_LEN,
            "alive bitset bytes too short: len={}, segment_id={}, max_doc={}",
            bytes.len(),
            self.segment_id.uuid_string(),
            self.max_doc
        );
        ensure!(
            (bytes.len() - Self::ENCODED_MAX_DOC_BYTES_LEN) % Self::ENCODED_ALIVE_BLOCK_BYTES_LEN
                == 0,
            "alive bitset bytes length {} is not aligned, segment_id={}",
            bytes.len(),
            self.segment_id.uuid_string()
        );
        let encoded_max_doc = DocId::from_le_bytes(
            bytes[0..Self::ENCODED_MAX_DOC_BYTES_LEN]
                .try_into()
                .context("slice to DocId failed")?,
        );
        ensure!(
            encoded_max_doc == self.max_doc,
            "alive bitset max_doc {} does not match expected {}, segment_id={}",
            encoded_max_doc,
            self.max_doc,
            self.segment_id.uuid_string()
        );
        Ok(())
    }

    pub(crate) fn validate_segment_meta(
        &self,
        expected_segment_id: SegmentId,
        expected_max_doc: DocId,
    ) -> Result<()> {
        ensure!(
            self.segment_id == expected_segment_id,
            "alive bitset segment_id {} does not match expected {}",
            self.segment_id.uuid_string(),
            expected_segment_id.uuid_string()
        );
        ensure!(
            self.max_doc == expected_max_doc,
            "alive bitset max_doc {} does not match segment {} max_doc {}",
            self.max_doc,
            expected_segment_id.uuid_string(),
            expected_max_doc
        );
        Ok(())
    }

    pub fn open_alive_bitset(&self) -> Result<Option<AliveBitSet>> {
        self.validate_alive_bitset_bytes()?;
        Ok(self
            .alive_bitset_bytes
            .as_ref()
            .map(|bytes| AliveBitSet::open(OwnedBytes::new(bytes.clone()))))
    }
}

impl AliveBitSetSegment for SegmentAliveBitSetBytes {
    fn has_deleted_docs(&self) -> bool {
        self.alive_bitset_bytes.is_some()
    }

    fn open_alive_bitset(&self) -> Result<Option<AliveBitSet>> {
        SegmentAliveBitSetBytes::open_alive_bitset(self)
    }
}

impl FragBitSet<SegmentDocBitSet> {
    pub fn has_deleted_docs(&self) -> bool {
        frag_has_deleted_docs(&self.segments)
    }

    pub fn open_alive_bitsets(&self) -> Result<Vec<Option<AliveBitSet>>> {
        open_frag_alive_bitsets(&self.segments)
    }

    pub(crate) fn empty_for_index(index: &Index) -> Result<Self> {
        let segments = index
            .searchable_segments()
            .context("load searchable segments failed")?
            .into_iter()
            .map(|segment| {
                SegmentDocBitSet::new(segment.id(), segment.meta().max_doc(), Vec::new())
            })
            .collect();
        Ok(Self { segments })
    }

    pub(crate) fn push_deleted_doc(&mut self, addr: DocAddress) -> Result<()> {
        let segment_count = self.segments.len();
        let segment = self
            .segments
            .get_mut(addr.segment_ord as usize)
            .with_context(|| {
                format!(
                    "segment ordinal {} out of range for doc bitset with {} segments",
                    addr.segment_ord, segment_count
                )
            })?;
        ensure!(
            addr.doc_id < segment.max_doc,
            "doc id {} exceeds max_doc {} for segment ordinal {}",
            addr.doc_id,
            segment.max_doc,
            addr.segment_ord
        );
        segment.deleted_doc_ids.push(addr.doc_id);
        Ok(())
    }

    pub(crate) fn sort_and_dedup(&mut self) {
        for segment in &mut self.segments {
            segment.deleted_doc_ids.sort_unstable();
            segment.deleted_doc_ids.dedup();
        }
    }

    pub fn into_alive_bitset_bytes(self) -> Result<FragAliveBitSetBytes> {
        let segments = self
            .segments
            .into_iter()
            .map(|segment| {
                let alive_bitset_bytes = segment.encode_alive_bitset_bytes()?;
                Ok(SegmentAliveBitSetBytes {
                    segment_id: segment.segment_id,
                    max_doc: segment.max_doc,
                    alive_bitset_bytes,
                })
            })
            .collect::<Result<Vec<_>>>()?;
        Ok(FragAliveBitSetBytes { segments })
    }
}

impl FragBitSet<SegmentAliveBitSetBytes> {
    pub fn has_deleted_docs(&self) -> bool {
        frag_has_deleted_docs(&self.segments)
    }

    pub fn open_alive_bitsets(&self) -> Result<Vec<Option<AliveBitSet>>> {
        open_frag_alive_bitsets(&self.segments)
    }

    pub fn from_doc_bitset(doc_bitset: FragDocBitSet) -> Result<Self> {
        doc_bitset.into_alive_bitset_bytes()
    }

    pub fn from_doc_bitsets(doc_bitsets: Vec<FragDocBitSet>) -> Result<Vec<Self>> {
        doc_bitsets
            .into_iter()
            .map(FragDocBitSet::into_alive_bitset_bytes)
            .collect()
    }

    pub(crate) fn validate_against_searchable_segments(
        &self,
        searchable_segments: &[Segment],
    ) -> Result<()> {
        ensure!(
            searchable_segments.len() == self.segments.len(),
            "segment count mismatch for alive bitsets, segments={}, bitsets={}",
            searchable_segments.len(),
            self.segments.len()
        );
        for (segment, alive_bitset_bytes) in searchable_segments.iter().zip(self.segments.iter()) {
            alive_bitset_bytes.validate_segment_meta(segment.id(), segment.meta().max_doc())?;
            alive_bitset_bytes
                .validate_alive_bitset_bytes()
                .with_context(|| {
                    format!(
                        "alive bitset bytes validation failed for segment {}",
                        segment.id().uuid_string()
                    )
                })?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use anyhow::Result;
    use tantivy::{
        DocAddress, Index, doc,
        index::SegmentId,
        schema::{STORED, STRING, Schema},
    };

    use super::*;

    fn build_test_index() -> Result<Index> {
        let mut schema_builder = Schema::builder();
        let field = schema_builder.add_text_field("body", STRING | STORED);
        let schema = schema_builder.build();

        let index = Index::create_in_ram(schema);
        let mut writer = index.writer(15_000_000)?;
        writer.add_document(doc!(field => "a"))?;
        writer.add_document(doc!(field => "b"))?;
        writer.commit()?;
        writer.wait_merging_threads()?;
        Ok(index)
    }

    #[test]
    fn test_segment_doc_bitset_converts_deleted_and_alive_bitsets() -> Result<()> {
        let segment_doc_bitset = SegmentDocBitSet::new(SegmentId::generate_random(), 5, vec![1, 3]);

        let alive_bitset = segment_doc_bitset
            .open_alive_bitset()?
            .expect("deleted docs should produce a custom alive bitset");
        assert!(alive_bitset.is_deleted(1));
        assert!(alive_bitset.is_deleted(3));
        assert!(!alive_bitset.is_deleted(0));
        assert!(!alive_bitset.is_deleted(2));
        assert!(!alive_bitset.is_deleted(4));
        Ok(())
    }

    #[test]
    fn test_segment_doc_bitset_rejects_out_of_range_doc_ids() {
        let segment_doc_bitset = SegmentDocBitSet::new(SegmentId::generate_random(), 2, vec![2]);

        assert!(segment_doc_bitset.open_alive_bitset().is_err());
    }

    #[test]
    fn test_frag_doc_bitset_push_sort_and_alive_helpers() -> Result<()> {
        let mut frag_doc_bitset = FragDocBitSet {
            segments: vec![
                SegmentDocBitSet::new(SegmentId::generate_random(), 4, Vec::new()),
                SegmentDocBitSet::new(SegmentId::generate_random(), 3, Vec::new()),
            ],
        };

        assert!(!frag_doc_bitset.has_deleted_docs());

        frag_doc_bitset.push_deleted_doc(DocAddress::new(0, 2))?;
        frag_doc_bitset.push_deleted_doc(DocAddress::new(0, 1))?;
        frag_doc_bitset.push_deleted_doc(DocAddress::new(0, 2))?;
        frag_doc_bitset.push_deleted_doc(DocAddress::new(1, 0))?;
        frag_doc_bitset.sort_and_dedup();

        assert!(frag_doc_bitset.has_deleted_docs());
        assert_eq!(frag_doc_bitset.segments[0].deleted_doc_ids, vec![1, 2]);
        assert_eq!(frag_doc_bitset.segments[1].deleted_doc_ids, vec![0]);

        let alive_bitsets = frag_doc_bitset.open_alive_bitsets()?;
        let first_alive_bitset = alive_bitsets[0]
            .as_ref()
            .expect("segment with deleted docs should have alive bitset");
        assert!(first_alive_bitset.is_deleted(1));
        assert!(first_alive_bitset.is_deleted(2));
        assert!(!first_alive_bitset.is_deleted(0));
        assert!(!first_alive_bitset.is_deleted(3));

        let second_alive_bitset = alive_bitsets[1]
            .as_ref()
            .expect("segment with deleted docs should have alive bitset");
        assert!(second_alive_bitset.is_deleted(0));
        assert!(!second_alive_bitset.is_deleted(1));
        assert!(!second_alive_bitset.is_deleted(2));

        assert!(
            frag_doc_bitset
                .push_deleted_doc(DocAddress::new(2, 0))
                .is_err()
        );
        assert!(
            frag_doc_bitset
                .push_deleted_doc(DocAddress::new(1, 3))
                .is_err()
        );
        Ok(())
    }

    #[test]
    fn test_frag_doc_bitset_empty_for_index_uses_segment_max_doc() -> Result<()> {
        let index = build_test_index()?;
        let expected_segment_id = index.searchable_segments()?[0].id();

        let frag_doc_bitset = FragDocBitSet::empty_for_index(&index)?;
        assert_eq!(frag_doc_bitset.segments.len(), 1);
        assert_eq!(frag_doc_bitset.segments[0].segment_id, expected_segment_id);
        assert_eq!(frag_doc_bitset.segments[0].max_doc, 2);
        assert!(frag_doc_bitset.segments[0].deleted_doc_ids.is_empty());
        Ok(())
    }

    #[test]
    fn test_frag_doc_bitset_into_alive_bitset_bytes_preserves_segment_metadata() -> Result<()> {
        let index = build_test_index()?;
        let expected_segment_id = index.searchable_segments()?[0].id();
        let mut frag_doc_bitset = FragDocBitSet::empty_for_index(&index)?;
        frag_doc_bitset.push_deleted_doc(DocAddress::new(0, 1))?;

        let bitset = frag_doc_bitset.into_alive_bitset_bytes()?;

        assert_eq!(bitset.segments[0].segment_id, expected_segment_id);
        assert_eq!(bitset.segments[0].max_doc, 2);
        assert!(bitset.has_deleted_docs());
        let alive_bitsets = bitset.open_alive_bitsets()?;
        let alive_bitset = alive_bitsets[0]
            .as_ref()
            .expect("segment with deleted docs should have alive bitset");
        assert!(alive_bitset.is_deleted(1));
        assert!(!alive_bitset.is_deleted(0));
        Ok(())
    }

    #[test]
    fn test_frag_alive_bitset_bytes_validate_against_searchable_segments_rejects_mismatch()
    -> Result<()> {
        let index = build_test_index()?;
        let mut bitset =
            FragAliveBitSetBytes::from_doc_bitset(FragDocBitSet::empty_for_index(&index)?)?;
        bitset.segments[0].segment_id = SegmentId::generate_random();

        let err = bitset
            .validate_against_searchable_segments(&index.searchable_segments()?)
            .expect_err("segment mismatch should be rejected");
        assert!(format!("{err:#}").contains("segment_id"));
        Ok(())
    }

    #[test]
    fn test_frag_alive_bitset_bytes_validate_against_searchable_segments_rejects_invalid_bytes()
    -> Result<()> {
        let index = build_test_index()?;
        let mut frag_doc_bitset = FragDocBitSet::empty_for_index(&index)?;
        frag_doc_bitset.push_deleted_doc(DocAddress::new(0, 1))?;
        let mut bitset = FragAliveBitSetBytes::from_doc_bitset(frag_doc_bitset)?;
        bitset.segments[0].alive_bitset_bytes = Some(vec![1, 2, 3]);

        let err = bitset
            .validate_against_searchable_segments(&index.searchable_segments()?)
            .expect_err("invalid alive bitset bytes should be rejected");
        assert!(format!("{err:#}").contains("alive bitset bytes"));
        Ok(())
    }
}