use chrono::{DateTime, NaiveDate, NaiveDateTime};
use once_cell::sync::Lazy;
use polars_arrow::export::arrow::array::PrimitiveArray;
use polars_core::prelude::*;
use polars_core::utils::arrow::types::NativeType;
use regex::Regex;

use super::patterns::{self, Pattern};
#[cfg(feature = "dtype-date")]
use crate::chunkedarray::date::naive_date_to_date;
use crate::chunkedarray::utf8::strptime;
use crate::prelude::utf8::strptime::StrpTimeState;

const DATETIME_DMY_PATTERN: &str = r#"(?x)
        ^
        ['"]?                        # optional quotes
        (?:\d{1,2})                  # day
        [-/]                         # separator
        (?P<month>[01]?\d{1})        # month
        [-/]                         # separator
        (?:\d{4,})                   # year
        (?:
            [T\ ]                    # separator
            (?:\d{2})                # hour
            :?                       # separator
            (?:\d{2})                # minute
            (?:
                :?                   # separator
                (?:\d{2})            # second
                (?:
                    \.(?:\d{1,9})    # subsecond
                )?
            )?
        )?
        ['"]?                        # optional quotes
        $
        "#;

static DATETIME_DMY_RE: Lazy<Regex> = Lazy::new(|| Regex::new(DATETIME_DMY_PATTERN).unwrap());
const DATETIME_YMD_PATTERN: &str = r#"(?x)
        ^
        ['"]?                      # optional quotes
        (?:\d{4,})                 # year
        [-/]                       # separator
        (?P<month>[01]?\d{1})      # month
        [-/]                       # separator
        (?:\d{1,2})                # day
        (?:
            [T\ ]                  # separator
            (?:\d{2})              # hour
            :?                     # separator
            (?:\d{2})              # minute
            (?:
                :?                 # separator
                (?:\d{2})          # seconds
                (?:
                    \.(?:\d{1,9})  # subsecond
                )?
            )?
        )?
        ['"]?                      # optional quotes
        $
        "#;
static DATETIME_YMD_RE: Lazy<Regex> = Lazy::new(|| Regex::new(DATETIME_YMD_PATTERN).unwrap());
const DATETIME_YMDZ_PATTERN: &str = r#"(?x)
        ^
        ['"]?                  # optional quotes
        (?:\d{4,})             # year
        [-/]                   # separator
        (?P<month>[01]?\d{1})  # month
        [-/]                   # separator
        (?:\d{1,2})            # year
        [T\ ]                  # separator
        (?:\d{2})              # hour
        :?                     # separator
        (?:\d{2})              # minute
        (?:
            :?                 # separator
            (?:\d{2})          # second
            (?:
                \.(?:\d{1,9})  # subsecond
            )?
        )?
        (?:
            # offset (e.g. +01:00)
            [+-](?:\d{2})
            :?
            (?:\d{2})
            # or Zulu suffix
            |Z
        )
        ['"]?                  # optional quotes
        $
        "#;
static DATETIME_YMDZ_RE: Lazy<Regex> = Lazy::new(|| Regex::new(DATETIME_YMDZ_PATTERN).unwrap());

impl Pattern {
    pub fn is_inferable(&self, val: &str) -> bool {
        match self {
            Pattern::DateDMY => true, // there are very few Date patterns, so it's cheaper
            Pattern::DateYMD => true, // to just try them
            Pattern::DatetimeDMY => match DATETIME_DMY_RE.captures(val) {
                Some(search) => (1..=12).contains(
                    &search
                        .name("month")
                        .unwrap()
                        .as_str()
                        .parse::<u8>()
                        .unwrap(),
                ),
                None => false,
            },
            Pattern::DatetimeYMD => match DATETIME_YMD_RE.captures(val) {
                Some(search) => (1..=12).contains(
                    &search
                        .name("month")
                        .unwrap()
                        .as_str()
                        .parse::<u8>()
                        .unwrap(),
                ),
                None => false,
            },
            Pattern::DatetimeYMDZ => match DATETIME_YMDZ_RE.captures(val) {
                Some(search) => (1..=12).contains(
                    &search
                        .name("month")
                        .unwrap()
                        .as_str()
                        .parse::<u8>()
                        .unwrap(),
                ),
                None => false,
            },
        }
    }
}

pub trait StrpTimeParser<T> {
    fn parse_bytes(&mut self, val: &[u8]) -> Option<T>;
}

#[cfg(feature = "dtype-datetime")]
impl StrpTimeParser<i64> for DatetimeInfer<i64> {
    fn parse_bytes(&mut self, val: &[u8]) -> Option<i64> {
        if self.fmt_len == 0 {
            self.fmt_len = strptime::fmt_len(self.latest_fmt.as_bytes())?;
        }
        unsafe {
            self.transform_bytes
                .parse(val, self.latest_fmt.as_bytes(), self.fmt_len)
                .map(datetime_to_timestamp_us)
                .or_else(|| {
                    // TODO! this will try all patterns.
                    // somehow we must early escape if value is invalid
                    for fmt in self.patterns {
                        if self.fmt_len == 0 {
                            self.fmt_len = strptime::fmt_len(fmt.as_bytes())?;
                        }
                        if let Some(parsed) = self
                            .transform_bytes
                            .parse(val, fmt.as_bytes(), self.fmt_len)
                            .map(datetime_to_timestamp_us)
                        {
                            self.latest_fmt = fmt;
                            return Some(parsed);
                        }
                    }
                    None
                })
        }
    }
}

#[cfg(feature = "dtype-date")]
impl StrpTimeParser<i32> for DatetimeInfer<i32> {
    fn parse_bytes(&mut self, val: &[u8]) -> Option<i32> {
        if self.fmt_len == 0 {
            self.fmt_len = strptime::fmt_len(self.latest_fmt.as_bytes())?;
        }
        unsafe {
            self.transform_bytes
                .parse(val, self.latest_fmt.as_bytes(), self.fmt_len)
                .map(|ndt| naive_date_to_date(ndt.date()))
                .or_else(|| {
                    // TODO! this will try all patterns.
                    // somehow we must early escape if value is invalid
                    for fmt in self.patterns {
                        if self.fmt_len == 0 {
                            self.fmt_len = strptime::fmt_len(fmt.as_bytes())?;
                        }
                        if let Some(parsed) = self
                            .transform_bytes
                            .parse(val, fmt.as_bytes(), self.fmt_len)
                            .map(|ndt| naive_date_to_date(ndt.date()))
                        {
                            self.latest_fmt = fmt;
                            return Some(parsed);
                        }
                    }
                    None
                })
        }
    }
}

#[derive(Clone)]
pub struct DatetimeInfer<T> {
    pub pattern: Pattern,
    patterns: &'static [&'static str],
    latest_fmt: &'static str,
    transform: fn(&str, &str) -> Option<T>,
    transform_bytes: StrpTimeState,
    fmt_len: u16,
    pub logical_type: DataType,
}

#[cfg(feature = "dtype-datetime")]
impl TryFrom<Pattern> for DatetimeInfer<i64> {
    type Error = PolarsError;

    fn try_from(value: Pattern) -> PolarsResult<Self> {
        match value {
            Pattern::DatetimeDMY => Ok(DatetimeInfer {
                pattern: Pattern::DatetimeDMY,
                patterns: patterns::DATETIME_D_M_Y,
                latest_fmt: patterns::DATETIME_D_M_Y[0],
                transform: transform_datetime_us,
                transform_bytes: StrpTimeState::default(),
                fmt_len: 0,
                logical_type: DataType::Datetime(TimeUnit::Microseconds, None),
            }),
            Pattern::DatetimeYMD => Ok(DatetimeInfer {
                pattern: Pattern::DatetimeYMD,
                patterns: patterns::DATETIME_Y_M_D,
                latest_fmt: patterns::DATETIME_Y_M_D[0],
                transform: transform_datetime_us,
                transform_bytes: StrpTimeState::default(),
                fmt_len: 0,
                logical_type: DataType::Datetime(TimeUnit::Microseconds, None),
            }),
            Pattern::DatetimeYMDZ => Ok(DatetimeInfer {
                pattern: Pattern::DatetimeYMDZ,
                patterns: patterns::DATETIME_Y_M_D_Z,
                latest_fmt: patterns::DATETIME_Y_M_D_Z[0],
                transform: transform_tzaware_datetime_us,
                transform_bytes: StrpTimeState::default(),
                fmt_len: 0,
                logical_type: DataType::Datetime(TimeUnit::Microseconds, None),
            }),
            _ => polars_bail!(ComputeError: "could not convert pattern"),
        }
    }
}

#[cfg(feature = "dtype-date")]
impl TryFrom<Pattern> for DatetimeInfer<i32> {
    type Error = PolarsError;

    fn try_from(value: Pattern) -> PolarsResult<Self> {
        match value {
            Pattern::DateDMY => Ok(DatetimeInfer {
                pattern: Pattern::DateDMY,
                patterns: patterns::DATE_D_M_Y,
                latest_fmt: patterns::DATE_D_M_Y[0],
                transform: transform_date,
                transform_bytes: StrpTimeState::default(),
                fmt_len: 0,
                logical_type: DataType::Date,
            }),
            Pattern::DateYMD => Ok(DatetimeInfer {
                pattern: Pattern::DateYMD,
                patterns: patterns::DATE_Y_M_D,
                latest_fmt: patterns::DATE_Y_M_D[0],
                transform: transform_date,
                transform_bytes: StrpTimeState::default(),
                fmt_len: 0,
                logical_type: DataType::Date,
            }),
            _ => polars_bail!(ComputeError: "could not convert pattern"),
        }
    }
}

impl<T: NativeType> DatetimeInfer<T> {
    pub fn parse(&mut self, val: &str) -> Option<T> {
        match (self.transform)(val, self.latest_fmt) {
            Some(parsed) => Some(parsed),
            // try other patterns
            None => {
                if !self.pattern.is_inferable(val) {
                    return None;
                }
                for fmt in self.patterns {
                    self.fmt_len = 0;
                    if let Some(parsed) = (self.transform)(val, fmt) {
                        self.latest_fmt = fmt;
                        return Some(parsed);
                    }
                }
                None
            }
        }
    }

    fn coerce_utf8(&mut self, ca: &Utf8Chunked) -> Series {
        let chunks = ca
            .downcast_iter()
            .map(|array| {
                let iter = array
                    .into_iter()
                    .map(|opt_val| opt_val.and_then(|val| self.parse(val)));
                Box::new(PrimitiveArray::from_trusted_len_iter(iter)) as ArrayRef
            })
            .collect();
        let mut out = match self.logical_type {
            DataType::Date => unsafe { Int32Chunked::from_chunks(ca.name(), chunks) }
                .into_series()
                .cast(&self.logical_type)
                .unwrap(),
            DataType::Datetime(_, _) => unsafe { Int64Chunked::from_chunks(ca.name(), chunks) }
                .into_series()
                .cast(&self.logical_type)
                .unwrap(),
            _ => unreachable!(),
        };
        out.rename(ca.name());
        out
    }
}

#[cfg(feature = "dtype-date")]
fn transform_date(val: &str, fmt: &str) -> Option<i32> {
    NaiveDate::parse_from_str(val, fmt)
        .ok()
        .map(naive_date_to_date)
}

#[cfg(feature = "dtype-datetime")]
pub(crate) fn transform_datetime_ns(val: &str, fmt: &str) -> Option<i64> {
    let out = NaiveDateTime::parse_from_str(val, fmt)
        .ok()
        .map(datetime_to_timestamp_ns);
    out.or_else(|| {
        NaiveDate::parse_from_str(val, fmt)
            .ok()
            .map(|nd| datetime_to_timestamp_ns(nd.and_hms_opt(0, 0, 0).unwrap()))
    })
}

fn transform_tzaware_datetime_ns(val: &str, fmt: &str) -> Option<i64> {
    let dt = DateTime::parse_from_str(val, fmt);
    dt.ok().map(|dt| datetime_to_timestamp_ns(dt.naive_utc()))
}

#[cfg(feature = "dtype-datetime")]
pub(crate) fn transform_datetime_us(val: &str, fmt: &str) -> Option<i64> {
    let out = NaiveDateTime::parse_from_str(val, fmt)
        .ok()
        .map(datetime_to_timestamp_us);
    out.or_else(|| {
        NaiveDate::parse_from_str(val, fmt)
            .ok()
            .map(|nd| datetime_to_timestamp_us(nd.and_hms_opt(0, 0, 0).unwrap()))
    })
}

fn transform_tzaware_datetime_us(val: &str, fmt: &str) -> Option<i64> {
    let dt = DateTime::parse_from_str(val, fmt);
    dt.ok().map(|dt| datetime_to_timestamp_us(dt.naive_utc()))
}

#[cfg(feature = "dtype-datetime")]
pub(crate) fn transform_datetime_ms(val: &str, fmt: &str) -> Option<i64> {
    let out = NaiveDateTime::parse_from_str(val, fmt)
        .ok()
        .map(datetime_to_timestamp_ms);
    out.or_else(|| {
        NaiveDate::parse_from_str(val, fmt)
            .ok()
            .map(|nd| datetime_to_timestamp_ms(nd.and_hms_opt(0, 0, 0).unwrap()))
    })
}

fn transform_tzaware_datetime_ms(val: &str, fmt: &str) -> Option<i64> {
    let dt = DateTime::parse_from_str(val, fmt);
    dt.ok().map(|dt| datetime_to_timestamp_ms(dt.naive_utc()))
}

pub fn infer_pattern_single(val: &str) -> Option<Pattern> {
    // Dates come first, because we see datetimes as superset of dates
    infer_pattern_date_single(val).or_else(|| infer_pattern_datetime_single(val))
}

fn infer_pattern_datetime_single(val: &str) -> Option<Pattern> {
    if patterns::DATETIME_D_M_Y.iter().any(|fmt| {
        NaiveDateTime::parse_from_str(val, fmt).is_ok()
            || NaiveDate::parse_from_str(val, fmt).is_ok()
    }) {
        Some(Pattern::DatetimeDMY)
    } else if patterns::DATETIME_Y_M_D.iter().any(|fmt| {
        NaiveDateTime::parse_from_str(val, fmt).is_ok()
            || NaiveDate::parse_from_str(val, fmt).is_ok()
    }) {
        Some(Pattern::DatetimeYMD)
    } else if patterns::DATETIME_Y_M_D_Z
        .iter()
        .any(|fmt| NaiveDateTime::parse_from_str(val, fmt).is_ok())
    {
        Some(Pattern::DatetimeYMDZ)
    } else {
        None
    }
}

fn infer_pattern_date_single(val: &str) -> Option<Pattern> {
    if patterns::DATE_D_M_Y
        .iter()
        .any(|fmt| NaiveDate::parse_from_str(val, fmt).is_ok())
    {
        Some(Pattern::DateDMY)
    } else if patterns::DATE_Y_M_D
        .iter()
        .any(|fmt| NaiveDate::parse_from_str(val, fmt).is_ok())
    {
        Some(Pattern::DateYMD)
    } else {
        None
    }
}

#[cfg(feature = "dtype-datetime")]
pub(crate) fn to_datetime(
    ca: &Utf8Chunked,
    tu: TimeUnit,
    tz: Option<&TimeZone>,
) -> PolarsResult<DatetimeChunked> {
    match ca.first_non_null() {
        None => Ok(Int64Chunked::full_null(ca.name(), ca.len()).into_datetime(tu, tz.cloned())),
        Some(idx) => {
            let subset = ca.slice(idx as i64, ca.len());
            let pattern = subset
                .into_iter()
                .find_map(|opt_val| opt_val.and_then(infer_pattern_datetime_single))
                .ok_or_else(|| polars_err!(parse_fmt_idk = "date"))?;
            let mut infer = DatetimeInfer::<i64>::try_from(pattern)?;
            match (tu, pattern) {
                (TimeUnit::Nanoseconds, Pattern::DatetimeYMDZ) => {
                    infer.transform = transform_tzaware_datetime_ns
                }
                (TimeUnit::Microseconds, Pattern::DatetimeYMDZ) => {
                    infer.transform = transform_tzaware_datetime_us
                }
                (TimeUnit::Milliseconds, Pattern::DatetimeYMDZ) => {
                    infer.transform = transform_tzaware_datetime_ms
                }
                (TimeUnit::Nanoseconds, _) => infer.transform = transform_datetime_ns,
                (TimeUnit::Microseconds, _) => infer.transform = transform_datetime_us,
                (TimeUnit::Milliseconds, _) => infer.transform = transform_datetime_ms,
            }
            if tz.is_some() && pattern == Pattern::DatetimeYMDZ {
                polars_bail!(ComputeError: "cannot parse tz-aware values with tz-aware dtype - please drop the time zone from the dtype.")
            }
            match pattern {
                #[cfg(feature = "timezones")]
                Pattern::DatetimeYMDZ => infer.coerce_utf8(ca).datetime().map(|ca| {
                    let mut ca = ca.clone();
                    ca.set_time_unit(tu);
                    ca.replace_time_zone(Some("UTC"), None)
                })?,
                _ => infer.coerce_utf8(ca).datetime().map(|ca| {
                    let mut ca = ca.clone();
                    ca.set_time_unit(tu);
                    match tz {
                        #[cfg(feature = "timezones")]
                        Some(tz) => ca.replace_time_zone(Some(tz), None),
                        _ => Ok(ca),
                    }
                })?,
            }
        }
    }
}
#[cfg(feature = "dtype-date")]
pub(crate) fn to_date(ca: &Utf8Chunked) -> PolarsResult<DateChunked> {
    match ca.first_non_null() {
        None => Ok(Int32Chunked::full_null(ca.name(), ca.len()).into_date()),
        Some(idx) => {
            let subset = ca.slice(idx as i64, ca.len());
            let pattern = subset
                .into_iter()
                .find_map(|opt_val| opt_val.and_then(infer_pattern_date_single))
                .ok_or_else(|| polars_err!(parse_fmt_idk = "date"))?;
            let mut infer = DatetimeInfer::<i32>::try_from(pattern).unwrap();
            infer.coerce_utf8(ca).date().cloned()
        }
    }
}
