use crate::config::injection::DIService;
use crate::config::{Config, ConfigObj, RepartitionStrategy};
use crate::import::ImportService;
use crate::metastore::job::{Job, JobType};
use crate::metastore::table::Table;
use crate::metastore::{MetaStore, RowKey, TableId};
use crate::queryplanner::trace_data_loaded::DataLoadedSize;
use crate::store::compaction::CompactionService;
use crate::store::ChunkDataStore;
use crate::{app_metrics, CubeError};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Duration;

#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct JobProcessResult {
    data_loaded_size: usize,
}

impl JobProcessResult {
    pub fn new(data_loaded_size: usize) -> Self {
        Self { data_loaded_size }
    }

    pub fn data_loaded_size(&self) -> usize {
        self.data_loaded_size
    }
}

impl Default for JobProcessResult {
    fn default() -> Self {
        Self {
            data_loaded_size: 0,
        }
    }
}

#[async_trait]
pub trait JobProcessor: DIService + Send + Sync {
    async fn wait_processing_loops(&self);
    async fn stop_processing_loops(&self) -> Result<(), CubeError>;
    async fn process_job(&self, job: Job) -> Result<JobProcessResult, CubeError>;
}

pub struct JobProcessorImpl {
    processor: Arc<JobIsolatedProcessor>,
}

impl JobProcessorImpl {
    pub fn new(
        meta_store: Arc<dyn MetaStore>,
        chunk_store: Arc<dyn ChunkDataStore>,
        compaction_service: Arc<dyn CompactionService>,
        import_service: Arc<dyn ImportService>,
        config_obj: Arc<dyn ConfigObj>,
    ) -> Arc<Self> {
        Arc::new(Self {
            processor: JobIsolatedProcessor::new(
                meta_store,
                chunk_store,
                compaction_service,
                import_service,
                config_obj,
            ),
        })
    }
}

#[async_trait]
impl JobProcessor for JobProcessorImpl {
    async fn wait_processing_loops(&self) {}

    async fn stop_processing_loops(&self) -> Result<(), CubeError> {
        Ok(())
    }

    async fn process_job(&self, job: Job) -> Result<JobProcessResult, CubeError> {
        self.processor.process_separate_job(&job).await
    }
}

crate::di_service!(JobProcessorImpl, [JobProcessor]);

pub struct JobIsolatedProcessor {
    meta_store: Arc<dyn MetaStore>,
    chunk_store: Arc<dyn ChunkDataStore>,
    compaction_service: Arc<dyn CompactionService>,
    import_service: Arc<dyn ImportService>,
    config_obj: Arc<dyn ConfigObj>,
}

impl JobIsolatedProcessor {
    pub fn new(
        meta_store: Arc<dyn MetaStore>,
        chunk_store: Arc<dyn ChunkDataStore>,
        compaction_service: Arc<dyn CompactionService>,
        import_service: Arc<dyn ImportService>,
        config_obj: Arc<dyn ConfigObj>,
    ) -> Arc<Self> {
        Arc::new(Self {
            meta_store,
            chunk_store,
            compaction_service,
            import_service,
            config_obj,
        })
    }

    pub async fn new_from_config(config: &Config) -> Arc<Self> {
        Self::new(
            config.injector().get_service_typed().await,
            config.injector().get_service_typed().await,
            config.injector().get_service_typed().await,
            config.injector().get_service_typed().await,
            config.injector().get_service_typed().await,
        )
    }

