use super::Triplestore;
use crate::errors::TriplestoreError;
use oxrdf::NamedNode;
use oxrdfio::{RdfFormat, RdfSerializer};
use polars::prelude::{by_name, col};
use polars_core::datatypes::DataType;
use polars_core::frame::DataFrame;
use polars_core::POOL;
use representation::cats::maybe_decode_complex_expr;
use representation::dataset::NamedGraph;
use representation::polars_to_rdf::{
    date_column_to_strings, datetime_column_to_strings, global_df_as_triples,
};
use representation::{
    LANG_STRING_LANG_FIELD, LANG_STRING_VALUE_FIELD, OBJECT_COL_NAME, SUBJECT_COL_NAME,
};
use std::collections::HashMap;
use std::io::Write;

mod fast_ntriples;
mod pretty_turtle;
mod serializers;

const CHUNK_SIZE: usize = 1_024;

impl Triplestore {
    pub fn write_triples<W: Write>(
        &mut self,
        buf: &mut W,
        format: RdfFormat,
        graph: &NamedGraph,
        prefixes: &HashMap<String, NamedNode>,
    ) -> Result<(), TriplestoreError> {
        self.check_graph_exists(graph)?;
        if RdfFormat::NTriples == format {
            let n_threads = POOL.current_num_threads();
            for (predicate, df_map) in self.graph_triples_map.get(graph).unwrap() {
                let predicate_string = predicate.to_string();
                let predicate_bytes = predicate_string.as_bytes();
                for ((subject_type, object_type), tt) in df_map {
                    let types = if object_type.is_lang_string() {
                        HashMap::from([
                            (SUBJECT_COL_NAME.to_string(), subject_type.clone()),
                            (LANG_STRING_VALUE_FIELD.to_string(), object_type.clone()),
                            (LANG_STRING_LANG_FIELD.to_string(), object_type.clone()),
                        ])
                    } else {
                        HashMap::from([
                            (SUBJECT_COL_NAME.to_string(), subject_type.clone()),
                            (OBJECT_COL_NAME.to_string(), object_type.clone()),
                        ])
                    };
                    for (mut lf, _) in tt.get_lazy_frames(&None, &None)? {
                        lf = lf.with_columns([
                            maybe_decode_complex_expr(
                                col(SUBJECT_COL_NAME),
                                subject_type,
                                &subject_type.default_stored_cat_state(),
                                self.global_cats.clone(),
                            )
                            .alias(SUBJECT_COL_NAME),
                            maybe_decode_complex_expr(
                                col(OBJECT_COL_NAME),
                                object_type,
                                &object_type.default_stored_cat_state(),
                                self.global_cats.clone(),
                            )
                            .alias(OBJECT_COL_NAME),
                        ]);
                        if object_type.is_lang_string() {
                            lf = lf
                                .unnest(by_name([OBJECT_COL_NAME], true, false), None)
                                .select([
                                    col(SUBJECT_COL_NAME),
                                    col(LANG_STRING_VALUE_FIELD),
                                    col(LANG_STRING_LANG_FIELD),
                                ]);
                        } else {
                            lf = lf.select([col(SUBJECT_COL_NAME), col(OBJECT_COL_NAME)]);
                        };
                        let mut df = lf.collect().unwrap();
                        if !object_type.is_lang_string() {
                            convert_datelike_to_string(&mut df, OBJECT_COL_NAME);
                        }

                        fast_ntriples::write_triples_in_df(
                            buf,
                            &df,
                            predicate_bytes,
                            &types,
                            CHUNK_SIZE,
                            n_threads,
                        )
                        .unwrap();
                    }
                }
            }
        } else if RdfFormat::Turtle == format {
            self.write_pretty_turtle(buf, graph, prefixes)?;
        } else {
            let mut writer = RdfSerializer::from_format(format).for_writer(buf);

            for (predicate, df_map) in self.graph_triples_map.get(graph).unwrap() {
                for ((subject_type, object_type), tt) in df_map {
                    for (lf, _) in tt.get_lazy_frames(&None, &None)? {
                        let triples = global_df_as_triples(
                            lf.collect().unwrap(),
                            subject_type.clone(),
                            object_type.clone(),
                            predicate,
                            self.global_cats.clone(),
                        );
                        for t in &triples {
                            writer.serialize_triple(t).map_err(|x| {
                                TriplestoreError::WriteTurtleError(format!(
                                    "Error serializing triple {}: {}",
                                    t.to_string(),
                                    x.to_string()
                                ))
                            })?;
                        }
                    }
                }
            }
            writer.finish().map_err(|x| {
                TriplestoreError::WriteTurtleError(format!(
                    "Error finishing turtle serialization: {}",
                    x.to_string()
                ))
            })?;
        }
        Ok(())
    }

    pub(crate) fn check_graph_exists(&self, graph: &NamedGraph) -> Result<(), TriplestoreError> {
        if !self.graph_triples_map.contains_key(graph) {
            Err(TriplestoreError::GraphDoesNotExist(graph.to_string()))
        } else {
            Ok(())
        }
    }
}

pub fn convert_datelike_to_string(df: &mut DataFrame, c: &str) {
    match df.column(c).unwrap().dtype() {
        DataType::Date => {
            df.with_column(date_column_to_strings(df.column(c).unwrap()))
                .unwrap();
        }
        DataType::Datetime(_, tz_opt) => {
            df.with_column(datetime_column_to_strings(df.column(c).unwrap(), tz_opt))
                .unwrap();
        }
        _ => {}
    }
}
