use std::cmp::Reverse;
use std::path::PathBuf;

use polars_core::frame::DataFrame;
use polars_core::schema::SchemaRef;
use polars_error::PolarsResult;
use polars_expr::state::ExecutionState;
use polars_io::SerWriter;
use polars_io::cloud::CloudOptions;
use polars_io::prelude::{CsvWriter, CsvWriterOptions};
use polars_io::utils::file::AsyncWriteable;
use polars_plan::dsl::SinkOptions;
use polars_utils::priority::Priority;

use super::{SinkInputPort, SinkNode};
use crate::async_executor::spawn;
use crate::async_primitives::connector::Receiver;
use crate::nodes::io_sinks::parallelize_receive_task;
use crate::nodes::{JoinHandle, PhaseOutcome, TaskPriority};

pub struct CsvSinkNode {
    path: PathBuf,
    schema: SchemaRef,
    sink_options: SinkOptions,
    write_options: CsvWriterOptions,
    cloud_options: Option<CloudOptions>,
}
impl CsvSinkNode {
    pub fn new(
        path: PathBuf,
        schema: SchemaRef,
        sink_options: SinkOptions,
        write_options: CsvWriterOptions,
        cloud_options: Option<CloudOptions>,
    ) -> Self {
        Self {
            path,
            schema,
            sink_options,
            write_options,
            cloud_options,
        }
    }
}

impl SinkNode for CsvSinkNode {
    fn name(&self) -> &str {
        "csv_sink"
    }

    fn is_sink_input_parallel(&self) -> bool {
        true
    }

    fn spawn_sink(
        &mut self,
        num_pipelines: usize,
        recv_port_rx: Receiver<(PhaseOutcome, SinkInputPort)>,
        _state: &ExecutionState,
        join_handles: &mut Vec<JoinHandle<PolarsResult<()>>>,
    ) {
        let (pass_rxs, mut io_rx) = parallelize_receive_task(
            join_handles,
            recv_port_rx,
            num_pipelines,
            self.sink_options.maintain_order,
        );

        // 16MB
        const DEFAULT_ALLOCATION_SIZE: usize = 1 << 24;

        // Encode task.
        //
        // Task encodes the columns into their corresponding CSV encoding.
        join_handles.extend(pass_rxs.into_iter().map(|mut pass_rx| {
            let schema = self.schema.clone();
            let options = self.write_options.clone();

            spawn(TaskPriority::High, async move {
                // Amortize the allocations over time. If we see that we need to do way larger
                // allocations, we adjust to that over time.
                let mut allocation_size = DEFAULT_ALLOCATION_SIZE;
                let options = options.clone();

                while let Ok((mut rx, mut lin_tx)) = pass_rx.recv().await {
                    while let Ok(morsel) = rx.recv().await {
                        let (df, seq, _, consume_token) = morsel.into_inner();

                        let mut buffer = Vec::with_capacity(allocation_size);
                        let mut writer = CsvWriter::new(&mut buffer)
                            .include_bom(false) // Handled once in the IO task.
                            .include_header(false) // Handled once in the IO task.
                            .with_separator(options.serialize_options.separator)
                            .with_line_terminator(options.serialize_options.line_terminator.clone())
                            .with_quote_char(options.serialize_options.quote_char)
                            .with_datetime_format(options.serialize_options.datetime_format.clone())
                            .with_date_format(options.serialize_options.date_format.clone())
                            .with_time_format(options.serialize_options.time_format.clone())
                            .with_float_scientific(options.serialize_options.float_scientific)
                            .with_float_precision(options.serialize_options.float_precision)
                            .with_null_value(options.serialize_options.null.clone())
                            .with_quote_style(options.serialize_options.quote_style)
                            .n_threads(1) // Disable rayon parallelism
                            .batched(&schema)?;

                        writer.write_batch(&df)?;

                        allocation_size = allocation_size.max(buffer.len());
                        if lin_tx.insert(Priority(Reverse(seq), buffer)).await.is_err() {
                            return Ok(());
                        }
                        drop(consume_token); // Keep the consume_token until here to increase the
                        // backpressure.
                    }
                }

                PolarsResult::Ok(())
            })
        }));

        // IO task.
        //
        // Task that will actually do write to the target file.
        let path = self.path.clone();
        let sink_options = self.sink_options.clone();
        let schema = self.schema.clone();
        let include_header = self.write_options.include_header;
        let include_bom = self.write_options.include_bom;
        let cloud_options = self.cloud_options.clone();
        let io_task = polars_io::pl_async::get_runtime().spawn(async move {
            use tokio::io::AsyncWriteExt;

            if sink_options.mkdir {
                polars_io::utils::mkdir::tokio_mkdir_recursive(path.as_path()).await?;
            }

            let mut file = polars_io::utils::file::Writeable::try_new(
                path.to_str().unwrap(),
                cloud_options.as_ref(),
            )?;

            // Write the header
            if include_header || include_bom {
                let mut writer = CsvWriter::new(&mut *file)
                    .include_bom(include_bom)
                    .include_header(include_header)
                    .n_threads(1) // Disable rayon parallelism
                    .batched(&schema)?;
                writer.write_batch(&DataFrame::empty_with_schema(&schema))?;
            }

            let mut file = file.try_into_async_writeable()?;

            while let Ok(mut lin_rx) = io_rx.recv().await {
                while let Some(Priority(_, buffer)) = lin_rx.get().await {
                    file.write_all(&buffer).await?;
                }
            }

            if let AsyncWriteable::Local(file) = &mut file {
                polars_io::utils::sync_on_close::tokio_sync_on_close(
                    sink_options.sync_on_close,
                    file,
                )
                .await?;
            }

            file.close().await?;

            PolarsResult::Ok(())
        });
        join_handles.push(spawn(TaskPriority::Low, async move {
            io_task
                .await
                .unwrap_or_else(|e| Err(std::io::Error::from(e).into()))
        }));
    }
}
