// Copyright 2026 PingCAP, Inc.

use std::{sync::Arc, time::Duration};

use anyhow::Result;
use clara_fts::test_util::{PlainFtsQueryInfo, make_unscored_query};

use super::{FtsDedSnapVersionAllocator, merge_fts_l2_files};
use crate::table::{
    file::{File, FileMmapGuard, InMemFile, MmapData},
    fts::{
        FtsCache, FtsIndexReader, FtsLevels,
        dedicated_file::{DedicatedFile, DedicatedFileBuilderOptions, EDedicatedFile},
        iter::{IntPk, PkReader, PkType},
        lp_key,
        test_util::new_ded,
    },
    table,
};

#[derive(Default)]
struct TestSnapVersionAllocator {
    next: u64,
}

impl FtsDedSnapVersionAllocator for TestSnapVersionAllocator {
    async fn alloc_snap_version(&mut self) -> Result<u64> {
        self.next += 1;
        Ok(self.next)
    }
}

#[tokio::test]
async fn merge_l2_dedicated_files_int_pk() -> Result<()> {
    let file1 = new_ded(1)
        .lp(100, 0, |d| {
            d(1, 100, false, "doc1");
            d(2, 90, false, "doc2");
        })
        .finish_as_file();

    let file2 = new_ded(2)
        .lp(100, 0, |d| {
            d(1, 110, false, "doc3");
            d(3, 80, true, "doc4");
        })
        .finish_as_file();

    let files = vec![file1, file2];

    let mut snap_alloc = TestSnapVersionAllocator {
        next: files
            .iter()
            .map(|f| f.props().get_snap_version())
            .max()
            .unwrap_or(0),
    };
    let output = merge_fts_l2_files(
        &files,
        0,
        DedicatedFileBuilderOptions::default(),
        false,
        &mut snap_alloc,
    )
    .await?
    .expect("should produce merged file");
    assert_eq!(output.summary.props.get_table_id(), 100);
    assert_eq!(output.summary.props.get_pk_total(), 4);
    assert_eq!(output.summary.props.get_lp_key(), lp_key(100, 0));
    Ok(())
}

#[derive(Clone)]
struct DelayFile {
    inner: InMemFile,
    delay: Duration,
}

impl DelayFile {
    fn new(id: u64, data: bytes::Bytes, delay: Duration) -> Self {
        Self {
            inner: InMemFile::new(id, data),
            delay,
        }
    }
}

#[async_trait::async_trait]
impl File for DelayFile {
    fn id(&self) -> u64 {
        self.inner.id()
    }

    fn size(&self) -> u64 {
        self.inner.size()
    }

    fn is_sync(&self) -> bool {
        self.inner.is_sync()
    }

    fn read(&self, off: u64, length: usize) -> table::Result<bytes::Bytes> {
        self.inner.read(off, length)
    }

    fn read_at(&self, buf: &mut [u8], offset: u64) -> table::Result<()> {
        self.inner.read_at(buf, offset)
    }

    fn read_table_meta(&self, off: u64, length: usize) -> table::Result<bytes::Bytes> {
        self.inner.read_table_meta(off, length)
    }

    fn mmap(&self) -> table::Result<MmapData> {
        self.inner.mmap()
    }

    fn mmap2(&self) -> table::Result<bytes::Bytes> {
        self.inner.mmap2()
    }

    async fn mmap_range(
        &self,
        offset: u64,
        length: usize,
    ) -> table::Result<(bytes::Bytes, FileMmapGuard)> {
        if !self.delay.is_zero() {
            tokio::time::sleep(self.delay).await;
        }
        self.inner.mmap_range(offset, length).await
    }

    fn storage_class(&self) -> schema::schema::StorageClass {
        self.inner.storage_class()
    }

    fn as_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync> {
        self
    }
}

fn search_doc_ids(reader: &clara_fts::IndexReader, term: &str) -> Result<Vec<u32>> {
    let mut info = PlainFtsQueryInfo::default();
    info.query = term.to_owned();
    let query = make_unscored_query(&info);
    let mut results = Vec::new();
    reader.search(&query, &mut results)?;
    Ok(results.into_iter().map(|r| r.doc_id).collect())
}

