unitforge 0.4.0

use std::fmt::Display;
use std::cmp::Ordering;
use std::fmt;
use std::fmt::Formatter;
use std::ops::{Add, Sub, Mul, Div, Neg};
#[cfg(feature = "pyo3")]
use pyo3::{Bound, PyAny, prelude::*};

#[derive(Debug)]
pub enum QuantityOperationError {
    AddError,
    SubError,
    MulError,
    DivError,
    SqrtError,
    ComparisonError,
}

impl fmt::Display for QuantityOperationError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            QuantityOperationError::AddError => write!(f, "Addition operation failed"),
            QuantityOperationError::SubError => write!(f, "Subtraction operation failed"),
            QuantityOperationError::MulError => write!(f, "Multiplication operation failed"),
            QuantityOperationError::DivError => write!(f, "Division operation failed"),
            QuantityOperationError::SqrtError => write!(f, "Sqrt operation failed"),
            QuantityOperationError::ComparisonError => write!(f, "Comparison operation failed"),
        }
    }
}

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum Quantity {
    //Is used for runtime checked operations with quantities
    FloatQuantity(f64),
    // __QUANTITY_VARIANTS__
}

impl Quantity {
    pub fn to(&self, unit: Unit) -> Result<f64, String> {
        match (self, unit) {
            (Quantity::FloatQuantity(value), Unit::NoUnit) => Ok(*value),
            // __QUANTITY_TO_VARIANTS__
            _ => Err("Cannot use given pair of quantity and unit.".to_string())
        }
    }

    pub fn abs(&self) -> Quantity {
        match self {
            Quantity::FloatQuantity(value) => Quantity::FloatQuantity(value.abs()),
            // __QUANTITY_ABS_VARIANTS__
        }
    }

    pub fn is_nan(&self) -> bool {
        match self {
            Quantity::FloatQuantity(value) => value.is_nan(),
            // __QUANTITY_NAN_VARIANTS__
        }
    }
}


impl Neg for Quantity {
    type Output = Quantity;
    fn neg(self) -> Quantity {
        match self {
            Quantity::FloatQuantity(value) => Quantity::FloatQuantity(-value),
            // __QUANTITY_NEG_VARIANTS__
        }
    }
}


#[derive(Clone, Copy, PartialEq, Debug)]
pub enum Unit {
    //Is used for runtime checked operations with quantities
    NoUnit,
    // __UNIT_VARIANTS__
}

impl Display for Unit {
    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
        write!(f, "{}", self.get_name())
    }
}

impl Unit {
    pub fn to_quantity(&self, value: f64) -> Quantity {
        match self {
            Unit::NoUnit => Quantity::FloatQuantity(value),
            // __TO_QUANTITY_VARIANTS__
        }
    }

    pub fn get_name(&self) -> &str {
        match self {
            Unit::NoUnit => "No Unit",
            // __TO_UNIT_NAME_VARIANTS__
        }
    }
}

impl Mul for Quantity {
    type Output = Result<Quantity, QuantityOperationError>;
    fn mul(self, other: Quantity) -> Result<Quantity, QuantityOperationError> {
        fn try_multiply(lhs: &Quantity, rhs: &Quantity) -> Result<Quantity, QuantityOperationError> {
            use Quantity::*;
            match (lhs, rhs) {
                (FloatQuantity(v_lhs), FloatQuantity(v_rhs)) => Ok(FloatQuantity(v_lhs * v_rhs)),
                // __MUL_MATCHES__
                _ => Err(QuantityOperationError::MulError)
            }
        }
        match try_multiply(&self, &other) {
            Ok(result) => Ok(result),
            Err(_) => try_multiply(&other, &self)
        }
    }
}

impl Div for Quantity {
    type Output = Result<Quantity, QuantityOperationError>;
    fn div(self, other: Quantity) -> Result<Quantity, QuantityOperationError> {
        use Quantity::*;
        match (self, other) {
            (FloatQuantity(v_lhs), FloatQuantity(v_rhs)) => Ok(FloatQuantity(v_lhs / v_rhs)),
            // __DIV_MATCHES__
            _ => Err(QuantityOperationError::DivError)
        }
    }
}

