use std::time::Instant;

use anyhow::{Context, Result, bail};
use log::info;
use tantivy::{Index, index::SegmentId, schema::*, tokenizer::TokenizerManager};
use tici_common::cdc_file_name::CDCFileName;
use tokio::time::Instant as TokioInstant;

use crate::{
    utils::{
        self, create_index_writer, metrics::WRITER_FRAG_OPERATION_LATENCY, row_entry::RowEntry,
        schema::SortField, tantivy_doc::TantivyDoc,
    },
    writer::tantivy_doc_buffer::{DocStat, HandleStat, TantivyDocBuffer},
};

pub struct FragWriterBuilder {
    frag_base_path: Option<String>,
    frag_path: Option<String>,
    shard_tracing_id: Option<String>,
    schema: Option<Schema>,
    index_num_threads: Option<usize>,
    index_mem_budget: Option<usize>,
    is_base: Option<bool>,
    analyzer_manager: Option<TokenizerManager>,
    sort_fields: Option<Vec<SortField>>,
}

impl FragWriterBuilder {
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        Self {
            frag_base_path: None,
            frag_path: None,
            shard_tracing_id: None,
            schema: None,
            index_num_threads: None,
            index_mem_budget: None,
            is_base: None,
            analyzer_manager: None,
            sort_fields: None,
        }
    }

    pub fn frag_base_path(mut self, path: &str) -> Self {
        self.frag_base_path = Some(path.to_string());
        self
    }

    pub fn frag_path(mut self, path: &str) -> Self {
        self.frag_path = Some(path.to_string());
        self
    }

    pub fn shard_tracing_id(mut self, shard_tracing_id: &str) -> Self {
        self.shard_tracing_id = Some(shard_tracing_id.to_string());
        self
    }

    pub fn schema(mut self, schema: Schema) -> Self {
        self.schema = Some(schema);
        self
    }

    pub fn index_num_threads(mut self, n: usize) -> Self {
        self.index_num_threads = Some(n);
        self
    }

    pub fn index_mem_budget(mut self, budget: usize) -> Self {
        self.index_mem_budget = Some(budget);
        self
    }

    pub fn is_base(mut self, is_base: bool) -> Self {
        self.is_base = Some(is_base);
        self
    }

    pub fn analyzer_manager(mut self, analyzer_manager: TokenizerManager) -> Self {
        self.analyzer_manager = Some(analyzer_manager);
        self
    }

    pub fn sort_fields(mut self, sort_fields: Vec<SortField>) -> Self {
        self.sort_fields = Some(sort_fields);
        self
    }

    pub fn build(self) -> Result<FragWriter> {
        let frag_base_path = self.frag_base_path.context("frag_base_path missing")?;
        let frag_path = self.frag_path.context("frag_path missing")?;
        let shard_tracing_id = self.shard_tracing_id;
        let schema = self.schema.context("schema missing")?;
        let index_num_threads = self
            .index_num_threads
            .context("index_num_threads missing")?;
        let index_mem_budget = self.index_mem_budget.context("index_mem_budget missing")?;
        let is_base = self.is_base.context("is_base missing")?;
        let analyzer_manager = self.analyzer_manager.context("analyzer_manager missing")?;
        let sort_columns = self.sort_fields.context("sort_columns missing")?;

        Ok(FragWriter {
            frag_base_path,
            frag_path,
            shard_tracing_id,
            last_cdc_file: None,
            pending_progress_since: None,
            is_base,
            schema,
            analyzer_manager,
            tantivy_doc_buffer: None,
            index_num_threads,
            index_mem_budget,
            sort_columns,
        })
    }
}

#[derive(Clone)]
pub struct FragResult {
    pub seg_ids: Vec<SegmentId>,
    pub handle_stat: HandleStat,
    pub finalize_elapsed_ms: u128,
}

pub struct FragWriter {
    frag_base_path: String,
    frag_path: String,
    shard_tracing_id: Option<String>,
    last_cdc_file: Option<CDCFileName>,
    pending_progress_since: Option<TokioInstant>,
    is_base: bool,
    schema: Schema,
    analyzer_manager: TokenizerManager,
    tantivy_doc_buffer: Option<TantivyDocBuffer>,
    index_num_threads: usize,
    index_mem_budget: usize,
    sort_columns: Vec<SortField>,
}

impl FragWriter {
    pub fn add_tantivy_doc(&mut self, tantivy_doc: TantivyDoc) {
        if self.tantivy_doc_buffer.is_none() {
            self.tantivy_doc_buffer = Some(TantivyDocBuffer::new());
        }
        let tantivy_doc_buffer = self.tantivy_doc_buffer.as_mut().unwrap();
        tantivy_doc_buffer.add(tantivy_doc);
    }