#[tokio::test]
async fn merge_l2_preserves_file_order_for_tantivy_segments() -> Result<()> {
    // Build 2 dedicated files with disjoint PK ranges so that the expected
    // merged doc order is unambiguous: PK 1,2,3,4.
    let bytes_slow = new_ded(1)
        .lp(100, 0, |d| {
            d(1, 100, false, "slow1");
            d(2, 90, false, "slow2");
        })
        .finish_as_bytes();
    let bytes_fast = new_ded(2)
        .lp(100, 0, |d| {
            d(3, 80, false, "fast3");
            d(4, 70, false, "fast4");
        })
        .finish_as_bytes();

    // Make the first file load slower, forcing `cached_read_index()` futures to
    // complete out-of-order. The merge implementation must still keep the
    // Tantivy segment order aligned with `files` order.
    let file_slow = EDedicatedFile::new(
        Arc::new(DelayFile::new(1, bytes_slow, Duration::from_millis(50))),
        FtsCache::disabled(),
    )?;
    let file_fast = EDedicatedFile::new(
        Arc::new(DelayFile::new(2, bytes_fast, Duration::from_millis(0))),
        FtsCache::disabled(),
    )?;
    let files = vec![file_slow, file_fast];

    let mut snap_alloc = TestSnapVersionAllocator {
        next: files
            .iter()
            .map(|f| f.props().get_snap_version())
            .max()
            .unwrap_or(0),
    };
    let output = merge_fts_l2_files(
        &files,
        0,
        DedicatedFileBuilderOptions::default(),
        false,
        &mut snap_alloc,
    )
    .await?
    .expect("should produce merged file");

    let merged = DedicatedFile::from_buffer(output.data.as_ref())?;
    let merged_int = merged.as_int().unwrap();
    let reader = merged.cached_read_index().await?;

    // Term from the "fast" file should map to doc_id=2 -> pk=3.
    let doc_ids = search_doc_ids(&reader, "fast3")?;
    assert_eq!(doc_ids, vec![2]);
    let (pk_bytes, ..) = merged_int
        .async_at(doc_ids[0] as usize, u64::MAX)
        .await?
        .expect("doc should be visible");
    assert_eq!(IntPk::decode(pk_bytes.as_ref())?, 3);

    // Term from the "slow" file should map to doc_id=0 -> pk=1.
    let doc_ids = search_doc_ids(&reader, "slow1")?;
    assert_eq!(doc_ids, vec![0]);
    let (pk_bytes, ..) = merged_int
        .async_at(doc_ids[0] as usize, u64::MAX)
        .await?
        .expect("doc should be visible");
    assert_eq!(IntPk::decode(pk_bytes.as_ref())?, 1);

    Ok(())
}

