use arrow_schema::DataType;
use serde::{Deserialize, Serialize};
use std::{
    collections::HashMap,
    fmt,
    fmt::{Display, Formatter},
    sync::Arc,
};

#[derive(thiserror::Error, Debug, PartialEq)]
#[error("Wrong type for property {name}: expected {expected:?} but actual type is {actual:?}")]
pub struct PropError {
    pub(crate) name: String,
    pub(crate) expected: PropType,
    pub(crate) actual: PropType,
}

impl PropError {
    pub fn with_name(self, name: String) -> PropError {
        Self { name, ..self }
    }
}

#[derive(Clone, PartialEq, Eq, Debug, Default, Serialize, Deserialize)]
pub enum PropType {
    #[default]
    Empty,
    Str,
    U8,
    U16,
    I32,
    I64,
    U32,
    U64,
    F32,
    F64,
    Bool,
    List(Box<PropType>),
    Map(Arc<HashMap<String, PropType>>),
    NDTime,
    DTime,
    Decimal {
        scale: i64,
    },
}

impl Display for PropType {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let type_str = match self {
            PropType::Empty => "Empty",
            PropType::Str => "Str",
            PropType::U8 => "U8",
            PropType::U16 => "U16",
            PropType::I32 => "I32",
            PropType::I64 => "I64",
            PropType::U32 => "U32",
            PropType::U64 => "U64",
            PropType::F32 => "F32",
            PropType::F64 => "F64",
            PropType::Bool => "Bool",
            PropType::List(p_type) => return write!(f, "List<{}>", p_type),
            PropType::Map(p_type) => {
                let mut types = p_type
                    .iter()
                    .map(|(k, v)| format!("{}: {}", k, v))
                    .collect::<Vec<String>>();
                types.sort();
                return write!(f, "Map{{ {} }}", types.join(", "));
            }
            PropType::NDTime => "NDTime",
            PropType::DTime => "DTime",
            PropType::Decimal { scale } => return write!(f, "Decimal({})", scale),
        };

        write!(f, "{}", type_str)
    }
}

impl PropType {
    pub fn inner(&self) -> Option<&PropType> {
        match self {
            PropType::List(inner) => Some(inner.as_ref()),
            _ => None,
        }
    }

    pub fn map(fields: impl IntoIterator<Item = (impl Into<String>, PropType)>) -> Self {
        let map: HashMap<_, _> = fields.into_iter().map(|(k, v)| (k.into(), v)).collect();
        PropType::Map(Arc::from(map))
    }

    pub fn is_numeric(&self) -> bool {
        matches!(
            self,
            PropType::U8
                | PropType::U16
                | PropType::U32
                | PropType::U64
                | PropType::I32
                | PropType::I64
                | PropType::F32
                | PropType::F64
                | PropType::Decimal { .. }
        )
    }

    pub fn is_str(&self) -> bool {
        matches!(self, PropType::Str)
    }

    pub fn is_bool(&self) -> bool {
        matches!(self, PropType::Bool)
    }

    pub fn is_date(&self) -> bool {
        matches!(self, PropType::DTime | PropType::NDTime)
    }

    pub fn has_add(&self) -> bool {
        self.is_numeric() || self.is_str()
    }

    pub fn has_divide(&self) -> bool {
        self.is_numeric()
    }

    pub fn has_cmp(&self) -> bool {
        self.is_bool() || self.is_numeric() || self.is_str() || self.is_date()
    }

    pub fn homogeneous_map_value_type(&self) -> Option<PropType> {
        if let PropType::Map(map) = self {
            let mut iter = map.values();
            if let Some(first) = iter.next() {
                if iter.all(|v| v == first) {
                    return Some(first.clone());
                }
            }
        }
        None
    }

    // This is the best guess for the size of one row of properties
    pub fn est_size(&self) -> usize {
        const CONTAINER_SIZE: usize = 64;
        match self {
            PropType::Str => CONTAINER_SIZE,
            PropType::U8 | PropType::Bool => 1,
            PropType::U16 => 2,
            PropType::I32 | PropType::F32 | PropType::U32 => 4,
            PropType::I64 | PropType::F64 | PropType::U64 => 8,
            PropType::NDTime | PropType::DTime => 8,
            PropType::List(p_type) => p_type.est_size() * CONTAINER_SIZE,
            PropType::Map(p_map) => {
                p_map.values().map(|v| v.est_size()).sum::<usize>() * CONTAINER_SIZE
            }
            PropType::Decimal { .. } => 16,
            PropType::Empty => 0,
        }
    }
}