    pub fn finalize(&mut self) -> Result<FragResult> {
        if self.tantivy_doc_buffer.is_none() {
            bail!("tantivy_doc_buffer is None, frag_path={}", self.frag_path);
        }

        let finalize_watcher = Instant::now();
        let tantivy_doc_buffer = self.tantivy_doc_buffer.take().unwrap();
        let tantivy_doc_res = tantivy_doc_buffer.finalize(&self.sort_columns)?;
        let idx_path = utils::frag_index_path(&self.frag_path);
        std::fs::create_dir_all(&idx_path)
            .with_context(|| format!("Failed to create index directory idx_path={}", idx_path))?;
        let mut index = Index::create_in_dir(&idx_path, self.schema.clone())
            .with_context(|| format!("Failed to create index idx_path={}", idx_path))?;
        index.set_tokenizers(self.analyzer_manager.clone());
        let mut index_writer = create_index_writer(
            &index,
            self.index_num_threads,
            self.index_mem_budget,
            &self.frag_path,
        )?;
        let watcher = Instant::now();
        for doc in tantivy_doc_res.docs {
            index_writer.add_document(doc)?;
        }
        let add_doc_elapsed_ms = watcher.elapsed().as_millis();

        let watcher = Instant::now();
        index_writer
            .commit()
            .with_context(|| format!("index_writer commit failed. frag_path={}", self.frag_path))?;
        let commit_elapsed_ms = watcher.elapsed().as_millis();

        let watcher = Instant::now();
        RowEntry::write_all(
            utils::frag_row_entry_path(&self.frag_path),
            &tantivy_doc_res.row_entries,
        )?;
        let write_handle_set_elapsed_ms = watcher.elapsed().as_millis();

        let watcher = Instant::now();
        index_writer.wait_merging_threads().with_context(|| {
            format!("wait_merging_threads failed. frag_path={}", self.frag_path)
        })?;
        let wait_merging_threads_elapsed_ms = watcher.elapsed().as_millis();

        let seg_ids = index.searchable_segment_ids().with_context(|| {
            format!(
                "searchable_segment_ids failed. frag_path={}",
                self.frag_path
            )
        })?;

        let finalize_elapsed_ms = finalize_watcher.elapsed().as_millis();
        WRITER_FRAG_OPERATION_LATENCY
            .with_label_values(&["finalize"])
            .observe(finalize_watcher.elapsed().as_secs_f64());
        info!(
            "finalize_frag_success, {:?}, frag_path={}, add_doc_ms={}, commit_ms={}, write_handle_set_ms={}, wait_merging_threads_ms={}, total_ms={}",
            self.shard_tracing_id,
            self.frag_path,
            add_doc_elapsed_ms,
            commit_elapsed_ms,
            write_handle_set_elapsed_ms,
            wait_merging_threads_elapsed_ms,
            finalize_elapsed_ms
        );
        Ok(FragResult {
            seg_ids,
            handle_stat: tantivy_doc_res.handle_stat,
            finalize_elapsed_ms,
        })
    }

    pub fn get_buffered_stat(&self) -> Option<DocStat> {
        self.tantivy_doc_buffer.as_ref().map(|b| b.stat())
    }

    pub fn get_s3_frag_path(&self, s3_prefix: &str) -> String {
        let short_frag_path = self
            .frag_path
            .strip_prefix(&self.frag_base_path)
            .unwrap()
            .to_string()
            .trim_start_matches('/')
            .to_string();
        format!("{}/{}", s3_prefix, short_frag_path)
    }

    pub fn set_last_cdc_file(&mut self, file: Option<CDCFileName>) {
        self.last_cdc_file = file;
        if self.is_empty() && self.last_cdc_file.is_some() && self.pending_progress_since.is_none()
        {
            self.pending_progress_since = Some(TokioInstant::now());
        } else if !self.is_empty() || self.last_cdc_file.is_none() {
            self.pending_progress_since = None;
        }
    }

    pub fn last_cdc_file(&self) -> &Option<CDCFileName> {
        &self.last_cdc_file
    }

    pub fn progress_flush_deadline(&self, interval: std::time::Duration) -> Option<TokioInstant> {
        if !self.is_empty() || self.last_cdc_file.is_none() {
            return None;
        }
        let pending_progress_since = self.pending_progress_since?;
        let elapsed = pending_progress_since.elapsed();
        if elapsed >= interval {
            Some(TokioInstant::now())
        } else {
            Some(TokioInstant::now() + (interval - elapsed))
        }
    }

    pub fn take_progress_file(&mut self) -> Option<CDCFileName> {
        if !self.is_empty() {
            return None;
        }
        self.pending_progress_since = None;
        self.last_cdc_file.take()
    }

    pub fn frag_path(&self) -> String {
        self.frag_path.clone()
    }

    pub fn is_base(&self) -> bool {
        self.is_base
    }

    pub fn frag_base_path(&self) -> String {
        self.frag_base_path.clone()
    }

    pub fn is_empty(&self) -> bool {
        self.tantivy_doc_buffer.is_none()
    }
}