use std::io::{Read, Write};

use anyhow::{Context, Result, ensure};
use flate2::{Compression, read::ZlibDecoder, write::ZlibEncoder};
pub use metadata_pb::{PersistedFragDedupBitset, PersistedShardDedupBitsets};
use prost::Message;
use tantivy::index::SegmentId;
use tici_common::index::IndexKey;
use tici_grpc::FragMeta;

use crate::{
    metadata::tici_shard_metadata as metadata_pb,
    reader::frag_reader::FragReaderMeta,
    utils::{
        dedup_bitset::{FragAliveBitSetBytes, SegmentAliveBitSetBytes},
        key::ShardingKeyRange,
    },
};

/// Persisted dedup snapshot for a shard after a refresh.
///
/// Lifecycle:
/// - `ShardMirror::refresh` first publishes the latest frag metadata, then runs
///   dedup and wraps the resulting `FragAliveBitSetBytes` in this type.
/// - The encoded payload is stored in metadata storage so a later process can
///   reuse the same per-frag alive bitsets without rerunning dedup immediately.
/// - On restore, callers decode the payload and validate the shard range, frag
///   order, segment metadata, and alive-bitset bytes before extracting the
///   bitsets. Any mismatch means the snapshot is stale or corrupted relative to
///   the current frag layout and must be discarded in favor of rebuilding dedup
///   state.
///
/// This type is therefore a cache of dedup results, not the source of truth.
/// It exists to bridge process restarts and the window between publishing new
/// frag metadata and reconstructing readers for that latest shard view.
impl PersistedShardDedupBitsets {
    pub fn new(
        index_key: &IndexKey,
        shard_id: u64,
        shard_range: &ShardingKeyRange,
        frag_metas: &[FragMeta],
        frag_bitsets: &[FragAliveBitSetBytes],
    ) -> Result<Self> {
        ensure!(
            frag_metas.len() == frag_bitsets.len(),
            "frag meta count {} does not match dedup bitset count {}",
            frag_metas.len(),
            frag_bitsets.len()
        );
        let frags = frag_metas
            .iter()
            .zip(frag_bitsets.iter())
            .map(|(frag_meta, bitset)| PersistedFragDedupBitset {
                frag_path: frag_meta.frag_path.clone(),
                bitset: Some(encode_frag_alive_bitset_bytes(bitset)),
            })
            .collect();
        Ok(Self {
            shard_start_key: shard_range.start().raw().to_vec(),
            shard_end_key: shard_range.end().raw().to_vec(),
            frags,
            keyspace_id: index_key.keyspace_id,
            table_id: index_key.table_id,
            index_id: index_key.index_id,
            shard_id,
        })
    }

    pub fn encode_compressed(&self) -> Result<Vec<u8>> {
        let encoded = self.encode_to_vec();
        let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
        encoder
            .write_all(&encoded)
            .context("compress shard dedup bitsets failed")?;
        encoder
            .finish()
            .context("finalize shard dedup bitset compression failed")
    }

    pub fn decode_compressed(bytes: &[u8]) -> Result<Self> {
        let mut decoder = ZlibDecoder::new(bytes);
        let mut decoded = Vec::new();
        decoder
            .read_to_end(&mut decoded)
            .context("decompress shard dedup bitsets failed")?;
        Self::decode(decoded.as_slice()).context("deserialize shard dedup bitsets failed")
    }

    pub fn validate_and_extract(
        &self,
        index_key: &IndexKey,
        shard_id: u64,
        shard_range: &ShardingKeyRange,
        frag_metas: &[FragMeta],
        frag_reader_metas: &[FragReaderMeta],
    ) -> Result<Vec<FragAliveBitSetBytes>> {
        self.validate_shard_identity(index_key, shard_id)?;
        ensure!(
            self.shard_start_key == shard_range.start().raw(),
            "persisted shard start key mismatch"
        );
        ensure!(
            self.shard_end_key == shard_range.end().raw(),
            "persisted shard end key mismatch"
        );
        ensure!(
            self.frags.len() == frag_metas.len(),
            "persisted frag count {} does not match current frag count {}",
            self.frags.len(),
            frag_metas.len()
        );
        ensure!(
            self.frags.len() == frag_reader_metas.len(),
            "persisted frag count {} does not match frag reader meta count {}",
            self.frags.len(),
            frag_reader_metas.len()
        );

        self.frags
            .iter()
            .zip(frag_metas.iter())
            .zip(frag_reader_metas.iter())
            .map(|((persisted_frag, frag_meta), frag_reader_meta)| {
                validate_frag_dedup_bitset(persisted_frag, frag_meta, frag_reader_meta)
            })
            .collect()
    }