impl Add for Quantity {
    type Output = Result<Quantity, QuantityOperationError>;
    fn add(self, other: Quantity) -> Result<Quantity, QuantityOperationError> {
        use Quantity::*;
        match (self, other) {
            (Quantity::FloatQuantity(v_lhs), Quantity::FloatQuantity(v_rhs)) => Ok(Quantity::FloatQuantity(v_lhs + v_rhs)),
            // __ADD_QUANTITY_MATCHES__
            _ => Err(QuantityOperationError::AddError)
        }
    }
}

impl Sub for Quantity {
    type Output = Result<Quantity, QuantityOperationError>;
    fn sub(self, other: Quantity) -> Result<Self, QuantityOperationError> {
        use Quantity::*;
        match (self, other) {
            (Quantity::FloatQuantity(v_lhs), Quantity::FloatQuantity(v_rhs)) => Ok(Quantity::FloatQuantity(v_lhs - v_rhs)),
            // __SUB_QUANTITY_MATCHES__
            _ => Err(QuantityOperationError::SubError)
        }
    }
}

impl Quantity {
    pub fn extract_float(&self) -> Result<f64, String> {
        match self {
            Quantity::FloatQuantity(v) => Ok(*v),
            _ => Err("Cannot extract float from Quantity enum".into()),
        }
    }

    // __BASE_QUANTITY_MATCHES__
    pub fn sqrt(&self) -> Result<Self, QuantityOperationError> {
        match self {
            Quantity::FloatQuantity(v) => Ok(Self::FloatQuantity(v.sqrt())),
            // __QUANTITY_SQRTS__
            _=> Err(QuantityOperationError::SqrtError)
        }
    }
}

impl PartialOrd for Quantity {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        use Quantity::*;
        match (self, other) {
            (FloatQuantity(lhs), FloatQuantity(rhs)) => lhs.partial_cmp(rhs),
            // __QUANTITY_COMPARISONS__
            _ => panic!("Cannot compare non matching quantities!")
        }
    }
}

// __QUANTITY_KIND_ID_IMPLS__

#[cfg(feature = "pyo3")]
fn extract_f64(v: &Bound<PyAny>) -> Option<f64> {
    if let Ok(inner) = v.extract::<f64>() {
        Some(inner)
    } else if let Ok(inner) = v.extract::<f32>() {
        Some(inner as f64)
    } else if let Ok(inner) = v.extract::<i32>() {
        Some(inner as f64)
    } else if let Ok(inner) = v.extract::<i64>() {
        Some(inner as f64)
    } else {
        None
    }
}

#[cfg(feature = "pyo3")]
impl Quantity {
    pub fn from_py_any(v: &Bound<PyAny>) -> Result<Self, String> {
        if let Some(inner) = extract_f64(v) {
            Ok(Quantity::FloatQuantity(inner))
        }
        // __EXTRACT_QUANTITY_MATCHES__
        else if let Ok(inner) = pyo3_capsule_api::extract_quantity_from_capsule(v) {
            Ok(inner)
        }
        else {
            Err("Cannot interpret given value as Quantity".to_string())
        }
    }

    pub fn to_pyobject(self, py: Python) -> PyResult<Py<PyAny>> {
        Ok(match self {
            Quantity::FloatQuantity(v) => v.into_pyobject(py).map(|obj| obj.into())?,
            // __TO_PYOBJECT_MATCHES__
        })
    }
}

impl Display for Quantity {
    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
        match self {
            Quantity::FloatQuantity(v) => write!(f, "{v}"),
            // __QUANTITY_FMT_MATCHES__
        }
    }
}

#[cfg(feature = "pyo3")]
impl Unit {

    pub fn from_py_any(v: &Bound<PyAny>) -> Result<Self, String> {
        // __EXTRACT_UNIT_MATCHES__
        else if let Ok(inner) = pyo3_capsule_api::extract_unit_from_capsule(v) {
            Ok(inner)
        }
        else {
            Err("Cannot interpret given value as Quantity".to_string())
        }
    }
}

#[cfg(feature = "pyo3")]
pub mod pyo3_capsule_api {
    use super::*;
    use pyo3::exceptions::PyValueError;
    use pyo3::ffi;
    use pyo3::types::{PyAnyMethods, PyCapsule, PyCapsuleMethods, PyList, PyModule};
    use pyo3::{Bound, PyErr, PyRef, PyResult, Python};
    use std::ffi::{c_int, CString};

    pub const UNITFORGE_API_VERSION: u64 = 1;
    pub const CAPSULE_NAME: &str = "_UNITFORGE_API";