pub fn data_type_as_prop_type(dt: &DataType) -> Result<PropType, InvalidPropertyTypeErr> {
    match dt {
        DataType::Boolean => Ok(PropType::Bool),
        DataType::Int32 => Ok(PropType::I32),
        DataType::Int64 => Ok(PropType::I64),
        DataType::UInt8 => Ok(PropType::U8),
        DataType::UInt16 => Ok(PropType::U16),
        DataType::UInt32 => Ok(PropType::U32),
        DataType::UInt64 => Ok(PropType::U64),
        DataType::Float32 => Ok(PropType::F32),
        DataType::Float64 => Ok(PropType::F64),
        DataType::Utf8 => Ok(PropType::Str),
        DataType::LargeUtf8 => Ok(PropType::Str),
        DataType::Utf8View => Ok(PropType::Str),
        DataType::Struct(fields) => Ok(PropType::map(fields.iter().filter_map(|f| {
            data_type_as_prop_type(f.data_type())
                .ok()
                .map(move |pt| (f.name(), pt))
        }))),
        DataType::List(v) => Ok(PropType::List(Box::new(data_type_as_prop_type(
            v.data_type(),
        )?))),
        DataType::FixedSizeList(v, _) => Ok(PropType::List(Box::new(data_type_as_prop_type(
            v.data_type(),
        )?))),
        DataType::LargeList(v) => Ok(PropType::List(Box::new(data_type_as_prop_type(
            v.data_type(),
        )?))),
        DataType::Timestamp(_, v) => match v {
            None => Ok(PropType::NDTime),
            Some(_) => Ok(PropType::DTime),
        },
        DataType::Date32 => Ok(PropType::NDTime),
        DataType::Date64 => Ok(PropType::NDTime),
        DataType::Decimal128(precision, scale) if *precision <= 38 => Ok(PropType::Decimal {
            scale: *scale as i64,
        }),
        DataType::Null => Ok(PropType::Empty),
        _ => Err(InvalidPropertyTypeErr(dt.clone())),
    }
}

#[derive(thiserror::Error, Debug)]
#[error("{0:?} not supported as property type")]
pub struct InvalidPropertyTypeErr(pub DataType);

pub mod arrow {
    use crate::core::entities::properties::prop::{PropType, EMPTY_MAP_FIELD_NAME};
    use arrow_schema::{DataType, TimeUnit};

    impl From<&DataType> for PropType {
        fn from(value: &DataType) -> Self {
            match value {
                DataType::Utf8View | DataType::LargeUtf8 | DataType::Utf8 => PropType::Str,
                DataType::UInt8 => PropType::U8,
                DataType::UInt16 => PropType::U16,
                DataType::Int32 => PropType::I32,
                DataType::Int64 => PropType::I64,
                DataType::UInt32 => PropType::U32,
                DataType::UInt64 => PropType::U64,
                DataType::Float32 => PropType::F32,
                DataType::Float64 => PropType::F64,
                DataType::Decimal128(_, scale) => PropType::Decimal {
                    scale: *scale as i64,
                },
                DataType::Boolean => PropType::Bool,
                DataType::Timestamp(TimeUnit::Millisecond, None) => PropType::NDTime,
                DataType::Timestamp(TimeUnit::Millisecond, tz) if tz.as_deref() == Some("UTC") => {
                    PropType::DTime
                }
                DataType::Struct(fields) => PropType::map(
                    fields
                        .iter()
                        .filter(|field| field.name() != EMPTY_MAP_FIELD_NAME)
                        .map(|f| (f.name().to_string(), PropType::from(f.data_type()))),
                ),
                DataType::List(field) | DataType::LargeList(field) => {
                    PropType::List(Box::new(PropType::from(field.data_type())))
                }
                DataType::Null => PropType::Empty,
                dtype => panic!("unsupported type {dtype:?}"),
            }
        }
    }
}

