use crate::metastore::Column;
use crate::queryplanner::udfs::{registerable_aggregate_udfs_iter, registerable_scalar_udfs_iter};
use crate::CubeError;
use async_trait::async_trait;
use chrono::{TimeZone, Utc};
use chrono_tz::Tz;
use datafusion::arrow::array::{
    Array, StringArray, StringBuilder, TimestampMicrosecondArray, TimestampMicrosecondBuilder,
};
use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
use datafusion::catalog::Session;
use datafusion::common::TableReference;
use datafusion::config::ConfigOptions;
use datafusion::datasource::{provider_as_source, TableProvider, TableType};
use datafusion::error::DataFusionError;
use datafusion::execution::SessionStateDefaults;
use datafusion::logical_expr::{
    AggregateUDF, Expr, ScalarUDF, ScalarUDFImpl, Signature, TableSource, TypeSignature,
    Volatility, WindowUDF,
};
use datafusion::physical_plan::empty::EmptyExec;
use datafusion::physical_plan::ColumnarValue;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::scalar::ScalarValue;
use datafusion::sql::planner::ContextProvider;
use std::any::Any;
use std::collections::HashMap;
use std::fmt::{Debug, Formatter};
use std::sync::Arc;

#[derive(Debug, Clone)]
pub struct TopicTableProvider {
    topic: String,
    schema: SchemaRef,
    config_options: ConfigOptions,
    udfs: HashMap<String, Arc<ScalarUDF>>,
    udafs: HashMap<String, Arc<AggregateUDF>>,
    udwfs: HashMap<String, Arc<WindowUDF>>,
}

impl TopicTableProvider {
    pub fn new(topic: String, columns: &Vec<Column>) -> Self {
        let schema = Arc::new(Schema::new(
            columns
                .iter()
                .map(|c| c.clone().into())
                .collect::<Vec<Field>>(),
        ));
        let mut udfs = SessionStateDefaults::default_scalar_functions();
        udfs.extend(registerable_scalar_udfs_iter().map(Arc::new));
        udfs.push(Arc::new(
            ScalarUDF::new_from_impl(ParseTimestampFunc::new()),
        ));
        udfs.push(Arc::new(ScalarUDF::new_from_impl(ConvertTzFunc::new())));
        udfs.push(Arc::new(ScalarUDF::new_from_impl(
            FormatTimestampFunc::new(),
        )));

        let udfs = udfs
            .into_iter()
            .map(|udf| (udf.name().to_owned(), udf))
            .collect();

        let mut udafs = SessionStateDefaults::default_aggregate_functions();
        udafs.extend(registerable_aggregate_udfs_iter().map(Arc::new));

        let udafs = udafs
            .into_iter()
            .map(|udaf| (udaf.name().to_owned(), udaf))
            .collect();

        let udwfs = SessionStateDefaults::default_window_functions();
        let udwfs = udwfs
            .into_iter()
            .map(|udwf| (udwf.name().to_owned(), udwf))
            .collect();
        Self {
            topic,
            schema,
            config_options: ConfigOptions::default(),
            udfs,
            udafs,
            udwfs,
        }
    }
}

impl ContextProvider for TopicTableProvider {
    fn get_table_source(
        &self,
        name: TableReference,
    ) -> Result<Arc<dyn TableSource>, DataFusionError> {
        match name {
            TableReference::Bare { table } if table.as_ref() == self.topic => {
                Ok(provider_as_source(Arc::new(self.clone())))
            }
            _ => Err(DataFusionError::Plan(format!(
                "Topic table {} is not found",
                name
            ))),
        }
    }

    fn get_function_meta(&self, name: &str) -> Option<Arc<ScalarUDF>> {
        self.udfs.get(&name.to_ascii_lowercase()).cloned()
    }

    fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
        self.udafs.get(&name.to_ascii_lowercase()).cloned()
    }

    fn get_window_meta(&self, name: &str) -> Option<Arc<WindowUDF>> {
        self.udwfs.get(&name.to_ascii_lowercase()).cloned()
    }

    fn get_variable_type(&self, _variable_names: &[String]) -> Option<DataType> {
        None
    }

    fn options(&self) -> &ConfigOptions {
        &self.config_options
    }

    fn udf_names(&self) -> Vec<String> {
        self.udfs.keys().cloned().collect()
    }

    fn udaf_names(&self) -> Vec<String> {
        self.udafs.keys().cloned().collect()
    }

    fn udwf_names(&self) -> Vec<String> {
        self.udwfs.keys().cloned().collect()
    }
}

#[async_trait]
impl TableProvider for TopicTableProvider {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn table_type(&self) -> TableType {
        TableType::Base
    }

    async fn scan(
        &self,
        _state: &dyn Session,
        _projection: Option<&Vec<usize>>,
        _filters: &[Expr],
        _limit: Option<usize>,
    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
        Ok(Arc::new(EmptyExec::new(self.schema())))
    }
}

fn sql_format_to_strformat(sql_format: &str) -> String {
    sql_format
        .replace("yyyy", "%Y")
        .replace("MM", "%m")
        .replace("dd", "%d")
        .replace("HH", "%H")
        .replace("mm", "%M")
        .replace("ss", "%S")
        .replace(".SSS", "%.3f")
        .replace("X", "Z") //FIXME there are no option for timezone offset with Z as zero in rust
        //so we alway expect Z at the end of the string
        .replace("'T'", "T")
}

fn parse_timestamp_array(
    input: &StringArray,
    tz: &Tz,
    format: &str,
) -> Result<TimestampMicrosecondArray, DataFusionError> {
    let mut result = TimestampMicrosecondBuilder::new();
    for i in 0..input.len() {
        if input.is_null(i) {
            result.append_null();
        } else {
            #[allow(deprecated)]
            let parse_result = tz.datetime_from_str(input.value(i), &format);

            let ts = match parse_result {
                Ok(ts) => ts,
                Err(e) => {
                    return Err(DataFusionError::Execution(format!(
                        "Error while parsing timestamp `{}`: {}",
                        input.value(i),
                        e
                    )));
                }
            };
            result.append_value(ts.timestamp_micros());
        }
    }
    Ok(result.finish())
}

fn convert_tz_array(
    input: &TimestampMicrosecondArray,
    from_tz: &Tz,
    to_tz: &Tz,
) -> Result<TimestampMicrosecondArray, DataFusionError> {
    let mut result = TimestampMicrosecondBuilder::new();
    if from_tz == to_tz {
        for i in 0..input.len() {
            if input.is_null(i) {
                result.append_null();
            } else {
                result.append_value(input.value(i));
            }
        }
    } else {
        for i in 0..input.len() {
            if input.is_null(i) {
                result.append_null();
            } else {
                let time = Utc
                    .timestamp_nanos(input.value(i) as i64 * 1000)
                    .naive_local();
                let from = match from_tz.from_local_datetime(&time).earliest() {
                    Some(t) => t,
                    None => {
                        return Err(DataFusionError::Execution(format!(
                            "Can't convert timezone for timestamp {}",
                            input.value(i)
                        )));
                    }
                };
                let res = from.with_timezone(to_tz);
                result.append_value(res.naive_local().and_utc().timestamp_micros());
            }
        }
    }
    Ok(result.finish())
}
fn format_timestamp_array(
    input: &TimestampMicrosecondArray,
    format: &str,
) -> Result<StringArray, DataFusionError> {
    let mut result = StringBuilder::new();
    for i in 0..input.len() {
        if input.is_null(i) {
            result.append_null();
        } else {
            let time = Utc
                .timestamp_nanos(input.value(i) as i64 * 1000)
                .naive_local();
            result.append_value(format!("{}", time.format(format)));
        }
    }
    Ok(result.finish())
}

struct ParseTimestampFunc {
    signature: Signature,
}

impl Debug for ParseTimestampFunc {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "ParseTimestampFunc")
    }
}

impl ParseTimestampFunc {
    fn new() -> ParseTimestampFunc {
        ParseTimestampFunc {
            signature: Signature::one_of(
                vec![
                    TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8, DataType::Utf8]),
                    TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8]),
                ],
                Volatility::Stable,
            ),
        }
    }
}