    #[repr(C)]
    #[derive(Clone, Copy, Debug)]
    pub struct QuantityRepr {
        pub kind_id: u32,
        pub multiplier: f64,
        pub power: i32,
    }

    #[repr(C)]
    #[derive(Clone, Copy, Debug)]
    pub struct UnitRepr {
        pub kind_id: u32,
        pub variant_id: i32,
    }

    pub trait UnitKindId {
        const KIND_ID: u32;
    }

    pub trait UnitVariantId {
        fn to_variant_id(self) -> i32;
        fn from_variant_id(variant_id: i32) -> Option<Self>
        where
            Self: Sized;
    }

    // __UNIT_KIND_ID_IMPLS__
    // __UNIT_VARIANT_ID_IMPLS__

    pub fn quantity_to_repr(quantity: Quantity) -> Result<QuantityRepr, String> {
        match quantity {
            Quantity::FloatQuantity(value) => Ok(QuantityRepr {
                kind_id: 0,
                multiplier: value,
                power: 0,
            }),
            // __QUANTITY_TO_REPR_MATCHES__
        }
    }

    fn quantity_from_repr(repr: QuantityRepr) -> Result<Quantity, String> {
        match repr.kind_id {
            0 => Ok(Quantity::FloatQuantity(repr.multiplier * 10f64.powi(repr.power))),
            // __QUANTITY_FROM_REPR_MATCHES__
            _ => Err(format!("Unknown quantity kind id: {}", repr.kind_id)),
        }
    }

    pub fn unit_to_repr(unit: Unit) -> UnitRepr {
        match unit {
            Unit::NoUnit => UnitRepr {
                kind_id: 0,
                variant_id: 0,
            },
            // __UNIT_TO_REPR_MATCHES__
        }
    }

    fn unit_from_repr(repr: UnitRepr) -> Result<Unit, String> {
        match repr.kind_id {
            0 => Ok(Unit::NoUnit),
            // __UNIT_FROM_REPR_MATCHES__
            _ => Err(format!("Unknown unit kind id: {}", repr.kind_id)),
        }
    }

    pub fn extract_quantity_from_capsule(obj: &Bound<PyAny>) -> PyResult<Quantity> {
        let py = obj.py();
        let module = PyModule::import(py, "unitforge")?;
        let api = get_api(py, &module)?;
        let mut repr = QuantityRepr {
            kind_id: 0,
            multiplier: 0.0,
            power: 0,
        };
        // SAFETY: `obj` is a valid python object and `repr` points to writable memory.
        let rc = unsafe { (api.extract_quantity_repr)(obj.as_ptr(), &mut repr as *mut QuantityRepr) };
        if rc != 0 {
            return Err(PyErr::fetch(py));
        }
        quantity_from_repr(repr).map_err(PyValueError::new_err)
    }

    pub fn extract_unit_from_capsule(obj: &Bound<PyAny>) -> PyResult<Unit> {
        let py = obj.py();
        let module = PyModule::import(py, "unitforge")?;
        let api = get_api(py, &module)?;
        let mut repr = UnitRepr {
            kind_id: 0,
            variant_id: 0,
        };
        // SAFETY: `obj` is a valid python object and `repr` points to writable memory.
        let rc = unsafe { (api.extract_unit_repr)(obj.as_ptr(), &mut repr as *mut UnitRepr) };
        if rc != 0 {
            return Err(PyErr::fetch(py));
        }
        unit_from_repr(repr).map_err(PyValueError::new_err)
    }

    #[repr(C)]
    #[derive(Clone, Copy)]
    pub struct UnitforgeApiV1 {
        pub version: u64,
        pub extract_quantity_repr:
            unsafe extern "C" fn(obj: *mut ffi::PyObject, out: *mut QuantityRepr) -> c_int,
        pub quantity_repr_to_object:
            unsafe extern "C" fn(repr: QuantityRepr) -> *mut ffi::PyObject,
        pub extract_unit_repr:
            unsafe extern "C" fn(obj: *mut ffi::PyObject, out: *mut UnitRepr) -> c_int,
        pub unit_repr_to_object: unsafe extern "C" fn(repr: UnitRepr) -> *mut ffi::PyObject,
        pub extract_vector3_repr:
            unsafe extern "C" fn(obj: *mut ffi::PyObject, out: *mut QuantityRepr) -> c_int,
        pub vector3_repr_to_object:
            unsafe extern "C" fn(data: *const QuantityRepr) -> *mut ffi::PyObject,
        pub extract_matrix3_repr:
            unsafe extern "C" fn(obj: *mut ffi::PyObject, out: *mut QuantityRepr) -> c_int,
        pub matrix3_repr_to_object:
            unsafe extern "C" fn(data: *const QuantityRepr) -> *mut ffi::PyObject,
    }

