use anyhow::Result;
use fts_index::index::IndexMeta;
use fts_index::index::Writer;
use fts_index::protocol::DocumentSet;
use log::*;
use serde::Serialize;
use tokio::sync::mpsc::Receiver;

#[derive(Debug, Serialize)]
pub struct SinkReport {
    pub successes: i64,
    pub deletes: i64,
    pub cost: std::time::Duration,
    #[serde(skip)]
    pub index_meta: Option<IndexMeta>,
    #[serde(skip)]
    pub deleteds: Vec<String>,
}

pub(crate) trait Sink {
    type Item;
    async fn sink(&mut self, receiver: Receiver<Self::Item>) -> Result<SinkReport>;
}

pub(crate) struct CollectionIndexSink {
    writer: Writer,
    concurrency: usize,
}

impl CollectionIndexSink {
    pub fn new(
        storage_path: String,
        index_name: String,
        start_version: u64,
        end_version: u64,
        concurrency: usize,
        memory_budget: impl Into<Option<usize>>,
    ) -> Self {
        debug!(
            "Creating CollectionIndexSink with storage_path: {}, index_name: {}, start_version: {}, end_version: {}, concurrency: {}",
            storage_path, index_name, start_version, end_version, concurrency,
        );
        Self {
            writer: Writer::new(
                &storage_path,
                &index_name,
                start_version,
                end_version,
                memory_budget,
            ),
            concurrency,
        }
    }
}

impl Sink for CollectionIndexSink {
    type Item = DocumentSet;
    async fn sink(&mut self, mut receiver: Receiver<Self::Item>) -> Result<SinkReport> {
        let instant = std::time::Instant::now();
        let mut successes = 0;
        let mut sets = Vec::new();
        let mut deleteds = Vec::new();
        while let Some(document_set) = receiver.recv().await {
            let docs = document_set.to_document_full();
            sets.extend(docs);

            if sets.len() >= self.concurrency * 10000 {
                successes += sets.len() as i64;
                info!("Indexing {} documents", sets.len());
                let result = self.writer.update(sets.clone()).await?;
                deleteds.extend(result.1);
                sets.clear();
            }
        }

        if sets.is_empty() {
            return Ok(SinkReport {
                successes,
                deletes: deleteds.len() as i64,
                cost: instant.elapsed(),
                index_meta: None,
                deleteds,
            });
        }

        successes += sets.len() as i64;
        info!("Indexing {} documents", sets.len());
        let result = self.writer.update(sets).await?;
        deleteds.extend(result.1);

        Ok(SinkReport {
            successes,
            deletes: deleteds.len() as i64,
            cost: instant.elapsed(),
            index_meta: Some(result.0),
            deleteds,
        })
    }
}