impl ScalarUDFImpl for ParseTimestampFunc {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn name(&self) -> &str {
        "parse_timestamp"
    }

    fn signature(&self) -> &Signature {
        &self.signature
    }

    fn return_type(&self, _: &[DataType]) -> datafusion::common::Result<DataType> {
        Ok(DataType::Timestamp(TimeUnit::Microsecond, None))
    }

    fn invoke(&self, inputs: &[ColumnarValue]) -> datafusion::common::Result<ColumnarValue> {
        if inputs.len() < 2 || inputs.len() > 3 {
            return Err(DataFusionError::Execution(
                "Expected 2 or 3 arguments in PARSE_TIMESTAMP".to_string(),
            ));
        }

        let format = match &inputs[1] {
            ColumnarValue::Scalar(ScalarValue::Utf8(Some(v))) => sql_format_to_strformat(v),
            _ => {
                return Err(DataFusionError::Execution(
                    "Only scalar arguments are supported as format in PARSE_TIMESTAMP".to_string(),
                ));
            }
        };
        let tz: Tz = if inputs.len() == 3 {
            match &inputs[2] {
                ColumnarValue::Scalar(ScalarValue::Utf8(Some(s))) => s.parse().map_err(|_| {
                    CubeError::user(format!("Incorrect timezone {} in PARSE_TIMESTAMP", s))
                })?,
                _ => {
                    return Err(DataFusionError::Execution(
                        "Only scalar arguments are supported as timezone in PARSE_TIMESTAMP"
                            .to_string(),
                    ));
                }
            }
        } else {
            Tz::UTC
        };

        match &inputs[0] {
            ColumnarValue::Scalar(ScalarValue::Utf8(Some(s))) => {
                #[allow(deprecated)]
                let parse_result = tz.datetime_from_str(s, &format);

                let ts = match parse_result {
                    Ok(ts) => ts,
                    Err(e) => {
                        return Err(DataFusionError::Execution(format!(
                            "Error while parsing timestamp: {}",
                            e
                        )));
                    }
                };
                Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
                    Some(ts.timestamp_micros()),
                    None,
                )))
            }
            ColumnarValue::Array(t) if t.as_any().is::<StringArray>() => {
                let t = t.as_any().downcast_ref::<StringArray>().unwrap();
                Ok(ColumnarValue::Array(Arc::new(parse_timestamp_array(
                    &t, &tz, &format,
                )?)))
            }
            _ => {
                return Err(DataFusionError::Execution(
                    "First argument in PARSE_TIMESTAMP must be string or array of strings"
                        .to_string(),
                ));
            }
        }
    }
}

struct ConvertTzFunc {
    signature: Signature,
}

impl ConvertTzFunc {
    fn new() -> ConvertTzFunc {
        ConvertTzFunc {
            signature: Signature::exact(
                vec![
                    DataType::Timestamp(TimeUnit::Microsecond, None),
                    DataType::Utf8,
                    DataType::Utf8,
                ],
                Volatility::Stable,
            ),
        }
    }
}

impl Debug for ConvertTzFunc {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "ConvertTzFunc")
    }
}

impl ScalarUDFImpl for ConvertTzFunc {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn name(&self) -> &str {
        "convert_tz_ksql"
    }

    fn signature(&self) -> &Signature {
        &self.signature
    }

    fn return_type(&self, _: &[DataType]) -> datafusion::common::Result<DataType> {
        Ok(DataType::Timestamp(TimeUnit::Microsecond, None))
    }

