mod cell;
pub mod dtype;
mod field;
mod scalar;
mod value;

use cell::GILOnceCell;

pub use dtype::*;
pub use field::Field;
pub use scalar::Scalar;
pub use value::AnyValue;

use pyo3::{
    Borrowed, Bound, IntoPyObjectExt, Py, PyAny, PyResult, Python,
    exceptions::{PyOverflowError, PyValueError},
    types::{
        PyAnyMethods, PyBool, PyDict, PyDictMethods, PyFloat, PyInt, PyList, PySequence, PyString,
        PyTuple, PyType, PyTypeMethods,
    },
};
use std::{
    borrow::{Borrow, Cow},
    collections::HashMap,
};

type InitFn = for<'py> fn(&Bound<'py, PyAny>, bool) -> PyResult<AnyValue<'py>>;

pub(crate) static LUT: GILOnceCell<HashMap<TypeObjectKey, InitFn>> = GILOnceCell::new();

pub fn any_value_into_py_object_ref<'py, 'a>(
    av: &'a AnyValue<'a>,
    py: Python<'py>,
) -> PyResult<Bound<'py, PyAny>> {
    use AnyValue::*;
    match av {
        Null => py.None().into_bound_py_any(py),
        Bool(v) => v.into_bound_py_any(py),
        Char(v) => v.into_bound_py_any(py),
        UInt8(v) => v.into_bound_py_any(py),
        UInt16(v) => v.into_bound_py_any(py),
        UInt32(v) => v.into_bound_py_any(py),
        UInt64(v) => v.into_bound_py_any(py),
        UInt128(v) => v.into_bound_py_any(py),
        Int8(v) => v.into_bound_py_any(py),
        Int16(v) => v.into_bound_py_any(py),
        Int32(v) => v.into_bound_py_any(py),
        Int64(v) => v.into_bound_py_any(py),
        Int128(v) => v.into_bound_py_any(py),
        Float32(v) => v.into_bound_py_any(py),
        Float64(v) => v.into_bound_py_any(py),
        Str(s) => s.into_bound_py_any(py),
        StrOwned(s) => s.into_bound_py_any(py),
        Vector(v) => Ok(PyList::new(
            py,
            v.iter()
                .map(|item| any_value_into_py_object_ref(item, py))
                .collect::<PyResult<Vec<_>>>()?,
        )?
        .into_any()),
        Struct(pairs) => {
            let dict = pyo3::types::PyDict::new(py);
            for (fld, val) in pairs.iter() {
                let key = fld.name().to_string();
                let value = any_value_into_py_object_ref(val, py)?;
                dict.set_item(key, value)?;
            }
            Ok(dict.into_any())
        }
    }
}

pub fn any_value_into_py_object<'py>(av: AnyValue, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
    match av {
        AnyValue::UInt8(v) => v.into_bound_py_any(py),
        AnyValue::UInt16(v) => v.into_bound_py_any(py),
        AnyValue::UInt32(v) => v.into_bound_py_any(py),
        AnyValue::UInt64(v) => v.into_bound_py_any(py),
        AnyValue::UInt128(v) => v.into_bound_py_any(py),
        AnyValue::Int8(v) => v.into_bound_py_any(py),
        AnyValue::Int16(v) => v.into_bound_py_any(py),
        AnyValue::Int32(v) => v.into_bound_py_any(py),
        AnyValue::Int64(v) => v.into_bound_py_any(py),
        AnyValue::Int128(v) => v.into_bound_py_any(py),
        AnyValue::Float32(v) => v.into_bound_py_any(py),
        AnyValue::Float64(v) => v.into_bound_py_any(py),
        AnyValue::Char(v) => v.into_bound_py_any(py),
        AnyValue::Vector(v) => Ok(PyList::new(
            py,
            v.into_iter()
                .map(|item| any_value_into_py_object(item, py))
                .collect::<PyResult<Vec<_>>>()?,
        )?
        .into_any()),
        AnyValue::Null => py.None().into_bound_py_any(py),
        AnyValue::Bool(v) => v.into_bound_py_any(py),
        AnyValue::Str(v) => v.into_bound_py_any(py),
        AnyValue::StrOwned(v) => v.into_bound_py_any(py),
        AnyValue::Struct(v) => {
            let dict = struct_dict(py, v.into_iter())?;
            dict.into_bound_py_any(py)
        }
    }
}

