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

use anyhow::{Context, Result, anyhow};
use backon::{ExponentialBuilder, Retryable};
use log::{info, warn};
use tici_common::index::IndexKey;
use tici_shard::{
    config::WorkerNodeConfig,
    metadata::metadata_store::MetadataStore,
    utils::{
        frag_path_generator::FragPathGenerator, import_task_tracing_id,
        meta_client_pool::MetaClientPool, tantivy_doc::TantivyDoc,
    },
    writer::{
        frag_batch_uploader::FragBatchUploader,
        frag_uploader::{FragUploadTask, FragUploader},
        frag_writer::FragWriterBuilder,
    },
};
use tokio::sync::mpsc;

use crate::{
    ImportTaskSummary,
    file::{DEFAULT_MAX_KV_BYTES, File},
    metrics::{IMPORT_INTO_BYTES_PROCESSED_TOTAL, IMPORT_INTO_ROWS_PROCESSED_TOTAL},
};

pub struct Writer {
    pub worker_config: WorkerNodeConfig,
    pub s3_bucket: String,
    pub import_file: String,
    pub frag_upload_prefix: String,
    pub index_key: IndexKey,
    pub frag_path_gen: Arc<FragPathGenerator>,
    pub task_id: u64,
    pub metadata_store: Arc<MetadataStore>,
    pub frag_batch_uploader: Arc<FragBatchUploader>,
}

impl Writer {
    pub async fn write(
        &self,
        meta_client_pool: Arc<dyn MetaClientPool>,
    ) -> Result<ImportTaskSummary> {
        let mut meta_client = meta_client_pool.get_meta_client().await?;
        let response = meta_client.get_index_meta(&self.index_key).await?;
        let index_info = response
            .index_info
            .ok_or_else(|| anyhow!("index_info not found, index_key={:?}", self.index_key))?;
        let parser_info = index_info.parser_info.clone();
        let index_created_ts = index_info.created_ts;

        let backend_storage = self.frag_batch_uploader.backend_storage();
        let (frag_writer_sender, frag_writer_receiver) = mpsc::channel(100);
        let frag_uploader = FragUploader::new(
            import_task_tracing_id(self.task_id),
            &self.index_key,
            0, // shard_id is not used for frags generated by import
            self.worker_config.server.get_advertise_addr().to_string(),
            meta_client_pool,
            self.frag_batch_uploader.clone(),
            self.s3_bucket.clone(),
            self.frag_upload_prefix.clone(),
            frag_writer_receiver,
            true,
            self.metadata_store.clone(),
        );

        let (frag_writer, summary, num_docs) = self
            .build_with_retry(backend_storage.clone(), parser_info, index_created_ts)
            .await?;
        if num_docs == 0 {
            info!(
                "ImportWriter - no docs read, skip uploading empty base fragment, task_id={} table_id={} index_id={} file={}",
                self.task_id, self.index_key.table_id, self.index_key.index_id, self.import_file
            );
            drop(frag_writer_sender);
            frag_uploader.wait_finished().await?;
            return Ok(summary);
        }
        info!(
            "ImportWriter - sending frag_writer to uploader, task_id={} table_id={} index_id={} n_docs={} frag_path={}",
            self.task_id,
            self.index_key.table_id,
            self.index_key.index_id,
            num_docs,
            frag_writer.frag_path()
        );
        frag_writer_sender
            .send(FragUploadTask::Frag(frag_writer))
            .await?;
        // close the channel so frag_uploader can finish
        drop(frag_writer_sender);
        // wait for the uploader finish uploading fragment
        frag_uploader.wait_finished().await?;
        Ok(summary)
    }

    async fn build_with_retry(
        &self,
        backend_storage: Arc<dyn storage_backend::storage::BackendStorage>,
        parser_info: Option<tici_grpc::ParserInfo>,
        index_created_ts: u64,
    ) -> Result<(
        tici_shard::writer::frag_writer::FragWriter,
        ImportTaskSummary,
        u64,
    )> {
        let max_attempts: usize = 10;
        let builder = ExponentialBuilder::default()
            .with_factor(2.0)
            .with_min_delay(Duration::from_millis(100))
            .with_max_delay(Duration::from_secs(3))
            .with_max_times(max_attempts - 1);
        let mut attempt = 0usize;

        let operation =
            (move || {
                self.build_once(
                    backend_storage.clone(),
                    parser_info.clone(),
                    index_created_ts,
                )
            })
                .retry(&builder)
                .notify(move |err: &anyhow::Error, dur| {
                    attempt += 1;
                    warn!(
                        "ImportWriter - read/build failed, retrying in {:?}, task_id={} file={} attempt={}/{} err={:?}",
                        dur,
                        self.task_id,
                        self.import_file,
                        attempt,
                        max_attempts,
                        err
                    );
                });

        operation.await.with_context(|| {
            format!(
                "ImportWriter - read/build failed after {} attempts",
                max_attempts
            )
        })
    }

    async fn build_once(
        &self,
        backend_storage: Arc<dyn storage_backend::storage::BackendStorage>,
        parser_info: Option<tici_grpc::ParserInfo>,
        index_created_ts: u64,
    ) -> Result<(
        tici_shard::writer::frag_writer::FragWriter,
        ImportTaskSummary,
        u64,
    )> {
        info!(
            "ImportWriter - open file from S3, task_id={} file={}",
            self.task_id, self.import_file
        );
        let mut file = File::open(
            self.s3_bucket.clone(),
            self.import_file.clone(),
            backend_storage,
            self.index_key.keyspace_id,
            self.index_key.index_id,
            DEFAULT_MAX_KV_BYTES,
            parser_info,
            index_created_ts,
        )
        .await?;
        let sort_fields = file.get_index_meta().get_sort_fields().to_vec();
        let mut frag_writer = {
            let frag_path = self.frag_path_gen.new_path(&self.index_key).await?;
            FragWriterBuilder::new()
                .frag_base_path(self.frag_path_gen.frag_base_path())
                .frag_path(&frag_path)
                .schema(file.get_index_meta().get_tantivy_schema().clone())
                .index_num_threads(self.worker_config.compaction.index_num_threads)
                .index_mem_budget(self.worker_config.compaction.index_mem_budget.0 as usize)
                .is_base(true)
                .analyzer_manager(file.get_index_meta().get_analyzer_manager().clone())
                .sort_fields(sort_fields)
                .build()?
        };

        let mut num_docs: u64 = 0;
        while let Some(doc) = file
            .next()
            .await
            .with_context(|| file_next_context(self.task_id, &self.index_key, &self.import_file))?
        {
            log::trace!("read doc from file, doc={:?}", doc);
            let tantivy_doc = TantivyDoc::new(&doc, file.get_index_meta())?;
            frag_writer.add_tantivy_doc(tantivy_doc);
            num_docs += 1;
        }
        if num_docs > 0 {
            IMPORT_INTO_ROWS_PROCESSED_TOTAL.inc_by(num_docs);
        }
        let bytes_processed = file.bytes_processed();
        if bytes_processed > 0 {
            IMPORT_INTO_BYTES_PROCESSED_TOTAL.inc_by(bytes_processed);
        }
        let summary = ImportTaskSummary {
            rows_processed: num_docs,
            bytes_processed,
            input_file_size: file.input_file_size(),
        };
        Ok((frag_writer, summary, num_docs))
    }
}

fn file_next_context(task_id: u64, index_key: &IndexKey, import_file: &str) -> String {
    format!(
        "ImportWriter - failed to read next doc, task_id={} index_key={:?} import_file={}",
        task_id, index_key, import_file
    )
}