use std::sync::Arc;

use polars_core::prelude::CompatLevel;
use polars_error::PolarsResult;
use polars_parquet::parquet::error::ParquetResult;
use polars_parquet::read::ParquetError;
use polars_parquet::write::{
    ColumnWriteOptions, CompressedPage, Compressor, SchemaDescriptor, WriteOptions,
    array_to_columns,
};
use polars_utils::UnitVec;

use crate::async_executor::{self, TaskPriority};
use crate::async_primitives::connector;
use crate::async_primitives::opt_spawned_future::parallelize_first_to_local;
use crate::nodes::io_sinks2::components::sink_morsel::SinkMorsel;
use crate::nodes::io_sinks2::writers::parquet::EncodedRowGroup;

pub struct RowGroupEncoder {
    pub morsel_rx: connector::Receiver<SinkMorsel>,
    pub encoded_row_group_tx:
        tokio::sync::mpsc::Sender<async_executor::AbortOnDropHandle<PolarsResult<EncodedRowGroup>>>,
    pub schema_descriptor: Arc<SchemaDescriptor>,
    pub write_options: WriteOptions,
    pub column_options: Arc<Vec<ColumnWriteOptions>>,
    pub num_leaf_columns: usize,
}

impl RowGroupEncoder {
    pub async fn run(self) -> PolarsResult<()> {
        let RowGroupEncoder {
            mut morsel_rx,
            encoded_row_group_tx,
            schema_descriptor,
            write_options,
            column_options,
            num_leaf_columns,
        } = self;

        while let Ok(morsel) = morsel_rx.recv().await {
            let schema_descriptor = Arc::clone(&schema_descriptor);
            let column_options = Arc::clone(&column_options);

            let row_group_encode_handle = async_executor::AbortOnDropHandle::new(
                async_executor::spawn(TaskPriority::High, async move {
                    let (df, morsel_permit) = morsel.into_inner();

                    let mut data: Vec<Vec<CompressedPage>> = Vec::with_capacity(num_leaf_columns);

                    for fut in parallelize_first_to_local(
                        TaskPriority::High,
                        df.into_columns().into_iter().enumerate().map(|(i, c)| {
                            let schema_descriptor = Arc::clone(&schema_descriptor);
                            let column_options = Arc::clone(&column_options);

                            async move {
                                let parquet_type = &schema_descriptor.fields()[i];
                                let column_options = &column_options[i];
                                let array = c
                                    .as_materialized_series()
                                    .rechunk()
                                    .to_arrow(0, CompatLevel::newest());

                                let mut data: UnitVec<Vec<CompressedPage>> =
                                    UnitVec::with_capacity(num_leaf_columns);

                                for encode_page_iter in array_to_columns(
                                    array,
                                    parquet_type.clone(),
                                    column_options,
                                    write_options,
                                )? {
                                    let compressed_pages: Vec<CompressedPage> =
                                        Compressor::new_from_vec(
                                            encode_page_iter.map(|result| {
                                                result.map_err(|e| {
                                                    ParquetError::FeatureNotSupported(format!(
                                                        "reraised in polars: {e}",
                                                    ))
                                                })
                                            }),
                                            write_options.compression,
                                            vec![],
                                        )
                                        .collect::<ParquetResult<_>>()?;

                                    data.push(compressed_pages)
                                }

                                PolarsResult::Ok(data)
                            }
                        }),
                    ) {
                        data.extend(fut.await?);
                    }

                    Ok(EncodedRowGroup {
                        data,
                        morsel_permit,
                    })
                }),
            );

            if encoded_row_group_tx
                .send(row_group_encode_handle)
                .await
                .is_err()
            {
                break;
            }
        }

        Ok(())
    }
}