    fn validate_shard_identity(&self, index_key: &IndexKey, shard_id: u64) -> Result<()> {
        ensure!(
            self.keyspace_id == index_key.keyspace_id,
            "persisted keyspace_id {} does not match current keyspace_id {}",
            self.keyspace_id,
            index_key.keyspace_id
        );
        ensure!(
            self.table_id == index_key.table_id,
            "persisted table_id {} does not match current table_id {}",
            self.table_id,
            index_key.table_id
        );
        ensure!(
            self.index_id == index_key.index_id,
            "persisted index_id {} does not match current index_id {}",
            self.index_id,
            index_key.index_id
        );
        ensure!(
            self.shard_id == shard_id,
            "persisted shard_id {} does not match current shard_id {}",
            self.shard_id,
            shard_id
        );
        Ok(())
    }
}

fn encode_frag_alive_bitset_bytes(
    bitset: &FragAliveBitSetBytes,
) -> metadata_pb::FragAliveBitSetBytes {
    metadata_pb::FragAliveBitSetBytes {
        segments: bitset
            .segments
            .iter()
            .map(encode_segment_alive_bitset)
            .collect(),
    }
}

fn decode_frag_alive_bitset_bytes(
    proto: &metadata_pb::FragAliveBitSetBytes,
) -> Result<FragAliveBitSetBytes> {
    let segments = proto
        .segments
        .iter()
        .map(decode_segment_alive_bitset)
        .collect::<Result<Vec<_>>>()?;
    Ok(FragAliveBitSetBytes { segments })
}

fn encode_segment_alive_bitset(
    segment: &SegmentAliveBitSetBytes,
) -> metadata_pb::SegmentAliveBitSetBytes {
    metadata_pb::SegmentAliveBitSetBytes {
        segment_id: segment.segment_id.uuid_string(),
        max_doc: segment.max_doc,
        alive_bitset_bytes: segment.alive_bitset_bytes.clone().unwrap_or_default(),
    }
}

fn decode_segment_alive_bitset(
    proto: &metadata_pb::SegmentAliveBitSetBytes,
) -> Result<SegmentAliveBitSetBytes> {
    let segment_id = SegmentId::from_uuid_string(&proto.segment_id)
        .with_context(|| format!("parse persisted segment_id {} failed", proto.segment_id))?;
    Ok(SegmentAliveBitSetBytes {
        segment_id,
        max_doc: proto.max_doc,
        alive_bitset_bytes: (!proto.alive_bitset_bytes.is_empty())
            .then(|| proto.alive_bitset_bytes.clone()),
    })
}

fn validate_frag_dedup_bitset(
    persisted_frag: &PersistedFragDedupBitset,
    frag_meta: &FragMeta,
    frag_reader_meta: &FragReaderMeta,
) -> Result<FragAliveBitSetBytes> {
    let persisted_bitset = decode_frag_alive_bitset_bytes(
        persisted_frag
            .bitset
            .as_ref()
            .context("missing frag dedup bitset payload")?,
    )?;
    ensure!(
        persisted_frag.frag_path == frag_meta.frag_path,
        "persisted frag path {} does not match current frag path {}",
        persisted_frag.frag_path,
        frag_meta.frag_path
    );

    let index = frag_reader_meta.load_index().with_context(|| {
        format!(
            "load index for persisted dedup validation failed: {}",
            frag_meta.frag_path
        )
    })?;
    let searchable_segments = index.searchable_segments().with_context(|| {
        format!(
            "load searchable segments for persisted dedup validation failed: {}",
            frag_meta.frag_path
        )
    })?;
    ensure!(
        searchable_segments.len() == persisted_bitset.segments.len(),
        "persisted segment count {} does not match local segment count {} for frag {}",
        persisted_bitset.segments.len(),
        searchable_segments.len(),
        frag_meta.frag_path
    );
    persisted_bitset
        .validate_against_searchable_segments(&searchable_segments)
        .with_context(|| {
            format!(
                "persisted alive bitset validation failed for frag {}",
                frag_meta.frag_path
            )
        })?;

    Ok(persisted_bitset)
}

