//! Parse timestamp UDF for DataFusion.
//!
//! Uses chrono to parse timestamps directly to microseconds, avoiding the
//! nanosecond limitation in DataFusion's built-in to_timestamp functions.

use std::any::Any;
use std::sync::Arc;

use datafusion::arrow::array::{Array, ArrayRef, AsArray, TimestampMicrosecondArray};
use datafusion::arrow::datatypes::{DataType, TimeUnit};
use datafusion::common::{exec_err, Result, ScalarValue};
use datafusion::logical_expr::{
    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature,
    Volatility,
};
use hamelin_lib::parse_timestamp_to_utc;

use super::string_utils::{scalar_to_str, STRING_TYPES};

#[derive(Debug, PartialEq, Eq, Hash)]
pub struct ParseTimestampUdf {
    signature: Signature,
}

impl Default for ParseTimestampUdf {
    fn default() -> Self {
        Self::new()
    }
}

impl ParseTimestampUdf {
    pub fn new() -> Self {
        let sigs: Vec<TypeSignature> = STRING_TYPES
            .iter()
            .map(|t| TypeSignature::Exact(vec![t.clone()]))
            .collect();
        Self {
            signature: Signature::new(TypeSignature::OneOf(sigs), Volatility::Immutable),
        }
    }
}

/// Parse a timestamp string using chrono, returning microseconds since epoch.
///
/// Delegates to `hamelin_lib::parse_timestamp_to_utc` so the format set stays
/// in lockstep with eval's native `string AS timestamp` cast.
pub fn parse_timestamp_micros(s: &str) -> Result<i64> {
    match parse_timestamp_to_utc(s) {
        Some(dt) => Ok(dt.timestamp_micros()),
        None => exec_err!("Failed to parse timestamp from string: {}", s),
    }
}

/// Parse timestamps from any string array type.
fn parse_string_array<T>(array: &T) -> ArrayRef
where
    T: Array + 'static,
    for<'a> &'a T: IntoIterator<Item = Option<&'a str>>,
{
    let result: TimestampMicrosecondArray = array
        .into_iter()
        .map(|opt_str| opt_str.and_then(|s| parse_timestamp_micros(s).ok()))
        .collect();
    Arc::new(result.with_timezone("+00:00"))
}

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

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

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

    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
        Ok(DataType::Timestamp(
            TimeUnit::Microsecond,
            Some("+00:00".into()),
        ))
    }

    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
        let args = args.args;
        if args.len() != 1 {
            return exec_err!(
                "hamelin_parse_timestamp expects exactly 1 argument, got {}",
                args.len()
            );
        }

        match &args[0] {
            ColumnarValue::Scalar(scalar) => {
                let micros = scalar_to_str(scalar)?.and_then(|s| parse_timestamp_micros(s).ok());
                Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
                    micros,
                    Some("+00:00".into()),
                )))
            }
            ColumnarValue::Array(array) => {
                let result = match array.data_type() {
                    DataType::Utf8 => parse_string_array(array.as_string::<i32>()),
                    DataType::LargeUtf8 => parse_string_array(array.as_string::<i64>()),
                    DataType::Utf8View => parse_string_array(array.as_string_view()),
                    other => {
                        return exec_err!(
                            "hamelin_parse_timestamp expects string array, got {}",
                            other
                        );
                    }
                };
                Ok(ColumnarValue::Array(result))
            }
        }
    }
}

/// Create the parse_timestamp UDF.
pub fn parse_timestamp_udf() -> ScalarUDF {
    ScalarUDF::from(ParseTimestampUdf::new())
}