#[tokio::test]
async fn merge_l2_keeps_tombstone_below_safe_ts() -> Result<()> {
    const TABLE_ID: i64 = 100;
    const INDEX_ID: i64 = 0;
    const SAFE_TS: u64 = 200;

    // Simulate a partial L2 compaction:
    // - An older L2 file contains a live version of pk=1.
    // - A newer L2 file contains a tombstone for pk=1, but it's < safe_ts.
    // - If we drop that tombstone while still keeping the older file, pk=1 can be
    //   resurrected during reads.
    let old = new_ded(1)
        .lp(TABLE_ID, INDEX_ID, |d| {
            d(1, 120, false, "alpha");
        })
        .finish_as_file();

    let tombstone = new_ded(2)
        .lp(TABLE_ID, INDEX_ID, |d| {
            d(1, 150, true, "alpha");
        })
        .finish_as_file();

    // An extra file to make the merge set non-trivial; the old file is
    // intentionally excluded from the merge input.
    let other = new_ded(3)
        .lp(TABLE_ID, INDEX_ID, |d| {
            d(2, 160, false, "beta");
        })
        .finish_as_file();

    let merge_inputs = vec![tombstone, other];
    let mut snap_alloc = TestSnapVersionAllocator {
        next: merge_inputs
            .iter()
            .map(|f| f.props().get_snap_version())
            .max()
            .unwrap_or(0),
    };
    let output = merge_fts_l2_files(
        &merge_inputs,
        SAFE_TS,
        DedicatedFileBuilderOptions::default(),
        false,
        &mut snap_alloc,
    )
    .await?
    .expect("should produce merged file");

    // In a real compaction, the output file gets a new (greater) file id.
    // Preserve this ordering so the output shadows remaining older L2 files.
    let merged = EDedicatedFile::new(
        Arc::new(InMemFile::new(10, output.data.clone())),
        FtsCache::disabled(),
    )?;

    let query = make_unscored_query(&PlainFtsQueryInfo {
        query: "alpha".into(),
        index_id: INDEX_ID,
        ..Default::default()
    });

    {
        let mut levels = FtsLevels::default();
        levels.track_index(TABLE_ID, INDEX_ID);
        levels.insert_l2_file(old.clone());
        let reader = FtsIndexReader::<IntPk>::new(
            Arc::new(levels),
            TABLE_ID,
            u64::MAX,
            None,
            None,
            query.clone(),
        );
        let hits = reader.search(&[]).await?;
        assert_eq!(hits.len(), 1);
        assert_eq!(IntPk::decode(hits[0].pk.as_ref())?, 1);
        assert_eq!(hits[0].version, 120);
    }

    {
        let mut levels = FtsLevels::default();
        levels.track_index(TABLE_ID, INDEX_ID);
        levels.insert_l2_file(old);
        levels.insert_l2_file(merged);

        let reader =
            FtsIndexReader::<IntPk>::new(Arc::new(levels), TABLE_ID, u64::MAX, None, None, query);
        let hits = reader.search(&[]).await?;
        assert!(
            hits.is_empty(),
            "tombstone below safe_ts should shadow older live versions"
        );
    }

    Ok(())
}

#[tokio::test]
async fn merge_l2_drops_tombstone_below_safe_ts_when_full_merge() -> Result<()> {
    const TABLE_ID: i64 = 100;
    const INDEX_ID: i64 = 0;
    const SAFE_TS: u64 = 200;

    // Simulate a full L2 compaction:
    // - An older file contains a live version of pk=1.
    // - A newer file contains a tombstone for pk=1, but it's < safe_ts.
    // When we merge all L2 files for the LP, it's safe to drop that tombstone
    // (and everything older), reclaiming space without resurrecting pk=1.
    let old = new_ded(1)
        .lp(TABLE_ID, INDEX_ID, |d| {
            d(1, 120, false, "alpha");
            d(2, 140, false, "beta");
        })
        .finish_as_file();

    let tombstone = new_ded(2)
        .lp(TABLE_ID, INDEX_ID, |d| {
            d(1, 150, true, "alpha");
            d(3, 160, false, "gamma");
        })
        .finish_as_file();

    let merge_inputs = vec![old, tombstone];
    let mut snap_alloc = TestSnapVersionAllocator {
        next: merge_inputs
            .iter()
            .map(|f| f.props().get_snap_version())
            .max()
            .unwrap_or(0),
    };
    let output = merge_fts_l2_files(
        &merge_inputs,
        SAFE_TS,
        DedicatedFileBuilderOptions::default(),
        true,
        &mut snap_alloc,
    )
    .await?
    .expect("should produce merged file");

    let merged = DedicatedFile::from_buffer(output.data.as_ref())?;
    let merged_int = merged.as_int().unwrap();
    let reader = merged.cached_read_index().await?;

    assert!(
        search_doc_ids(&reader, "alpha")?.is_empty(),
        "tombstone below safe_ts should be dropped in full merge"
    );

    let beta = search_doc_ids(&reader, "beta")?;
    assert_eq!(beta.len(), 1);
    let (pk_bytes, version, is_deleted) = merged_int
        .async_at(beta[0] as usize, u64::MAX)
        .await?
        .expect("beta should exist");
    assert_eq!(IntPk::decode(pk_bytes.as_ref())?, 2);
    assert_eq!(version, 140);
    assert_eq!(is_deleted, 0);

    Ok(())
}