    unsafe extern "C" fn extract_quantity_repr_impl(
        obj: *mut ffi::PyObject,
        out: *mut QuantityRepr,
    ) -> c_int {
        let py = Python::assume_attached();
        let result: PyResult<()> = (|| {
            if out.is_null() {
                return Err(PyValueError::new_err("output pointer for quantity repr is null"));
            }
            let obj_bound = Bound::from_borrowed_ptr(py, obj);
            let quantity = Quantity::from_py_any(&obj_bound).map_err(PyValueError::new_err)?;
            unsafe {
                *out = quantity_to_repr(quantity).map_err(PyValueError::new_err)?;
            }
            Ok(())
        })();

        match result {
            Ok(_) => 0,
            Err(err) => {
                err.restore(py);
                -1
            }
        }
    }

    unsafe extern "C" fn quantity_repr_to_object_impl(repr: QuantityRepr) -> *mut ffi::PyObject {
        let py = Python::assume_attached();
        let result: PyResult<Py<PyAny>> = (|| {
            quantity_from_repr(repr)
                .map_err(PyValueError::new_err)?
                .to_pyobject(py)
        })();
        match result {
            Ok(obj) => obj.into_ptr(),
            Err(err) => {
                err.restore(py);
                std::ptr::null_mut()
            }
        }
    }

    unsafe extern "C" fn extract_unit_repr_impl(
        obj: *mut ffi::PyObject,
        out: *mut UnitRepr,
    ) -> c_int {
        let py = Python::assume_attached();
        let result: PyResult<()> = (|| {
            if out.is_null() {
                return Err(PyValueError::new_err("output pointer for unit repr is null"));
            }
            let obj_bound = Bound::from_borrowed_ptr(py, obj);
            let unit = Unit::from_py_any(&obj_bound).map_err(PyValueError::new_err)?;
            unsafe {
                *out = unit_to_repr(unit);
            }
            Ok(())
        })();

        match result {
            Ok(_) => 0,
            Err(err) => {
                err.restore(py);
                -1
            }
        }
    }

    unsafe extern "C" fn unit_repr_to_object_impl(repr: UnitRepr) -> *mut ffi::PyObject {
        let py = Python::assume_attached();
        let result: PyResult<Py<PyAny>> = (|| {
            let unit = unit_from_repr(repr).map_err(PyValueError::new_err)?;
            match unit {
                Unit::NoUnit => {
                    Err(PyValueError::new_err("NoUnit cannot be converted into a Python object"))
                }
                // __UNIT_TO_PYOBJECT_MATCHES__
            }
        })();
        match result {
            Ok(obj) => obj.into_ptr(),
            Err(err) => {
                err.restore(py);
                std::ptr::null_mut()
            }
        }
    }

    unsafe extern "C" fn extract_vector3_repr_impl(
        obj: *mut ffi::PyObject,
        out: *mut QuantityRepr,
    ) -> c_int {
        let py = Python::assume_attached();
        let result: PyResult<()> = (|| {
            if out.is_null() {
                return Err(PyValueError::new_err("output pointer for vector repr is null"));
            }
            let obj_bound = Bound::from_borrowed_ptr(py, obj);
            let vec = obj_bound.extract::<Vector3Py>()?;
            for i in 0..3 {
                let repr = quantity_to_repr(vec.data[i]).map_err(PyValueError::new_err)?;
                unsafe {
                    *out.add(i) = repr;
                }
            }
            Ok(())
        })();
        match result {
            Ok(_) => 0,
            Err(err) => {
                err.restore(py);
                -1
            }
        }
    }