#[cfg(test)]
mod tests {
    use anyhow::{Context, Result};
    use tantivy::{
        DocAddress, Index, doc,
        index::SegmentId,
        schema::{STORED, STRING, Schema},
    };
    use tempfile::TempDir;
    use tici_common::index::IndexKey;
    use tici_grpc::FragProperty;

    use super::*;
    use crate::{
        test_utils::{create_test_frag_path, test_mock_frag_storage_refs},
        utils::dedup_bitset::{FragDocBitSet, SegmentDocBitSet},
    };

    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)
    }

    fn build_test_frag() -> Result<(TempDir, FragMeta, FragReaderMeta)> {
        let temp_dir = TempDir::new()?;
        let frag_path = create_test_frag_path(temp_dir.path())?;
        let index_path = crate::utils::frag_index_path(&frag_path);
        std::fs::create_dir_all(&index_path)
            .with_context(|| format!("create_dir_all failed, path={index_path}"))?;

        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_dir(&index_path, 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()?;

        let frag_meta = FragMeta::new(
            1,
            "frag/f_1",
            FragProperty {
                size: 1,
                count: 2,
                min_handle: vec![0],
                split_handles: vec![vec![1]],
                max_handle: vec![2],
                delete_count: 0,
            },
            test_mock_frag_storage_refs(),
            1,
        )?;
        let frag_reader_meta = FragReaderMeta::new(&frag_path, vec![0], vec![2])?;
        Ok((temp_dir, frag_meta, frag_reader_meta))
    }

    fn test_index_key() -> IndexKey {
        IndexKey::new(11, 22, 33)
    }

    #[test]
    fn test_frag_alive_bitset_bytes_from_doc_bitset() -> 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 = FragAliveBitSetBytes::from_doc_bitset(frag_doc_bitset)?;

        assert_eq!(bitset.segments[0].segment_id, expected_segment_id);
        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_default_frag_alive_bitset_bytes_has_no_deleted_docs() -> Result<()> {
        let bitset = FragAliveBitSetBytes::default();
        assert!(!bitset.has_deleted_docs());
        assert!(bitset.open_alive_bitsets()?.is_empty());
        Ok(())
    }

    #[test]
    fn test_persisted_shard_dedup_bitsets_round_trip() -> Result<()> {
        let index = build_test_index()?;
        let segment = index.searchable_segments()?[0].id();
        let bitset = FragAliveBitSetBytes::from_doc_bitset(FragDocBitSet {
            segments: vec![SegmentDocBitSet::new(segment, 2, vec![1])],
        })?;
        let index_key = test_index_key();
        let shard_id = 44;
        let shard_range = ShardingKeyRange::half_open(vec![1], vec![9])?;
        let persisted = PersistedShardDedupBitsets {
            shard_start_key: shard_range.start().raw().to_vec(),
            shard_end_key: shard_range.end().raw().to_vec(),
            frags: vec![PersistedFragDedupBitset {
                frag_path: "frag".to_string(),
                bitset: Some(encode_frag_alive_bitset_bytes(&bitset)),
            }],
            keyspace_id: index_key.keyspace_id,
            table_id: index_key.table_id,
            index_id: index_key.index_id,
            shard_id,
        };
        let encoded = persisted.encode_compressed()?;
        let decoded = PersistedShardDedupBitsets::decode_compressed(&encoded)?;
        assert_eq!(decoded, persisted);
        Ok(())
    }

    #[test]
    fn test_validate_and_extract_rejects_bitset_segment_id_mismatch() -> Result<()> {
        let (_temp_dir, frag_meta, frag_reader_meta) = build_test_frag()?;
        let index_key = test_index_key();
        let shard_id = 44;
        let shard_range = ShardingKeyRange::half_open(vec![0], vec![9])?;
        let mut bitset = FragAliveBitSetBytes::from_doc_bitset(FragDocBitSet::empty_for_index(
            &frag_reader_meta.load_index()?,
        )?)?;
        bitset.segments[0].segment_id = SegmentId::generate_random();
        let persisted = PersistedShardDedupBitsets {
            shard_start_key: shard_range.start().raw().to_vec(),
            shard_end_key: shard_range.end().raw().to_vec(),
            frags: vec![PersistedFragDedupBitset {
                frag_path: frag_meta.frag_path.clone(),
                bitset: Some(encode_frag_alive_bitset_bytes(&bitset)),
            }],
            keyspace_id: index_key.keyspace_id,
            table_id: index_key.table_id,
            index_id: index_key.index_id,
            shard_id,
        };

        let err = persisted
            .validate_and_extract(
                &index_key,
                shard_id,
                &shard_range,
                &[frag_meta],
                &[frag_reader_meta],
            )
            .expect_err("segment id mismatch should be rejected");
        assert!(format!("{err:#}").contains("segment_id"));
        Ok(())
    }

    #[test]
    fn test_validate_and_extract_rejects_invalid_alive_bitset_bytes() -> Result<()> {
        let (_temp_dir, frag_meta, frag_reader_meta) = build_test_frag()?;
        let index_key = test_index_key();
        let shard_id = 44;
        let shard_range = ShardingKeyRange::half_open(vec![0], vec![9])?;
        let index = frag_reader_meta.load_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 persisted = PersistedShardDedupBitsets {
            shard_start_key: shard_range.start().raw().to_vec(),
            shard_end_key: shard_range.end().raw().to_vec(),
            frags: vec![PersistedFragDedupBitset {
                frag_path: frag_meta.frag_path.clone(),
                bitset: Some(encode_frag_alive_bitset_bytes(&bitset)),
            }],
            keyspace_id: index_key.keyspace_id,
            table_id: index_key.table_id,
            index_id: index_key.index_id,
            shard_id,
        };

        let err = persisted
            .validate_and_extract(
                &index_key,
                shard_id,
                &shard_range,
                &[frag_meta],
                &[frag_reader_meta],
            )
            .expect_err("invalid alive bitset bytes should be rejected");
        assert!(format!("{err:#}").contains("alive bitset bytes"));
        Ok(())
    }

    #[test]
    fn test_validate_and_extract_rejects_shard_identity_mismatch() -> Result<()> {
        let (_temp_dir, frag_meta, frag_reader_meta) = build_test_frag()?;
        let index_key = test_index_key();
        let shard_id = 44;
        let shard_range = ShardingKeyRange::half_open(vec![0], vec![9])?;
        let bitset = FragAliveBitSetBytes::from_doc_bitset(FragDocBitSet::empty_for_index(
            &frag_reader_meta.load_index()?,
        )?)?;
        let persisted = PersistedShardDedupBitsets {
            shard_start_key: shard_range.start().raw().to_vec(),
            shard_end_key: shard_range.end().raw().to_vec(),
            frags: vec![PersistedFragDedupBitset {
                frag_path: frag_meta.frag_path.clone(),
                bitset: Some(encode_frag_alive_bitset_bytes(&bitset)),
            }],
            keyspace_id: index_key.keyspace_id,
            table_id: index_key.table_id,
            index_id: index_key.index_id,
            shard_id,
        };

        let err = persisted
            .validate_and_extract(
                &IndexKey::new(
                    index_key.keyspace_id,
                    index_key.table_id,
                    index_key.index_id + 1,
                ),
                shard_id,
                &shard_range,
                &[frag_meta],
                &[frag_reader_meta],
            )
            .expect_err("shard identity mismatch should be rejected");
        assert!(format!("{err:#}").contains("index_id"));
        Ok(())
    }
}