use arrow::array::{
    Array, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
    TimestampSecondArray,
};
use arrow::datatypes::TimeUnit;
use pyo3::exceptions::PyTypeError;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList, PyModule};

pub(crate) fn append_timestamp_values<'py>(
    py: Python<'py>,
    out: &Bound<'py, PyList>,
    array: &dyn Array,
    unit: &TimeUnit,
    tz_aware: bool,
) -> PyResult<()> {
    let datetime = PyModule::import(py, "datetime")?;
    let datetime_cls = datetime.getattr("datetime")?;
    let timezone = datetime.getattr("timezone")?;
    let utc = timezone.getattr("utc")?;

    macro_rules! append_timestamp_array {
        ($arr_ty:ty) => {{
            let a = array
                .as_any()
                .downcast_ref::<$arr_ty>()
                .ok_or_else(|| PyTypeError::new_err(concat!("expected ", stringify!($arr_ty))))?;
            for i in 0..a.len() {
                if a.is_null(i) {
                    out.append(py.None())?;
                } else {
                    let value = a.value(i);
                    let (seconds, micros) = timestamp_parts(value, unit);
                    let tz_kwargs = PyDict::new(py);
                    if tz_aware {
                        tz_kwargs.set_item("tz", utc.clone())?;
                    }
                    let dt = if tz_aware {
                        datetime_cls.call_method("fromtimestamp", (seconds,), Some(&tz_kwargs))?
                    } else {
                        datetime_cls.call_method("utcfromtimestamp", (seconds,), None)?
                    };
                    if micros != 0 {
                        let replace_kwargs = PyDict::new(py);
                        replace_kwargs.set_item("microsecond", micros)?;
                        let dt = dt.call_method("replace", (), Some(&replace_kwargs))?;
                        out.append(dt)?;
                    } else {
                        out.append(dt)?;
                    }
                }
            }
        }};
    }

    match unit {
        TimeUnit::Second => append_timestamp_array!(TimestampSecondArray),
        TimeUnit::Millisecond => append_timestamp_array!(TimestampMillisecondArray),
        TimeUnit::Microsecond => append_timestamp_array!(TimestampMicrosecondArray),
        TimeUnit::Nanosecond => append_timestamp_array!(TimestampNanosecondArray),
    }

    Ok(())
}

fn timestamp_parts(value: i64, unit: &TimeUnit) -> (f64, u32) {
    match unit {
        TimeUnit::Second => (value as f64, 0),
        TimeUnit::Millisecond => (
            value.div_euclid(1_000) as f64 + value.rem_euclid(1_000) as f64 / 1_000.0,
            (value.rem_euclid(1_000) as u32) * 1_000,
        ),
        TimeUnit::Microsecond => (
            value.div_euclid(1_000_000) as f64 + value.rem_euclid(1_000_000) as f64 / 1_000_000.0,
            value.rem_euclid(1_000_000) as u32,
        ),
        TimeUnit::Nanosecond => (
            value.div_euclid(1_000_000_000) as f64
                + value.rem_euclid(1_000_000_000) as f64 / 1_000_000_000.0,
            (value.rem_euclid(1_000_000_000) / 1_000) as u32,
        ),
    }
}