// step through these types trees and check they are structurally the same
// if we encounter an empty we replace it with the other type
// the result is the unified type or err if the types are not compatible
pub fn unify_types(l: &PropType, r: &PropType, unified: &mut bool) -> Result<PropType, PropError> {
    match (l, r) {
        (PropType::Empty, r) => {
            *unified = true;
            Ok(r.clone())
        }
        (l, PropType::Empty) => {
            *unified = true;
            Ok(l.clone())
        }
        (PropType::Str, PropType::Str) => Ok(PropType::Str),
        (PropType::U8, PropType::U8) => Ok(PropType::U8),
        (PropType::U16, PropType::U16) => Ok(PropType::U16),
        (PropType::I32, PropType::I32) => Ok(PropType::I32),
        (PropType::I64, PropType::I64) => Ok(PropType::I64),
        (PropType::U32, PropType::U32) => Ok(PropType::U32),
        (PropType::U64, PropType::U64) => Ok(PropType::U64),
        (PropType::F32, PropType::F32) => Ok(PropType::F32),
        (PropType::F64, PropType::F64) => Ok(PropType::F64),
        (PropType::Bool, PropType::Bool) => Ok(PropType::Bool),
        (PropType::NDTime, PropType::NDTime) => Ok(PropType::NDTime),
        (PropType::DTime, PropType::DTime) => Ok(PropType::DTime),
        (PropType::List(l_type), PropType::List(r_type)) => {
            unify_types(l_type, r_type, unified).map(|t| PropType::List(Box::new(t)))
        }
        (PropType::Map(l_map), PropType::Map(r_map)) => {
            // maps need to be merged and only overlapping keys need to be unified

            let mut merged = HashMap::new();
            for (k, v) in l_map.iter() {
                if let Some(r_v) = r_map.get(k) {
                    let merged_prop = unify_types(v, r_v, unified)?;
                    merged.insert(k.clone(), merged_prop);
                } else {
                    merged.insert(k.clone(), v.clone());
                    *unified = true;
                }
            }
            for (k, v) in r_map.iter() {
                if !merged.contains_key(k) {
                    merged.insert(k.clone(), v.clone());
                    *unified = true;
                }
            }
            Ok(PropType::Map(merged.into()))
        }
        (PropType::Decimal { scale: l_scale }, PropType::Decimal { scale: r_scale })
            if l_scale == r_scale =>
        {
            Ok(PropType::Decimal { scale: *l_scale })
        }
        (_, _) => Err(PropError {
            name: "unknown".to_string(),
            expected: l.clone(),
            actual: r.clone(),
        }),
    }
}