    pub async fn process_separate_job(&self, job: &Job) -> Result<JobProcessResult, CubeError> {
        match job.job_type() {
            JobType::PartitionCompaction => {
                if let RowKey::Table(TableId::Partitions, partition_id) = job.row_reference() {
                    let compaction_service = self.compaction_service.clone();
                    let partition_id = *partition_id;
                    let data_loaded_size = DataLoadedSize::new();
                    app_metrics::JOBS_PARTITION_COMPACTION.add(1);
                    let r = compaction_service
                        .compact(partition_id, data_loaded_size.clone())
                        .await;
                    if let Err(e) = r {
                        app_metrics::JOBS_PARTITION_COMPACTION_FAILURES.add(1);
                        return Err(e);
                    }
                    app_metrics::JOBS_PARTITION_COMPACTION_COMPLETED.add(1);

                    Ok(JobProcessResult::new(data_loaded_size.get()))
                } else {
                    Self::fail_job_row_key(job)
                }
            }
            JobType::MultiPartitionSplit => {
                if let RowKey::Table(TableId::MultiPartitions, id) = job.row_reference() {
                    let compaction_service = self.compaction_service.clone();
                    let id = *id;
                    app_metrics::JOBS_MULTI_PARTITION_SPLIT.add(1);
                    let r = compaction_service.split_multi_partition(id).await;
                    if let Err(e) = r {
                        app_metrics::JOBS_MULTI_PARTITION_SPLIT_FAILURES.add(1);
                        return Err(e);
                    }
                    app_metrics::JOBS_MULTI_PARTITION_SPLIT_COMPLETED.add(1);
                    Ok(JobProcessResult::default())
                } else {
                    Self::fail_job_row_key(job)
                }
            }
            JobType::FinishMultiSplit => {
                if let RowKey::Table(TableId::MultiPartitions, multi_part_id) = job.row_reference()
                {
                    let meta_store = self.meta_store.clone();
                    let compaction_service = self.compaction_service.clone();
                    let multi_part_id = *multi_part_id;
                    for p in meta_store.find_unsplit_partitions(multi_part_id).await? {
                        app_metrics::JOBS_FINISH_MULTI_SPLIT.add(1);
                        let r = compaction_service
                            .finish_multi_split(multi_part_id, p)
                            .await;
                        if let Err(e) = r {
                            app_metrics::JOBS_FINISH_MULTI_SPLIT_FAILURES.add(1);
                            return Err(e);
                        }
                        app_metrics::JOBS_FINISH_MULTI_SPLIT_COMPLETED.add(1);
                    }

                    Ok(JobProcessResult::default())
                } else {
                    Self::fail_job_row_key(job)
                }
            }
            JobType::TableImport => {
                if let RowKey::Table(TableId::Tables, table_id) = job.row_reference() {
                    let import_service = self.import_service.clone();
                    let table_id = *table_id;
                    import_service.import_table(table_id).await?;
                    Ok(JobProcessResult::default())
                } else {
                    Self::fail_job_row_key(job)
                }
            }
            JobType::TableImportCSV(location) => {
                if Table::is_stream_location(&location) {
                    return Err(CubeError::internal(
                        "Streaming import cannot be processed in separate process".to_string(),
                    ));
                }
                if let RowKey::Table(TableId::Tables, table_id) = job.row_reference() {
                    let table_id = *table_id;
                    let import_service = self.import_service.clone();
                    let location = location.to_string();
                    let data_loaded_size = Some(DataLoadedSize::new());
                    import_service
                        .clone()
                        .import_table_part(table_id, &location, data_loaded_size.clone())
                        .await?;
                    Ok(JobProcessResult::new(
                        data_loaded_size.map_or(0, |d| d.get()),
                    ))
                } else {
                    Self::fail_job_row_key(job)
                }
            }
            JobType::RepartitionChunk => {
                if let RowKey::Table(TableId::Chunks, chunk_id) = job.row_reference() {
                    let chunk_id = *chunk_id;
                    let chunk = self.meta_store.get_chunk(chunk_id).await?;
                    if chunk.get_row().in_memory() {
                        return Err(CubeError::internal(
                            "In-memory chunk cannot be repartitioned in separate process"
                                .to_string(),
                        ));
                    }
                    let data_loaded_size = DataLoadedSize::new();
                    app_metrics::JOBS_REPARTITION_CHUNK.add(1);
                    // PerPartition reuses the RepartitionChunk job with the smallest chunk
                    // as the anchor (one job per partition) and merges all its chunks;
                    // PerChunk gets one job per chunk and repartitions just that chunk. We
                    // overload the existing RepartitionChunk type (no new JobType) so an
                    // older binary stays able to deserialize it across channel switches.
                    let r = if self.config_obj.repartition_strategy()
                        == RepartitionStrategy::PerPartition
                    {
                        let partition_id = chunk.get_row().get_partition_id();
                        let time_budget = Duration::from_secs(
                            self.config_obj.repartition_chunks_time_budget_secs(),
                        );
                        self.chunk_store
                            .repartition_partition_chunks(
                                partition_id,
                                chunk_id,
                                time_budget,
                                data_loaded_size.clone(),
                            )
                            .await
                    } else {
                        self.chunk_store
                            .repartition_chunk(chunk_id, data_loaded_size.clone())
                            .await
                    };
                    if let Err(e) = r {
                        app_metrics::JOBS_REPARTITION_CHUNK_FAILURES.add(1);
                        return Err(e);
                    }
                    app_metrics::JOBS_REPARTITION_CHUNK_COMPLETED.add(1);
                    Ok(JobProcessResult::new(data_loaded_size.get()))
                } else {
                    Self::fail_job_row_key(job)
                }
            }
            JobType::RepartitionRange(end_chunk_id) => {
                if let RowKey::Table(TableId::Chunks, start_chunk_id) = job.row_reference() {
                    let start_chunk_id = *start_chunk_id;
                    let end_chunk_id = *end_chunk_id;
                    let data_loaded_size = DataLoadedSize::new();
                    app_metrics::JOBS_REPARTITION_CHUNK.add(1);
                    let r = self
                        .chunk_store
                        .repartition_chunk_range(
                            start_chunk_id,
                            end_chunk_id,
                            data_loaded_size.clone(),
                        )
                        .await;
                    if let Err(e) = r {
                        app_metrics::JOBS_REPARTITION_CHUNK_FAILURES.add(1);
                        return Err(e);
                    }
                    app_metrics::JOBS_REPARTITION_CHUNK_COMPLETED.add(1);
                    Ok(JobProcessResult::new(data_loaded_size.get()))
                } else {
                    Self::fail_job_row_key(job)
                }
            }
            _ => Err(CubeError::internal(format!(
                "Job {:?} cannot be processed in separate process",
                job.job_type()
            ))),
        }
    }

    fn fail_job_row_key(job: &Job) -> Result<JobProcessResult, CubeError> {
        Err(CubeError::internal(format!(
            "Incorrect row key for {:?}: {:?}",
            job,
            job.row_reference()
        )))
    }
}