    unsafe extern "C" fn vector3_repr_to_object_impl(
        data: *const QuantityRepr,
    ) -> *mut ffi::PyObject {
        let py = Python::assume_attached();
        let result: PyResult<Py<PyAny>> = (|| {
            if data.is_null() {
                return Err(PyValueError::new_err("input pointer for vector repr is null"));
            }
            let mut quantity_data = [Quantity::FloatQuantity(0.0); 3];
            for i in 0..3 {
                let repr = unsafe { *data.add(i) };
                quantity_data[i] = quantity_from_repr(repr).map_err(PyValueError::new_err)?;
            }
            let vec = Vector3Py { data: quantity_data };
            Ok(vec.into_pyobject(py)?.into_any().unbind())
        })();

        match result {
            Ok(obj) => obj.into_ptr(),
            Err(err) => {
                err.restore(py);
                std::ptr::null_mut()
            }
        }
    }

    unsafe extern "C" fn extract_matrix3_repr_impl(
        obj: *mut ffi::PyObject,
        out: *mut QuantityRepr,
    ) -> c_int {
        let py = Python::assume_attached();
        let result: PyResult<()> = (|| {
            if out.is_null() {
                return Err(PyValueError::new_err("output pointer for matrix repr is null"));
            }
            let obj_bound = Bound::from_borrowed_ptr(py, obj);
            let mat = obj_bound.extract::<PyRef<Matrix3Py>>()?;
            for i in 0..3 {
                for j in 0..3 {
                    let repr =
                        quantity_to_repr(mat.data[i][j]).map_err(PyValueError::new_err)?;
                    unsafe {
                        *out.add(i * 3 + j) = repr;
                    }
                }
            }
            Ok(())
        })();
        match result {
            Ok(_) => 0,
            Err(err) => {
                err.restore(py);
                -1
            }
        }
    }

    unsafe extern "C" fn matrix3_repr_to_object_impl(
        data: *const QuantityRepr,
    ) -> *mut ffi::PyObject {
        let py = Python::assume_attached();
        let result: PyResult<Py<PyAny>> = (|| {
            if data.is_null() {
                return Err(PyValueError::new_err("input pointer for matrix repr is null"));
            }
            let mut quantity_data = [[Quantity::FloatQuantity(0.0); 3]; 3];
            for i in 0..3 {
                for j in 0..3 {
                    let repr = unsafe { *data.add(i * 3 + j) };
                    quantity_data[i][j] =
                        quantity_from_repr(repr).map_err(PyValueError::new_err)?;
                }
            }
            let mat = Matrix3Py { data: quantity_data };
            Ok(mat.into_pyobject(py)?.into_any().unbind())
        })();
        match result {
            Ok(obj) => obj.into_ptr(),
            Err(err) => {
                err.restore(py);
                std::ptr::null_mut()
            }
        }
    }

    pub fn register_unitforge_api(py: Python<'_>, m: &Bound<PyModule>) -> PyResult<()> {
        let api = UnitforgeApiV1 {
            version: UNITFORGE_API_VERSION,
            extract_quantity_repr: extract_quantity_repr_impl,
            quantity_repr_to_object: quantity_repr_to_object_impl,
            extract_unit_repr: extract_unit_repr_impl,
            unit_repr_to_object: unit_repr_to_object_impl,
            extract_vector3_repr: extract_vector3_repr_impl,
            vector3_repr_to_object: vector3_repr_to_object_impl,
            extract_matrix3_repr: extract_matrix3_repr_impl,
            matrix3_repr_to_object: matrix3_repr_to_object_impl,
        };

        let capsule = PyCapsule::new_with_destructor(
            py,
            api,
            Some(CString::new(CAPSULE_NAME).unwrap()),
            |_api, _ctx| {},
        )?;
        m.add(CAPSULE_NAME, capsule)?;
        Ok(())
    }

    pub fn get_api<'py>(py: Python<'py>, module: &Bound<'py, PyModule>) -> PyResult<&'py UnitforgeApiV1> {
        let capsule = module.getattr(CAPSULE_NAME)?.cast_into::<PyCapsule>()?;
        if !capsule.is_valid() {
            return Err(PyValueError::new_err(format!(
                "Invalid capsule in unitforge module at attribute '{CAPSULE_NAME}'"
            )));
        }
        let ptr = capsule.pointer();
        if ptr.is_null() {
            return Err(PyValueError::new_err(format!(
                "Null pointer in unitforge capsule '{CAPSULE_NAME}'"
            )));
        }
        unsafe { Ok(&*ptr.cast::<UnitforgeApiV1>()) }
    }
}