// fast check before we actually unify, 99% of the time this will be enough
// there are 3 outcomes,
// types are identical so no unification needed, None
// types can be unified Some(true)
// and types cannot be unified Some(false)
pub fn check_for_unification(l: &PropType, r: &PropType) -> Option<bool> {
    match (l, r) {
        (PropType::Empty, _) => Some(true),
        (_, PropType::Empty) => Some(true),
        (PropType::Str, PropType::Str) => None,
        (PropType::U8, PropType::U8) => None,
        (PropType::U16, PropType::U16) => None,
        (PropType::I32, PropType::I32) => None,
        (PropType::I64, PropType::I64) => None,
        (PropType::U32, PropType::U32) => None,
        (PropType::U64, PropType::U64) => None,
        (PropType::F32, PropType::F32) => None,
        (PropType::F64, PropType::F64) => None,
        (PropType::Bool, PropType::Bool) => None,
        (PropType::NDTime, PropType::NDTime) => None,
        (PropType::DTime, PropType::DTime) => None,
        (PropType::List(l_type), PropType::List(r_type)) => check_for_unification(l_type, r_type),
        (PropType::Map(l_map), PropType::Map(r_map)) => {
            let keys_check = l_map
                .keys()
                .any(|k| !r_map.contains_key(k))
                .then_some(true)
                .or_else(|| r_map.keys().any(|k| !l_map.contains_key(k)).then_some(true));

            // check for unification of the values
            let inner_checks = l_map
                .iter()
                .filter_map(|(l_key, l_d_type)| {
                    r_map
                        .get(l_key)
                        .and_then(|r_d_type| check_for_unification(r_d_type, l_d_type))
                })
                .chain(r_map.iter().filter_map(|(r_key, r_d_type)| {
                    l_map
                        .get(r_key)
                        .and_then(|l_d_type| check_for_unification(r_d_type, l_d_type))
                }));
            for check in inner_checks {
                if check {
                    return Some(true);
                }
            }
            keys_check
        }
        (PropType::Decimal { scale: l_scale }, PropType::Decimal { scale: r_scale })
            if l_scale == r_scale =>
        {
            None
        }
        _ => Some(false),
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_unify_types_ne() {
        let l = PropType::List(Box::new(PropType::U8));
        let r = PropType::List(Box::new(PropType::U16));
        assert!(unify_types(&l, &r, &mut false).is_err());

        let l = PropType::map([("a".to_string(), PropType::U8)]);
        let r = PropType::map([("a".to_string(), PropType::U16)]);
        assert!(unify_types(&l, &r, &mut false).is_err());

        let l = PropType::List(Box::new(PropType::U8));
        let r = PropType::List(Box::new(PropType::U16));
        assert!(unify_types(&l, &r, &mut false).is_err());
    }

    #[test]
    fn test_unify_types_eq() {
        let l = PropType::List(Box::new(PropType::U8));
        let r = PropType::List(Box::new(PropType::U8));
        assert_eq!(
            unify_types(&l, &r, &mut false),
            Ok(PropType::List(Box::new(PropType::U8)))
        );

        let l = PropType::map([("a".to_string(), PropType::U8)]);
        let r = PropType::map([("a".to_string(), PropType::U8)]);
        assert_eq!(
            unify_types(&l, &r, &mut false),
            Ok(PropType::map([("a".to_string(), PropType::U8)]))
        );
    }

    #[test]
    fn test_unify_maps() {
        let l = PropType::map([("a".to_string(), PropType::U8)]);
        let r = PropType::map([("a".to_string(), PropType::U16)]);
        assert!(unify_types(&l, &r, &mut false).is_err());

        let l = PropType::map([("a".to_string(), PropType::U8)]);
        let r = PropType::map([("b".to_string(), PropType::U16)]);
        let mut unify = false;
        assert_eq!(
            unify_types(&l, &r, &mut unify),
            Ok(PropType::map([
                ("a".to_string(), PropType::U8),
                ("b".to_string(), PropType::U16)
            ]))
        );
        assert!(unify);

        let l = PropType::map([("a".to_string(), PropType::U8)]);
        let r = PropType::map([
            ("a".to_string(), PropType::U8),
            ("b".to_string(), PropType::U16),
        ]);
        let mut unify = false;
        assert_eq!(
            unify_types(&l, &r, &mut unify),
            Ok(PropType::map([
                ("a".to_string(), PropType::U8),
                ("b".to_string(), PropType::U16)
            ]))
        );
        assert!(unify);

        let l = PropType::map([
            ("a".to_string(), PropType::U8),
            ("b".to_string(), PropType::U16),
        ]);
        let r = PropType::map([("a".to_string(), PropType::U8)]);
        let mut unify = false;
        assert_eq!(
            unify_types(&l, &r, &mut unify),
            Ok(PropType::map([
                ("a".to_string(), PropType::U8),
                ("b".to_string(), PropType::U16)
            ]))
        );
        assert!(unify);
    }

    #[test]
    fn test_unify() {
        let l = PropType::Empty;
        let r = PropType::U8;
        let mut unify = false;
        assert_eq!(unify_types(&l, &r, &mut unify), Ok(PropType::U8));
        assert!(unify);

        let l = PropType::Str;
        let r = PropType::Empty;
        let mut unify = false;
        assert_eq!(unify_types(&l, &r, &mut unify), Ok(PropType::Str));
        assert!(unify);

        let l = PropType::List(Box::new(PropType::List(Box::new(PropType::U8))));
        let r = PropType::List(Box::new(PropType::Empty));
        let mut unify = false;
        assert_eq!(
            unify_types(&l, &r, &mut unify),
            Ok(PropType::List(Box::new(PropType::List(Box::new(
                PropType::U8
            )))))
        );
        assert!(unify);

        let l = PropType::List(Box::new(PropType::map([("a".to_string(), PropType::U8)])));
        let r = PropType::List(Box::new(PropType::map([
            ("a".to_string(), PropType::Empty),
            ("b".to_string(), PropType::Str),
        ])));
        let mut unify = false;
        assert_eq!(
            unify_types(&l, &r, &mut unify),
            Ok(PropType::List(Box::new(PropType::map([
                ("a".to_string(), PropType::U8),
                ("b".to_string(), PropType::Str)
            ]))))
        );
        assert!(unify);
    }

    #[test]
    fn size_of_proptype() {
        let size = size_of::<PropType>();
        println!("PropType = {size}");
        let size = size_of::<HashMap<String, PropType>>();
        println!("Map = {size}");
        let size = size_of::<PropError>();
        println!("PropError = {size}")
    }
}