    fn invoke(&self, inputs: &[ColumnarValue]) -> datafusion::common::Result<ColumnarValue> {
        if inputs.len() != 3 {
            return Err(DataFusionError::Execution(
                "Expected 3 arguments in CONVERT_TZ_KSQL".to_string(),
            ));
        }

        let from_tz: Tz = match &inputs[1] {
            ColumnarValue::Scalar(ScalarValue::Utf8(Some(s))) => s.parse().map_err(|_| {
                CubeError::user(format!("Incorrect timezone {} in CONVERT_TZ_KSQL", s))
            })?,
            _ => {
                return Err(DataFusionError::Execution(
                    "Only scalar arguments are supported as from_timezone in CONVERT_TZ_KSQL"
                        .to_string(),
                ));
            }
        };

        let to_tz: Tz = match &inputs[2] {
            ColumnarValue::Scalar(ScalarValue::Utf8(Some(s))) => s.parse().map_err(|_| {
                CubeError::user(format!("Incorrect timezone {} in CONVERT_TZ_KSQL", s))
            })?,
            _ => {
                return Err(DataFusionError::Execution(
                    "Only scalar arguments are supported as to_timezone in CONVERT_TZ_KSQL"
                        .to_string(),
                ));
            }
        };
        match &inputs[0] {
            ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(Some(t), None)) => {
                if from_tz == to_tz {
                    Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
                        Some(*t),
                        None,
                    )))
                } else {
                    let time = Utc.timestamp_nanos(*t * 1000).naive_local();
                    let from = match from_tz.from_local_datetime(&time).earliest() {
                        Some(t) => t,
                        None => {
                            return Err(DataFusionError::Execution(format!(
                                "Can't convert timezone for timestamp {}",
                                t
                            )));
                        }
                    };
                    let result = from.with_timezone(&to_tz);
                    Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
                        Some(result.naive_local().and_utc().timestamp_micros()),
                        None,
                    )))
                }
            }
            ColumnarValue::Array(t) if t.as_any().is::<TimestampMicrosecondArray>() => {
                let t = t
                    .as_any()
                    .downcast_ref::<TimestampMicrosecondArray>()
                    .unwrap();
                Ok(ColumnarValue::Array(Arc::new(convert_tz_array(
                    t, &from_tz, &to_tz,
                )?)))
            }
            _ => {
                return Err(DataFusionError::Execution(
                    "First argument in CONVERT_TZ_KSQL must be timestamp or array of timestamps"
                        .to_string(),
                ));
            }
        }
    }
}

struct FormatTimestampFunc {
    signature: Signature,
}

impl FormatTimestampFunc {
    fn new() -> FormatTimestampFunc {
        FormatTimestampFunc {
            signature: Signature::exact(
                vec![
                    DataType::Timestamp(TimeUnit::Microsecond, None),
                    DataType::Utf8,
                ],
                Volatility::Stable,
            ),
        }
    }
}

impl Debug for FormatTimestampFunc {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "FormatTimestampFunc")
    }
}

impl ScalarUDFImpl for FormatTimestampFunc {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn name(&self) -> &str {
        "format_timestamp"
    }

    fn signature(&self) -> &Signature {
        &self.signature
    }

    fn return_type(&self, _: &[DataType]) -> datafusion::common::Result<DataType> {
        Ok(DataType::Utf8)
    }

    fn invoke(&self, inputs: &[ColumnarValue]) -> datafusion::common::Result<ColumnarValue> {
        if inputs.len() != 2 {
            return Err(DataFusionError::Execution(
                "Expected 2 arguments in FORMAT_TIMESTAMP".to_string(),
            ));
        }

        let format = match &inputs[1] {
            ColumnarValue::Scalar(ScalarValue::Utf8(Some(v))) => sql_format_to_strformat(v),
            _ => {
                return Err(DataFusionError::Execution(
                    "Only scalar arguments are supported as format in FORMAT_TIMESTAMP".to_string(),
                ));
            }
        };

        match &inputs[0] {
            ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(Some(t), None)) => {
                let time = Utc.timestamp_nanos(*t * 1000).naive_local();
                Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(format!(
                    "{}",
                    time.format(&format)
                )))))
            }
            ColumnarValue::Array(t) if t.as_any().is::<TimestampMicrosecondArray>() => {
                let t = t
                    .as_any()
                    .downcast_ref::<TimestampMicrosecondArray>()
                    .unwrap();
                Ok(ColumnarValue::Array(Arc::new(format_timestamp_array(
                    &t, &format,
                )?)))
            }
            _ => {
                return Err(DataFusionError::Execution(
                    "First argument in FORMAT_TIMESTAMP must be timestamp or array of timestamps"
                        .to_string(),
                ));
            }
        }
    }
}
