use std::collections::HashMap;
use std::fs::File;
use std::io::BufWriter;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};

use arrow::csv::{Writer, WriterBuilder};
use arrow::datatypes::Schema;
use arrow::error::ArrowError;
use arrow::record_batch::RecordBatch;
use mpera::programoutput::{OutputValue, ProgramOutput};

use crate::sink::Sink;

pub struct CsvSink {
    base: PathBuf,
    inner: Mutex<Inner>,
}

struct Inner {
    writers: HashMap<String, (Writer<BufWriter<File>>, Arc<Schema>)>,
}

impl CsvSink {
    pub fn new(base: PathBuf) -> Result<Self, ArrowError> {
        Ok(Self {
            base,
            inner: Mutex::new(Inner {
                writers: HashMap::new(),
            }),
        })
    }

    fn write_named_batch(&self, name: &str, batch: &RecordBatch) -> Result<(), ArrowError> {
        let _ = std::fs::create_dir_all(&self.base);
        let mut inner = self.inner.lock().unwrap();

        if let Some((w, _schema)) = inner.writers.get_mut(name) {
            w.write(batch)?;
            return Ok(());
        }

        let path = self.base.join(format!("{name}.csv"));
        let file = BufWriter::new(File::create(&path)?);

        let mut writer = WriterBuilder::new().with_header(true).build(file);

        writer.write(batch)?;
        inner
            .writers
            .insert(name.to_string(), (writer, batch.schema()));

        Ok(())
    }
}

impl Sink for CsvSink {
    fn sink(&self, output: ProgramOutput) -> crate::Result<()> {
        for (name, value) in output.0 {
            match value {
                OutputValue::Frame(batch) => self.write_named_batch(&name, &batch)?,
                OutputValue::Groups(groups) => {
                    for (group_key, batch) in groups {
                        let safe_key = group_key
                            .chars()
                            .map(|ch| match ch {
                                '/' | '\\' | ':' | '\u{1f}' => '_',
                                _ => ch,
                            })
                            .collect::<String>();
                        self.write_named_batch(&format!("{name}__{safe_key}"), &batch)?;
                    }
                }
            }
        }
        Ok(())
    }

    fn finish(&mut self) -> crate::Result<()> {
        Ok(())
    }
}