pub fn py_object_to_any_value<'a, 'py>(
    ob: Borrowed<'a, 'py, PyAny>,
    strict: bool,
) -> PyResult<AnyValue<'py>> {
    fn get_null(_ob: &Bound<'_, PyAny>, _strict: bool) -> PyResult<AnyValue<'static>> {
        Ok(AnyValue::Null)
    }

    fn get_bool(ob: &Bound<'_, PyAny>, _strict: bool) -> PyResult<AnyValue<'static>> {
        let b = ob.extract::<bool>()?;
        Ok(AnyValue::Bool(b))
    }

    fn get_int(ob: &Bound<'_, PyAny>, strict: bool) -> PyResult<AnyValue<'static>> {
        if let Ok(v) = ob.extract::<i64>() {
            Ok(AnyValue::Int64(v))
        } else if let Ok(v) = ob.extract::<i128>() {
            Ok(AnyValue::Int128(v))
        } else if !strict {
            let f = ob.extract::<f64>()?;
            Ok(AnyValue::Float64(f))
        } else {
            Err(PyOverflowError::new_err(format!(
                "int value too large for Polars integer types: {ob}"
            )))
        }
    }

    fn get_float(ob: &Bound<'_, PyAny>, _strict: bool) -> PyResult<AnyValue<'static>> {
        Ok(AnyValue::Float64(ob.extract::<f64>()?))
    }

    fn get_str(ob: &Bound<'_, PyAny>, _strict: bool) -> PyResult<AnyValue<'static>> {
        Ok(AnyValue::StrOwned(ob.extract::<String>()?.into()))
    }

    fn get_list(ob: &Bound<'_, PyAny>, strict: bool) -> PyResult<AnyValue<'static>> {
        let seq = ob.cast::<PySequence>()?;
        let mut out: Vec<AnyValue<'static>> = Vec::with_capacity(seq.len().unwrap_or(0).max(0));
        for item in seq.try_iter()? {
            let av = py_object_to_any_value(item?.as_borrowed(), strict)?;
            out.push(av.into_static());
        }
        Ok(AnyValue::Vector(Box::new(out)))
    }

    fn get_conversion_function(ob: &Bound<'_, PyAny>, strict: bool) -> PyResult<InitFn> {
        if ob.is_none() {
            Ok(get_null)
        } else if ob.is_instance_of::<PyBool>() {
            Ok(get_bool)
        } else if ob.is_instance_of::<PyInt>() {
            Ok(get_int)
        } else if ob.is_instance_of::<PyFloat>() {
            Ok(get_float)
        } else if ob.is_instance_of::<PyString>() {
            Ok(get_str)
        } else if ob.is_instance_of::<PyList>() || ob.is_instance_of::<PyTuple>() {
            Ok(get_list)
        } else if ob.is_instance_of::<PyDict>() {
            Ok(|ob, strict| {
                let dict = ob.cast::<PyDict>().unwrap();
                let len = dict.len();
                let mut key_value_pairs = Vec::with_capacity(len);
                for (k, v) in dict.into_iter() {
                    let key = k.extract::<Cow<str>>()?;
                    let val = py_object_to_any_value(v.as_borrowed(), strict)?;
                    key_value_pairs.push((Field::from((key.into_owned(), val.dtype())), val));
                }

                Ok(AnyValue::Struct(key_value_pairs))
            })
        } else {
            let ob_type = ob.get_type();
            let type_name = ob_type.qualname()?.to_string();
            match type_name.as_str() {
                "range" => Ok(get_list as InitFn),
                _ => {
                    if !strict {
                        Ok(get_null)
                    } else {
                        Err(PyValueError::new_err(format!(
                            "Cannot convert Python object of type {type_name} to AnyValue"
                        )))
                    }
                }
            }
        }
    }

    let py_type = ob.get_type();
    let py_type_address = py_type.as_ptr() as usize;

    let py = ob.py();

    if !LUT.is_initialized() {
        LUT.set(py, Default::default()).unwrap();
    }

    LUT.with_gil(py, |lut| {
        if !lut.contains_key(&py_type_address) {
            let k = TypeObjectKey::new(py_type.clone().unbind());

            assert_eq!(k.address, py_type_address);

            lut.insert(k, get_conversion_function(ob.as_any(), strict)?);
        }

        let conversion_func = lut.get(&py_type_address).unwrap();
        conversion_func(ob.as_any(), strict)
    })
}

fn struct_dict<'py, 'a>(
    py: Python<'py>,
    vals: impl Iterator<Item = (Field, AnyValue<'a>)>,
) -> PyResult<Bound<'py, PyDict>> {
    let dict = PyDict::new(py);

    for (fld, val) in vals {
        let key = fld.name().to_string();
        let value = any_value_into_py_object(val, py)?;
        dict.set_item(key, value)?;
    }

    Ok(dict)
}

/// Holds a Python type object and implements hashing / equality based on the pointer address of the
/// type object. This is used as a hashtable key instead of only the `usize` pointer value, as we
/// need to hold a ref to the Python type object to keep it alive.
#[derive(Debug)]
pub struct TypeObjectKey {
    #[allow(unused)]
    type_object: Py<PyType>,
    /// We need to store this in a field for `Borrow<usize>`
    address: usize,
}

impl TypeObjectKey {
    fn new(type_object: Py<PyType>) -> Self {
        let address = type_object.as_ptr() as usize;
        Self {
            type_object,
            address,
        }
    }
}

impl PartialEq for TypeObjectKey {
    fn eq(&self, other: &Self) -> bool {
        self.address == other.address
    }
}

impl Eq for TypeObjectKey {}

impl std::borrow::Borrow<usize> for TypeObjectKey {
    fn borrow(&self) -> &usize {
        &self.address
    }
}

impl std::hash::Hash for TypeObjectKey {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        let v: &usize = self.borrow();
        v.hash(state)
    }
}
