genius-core-client 0.4.0

use crate::auth::retrieve_auth_token_client_credentials as retrieve_auth_token_client_credentials_rs;
use crate::client::inference::{
    clear_observations as clear_observations_rs, get_probability as get_probability_rs,
    ObservationValue,
};
use crate::client::{Client, TimeoutAndRetries};
use crate::types::entity::HSMLEntity;
use crate::types::static_schema::entity_schema::ENTITY_SCHEMA_SWID;
use crate::types::static_schema::link_schema::LINK_SCHEMA_SWID;
use crate::utils;
use once_cell::sync::Lazy;
use pyo3::prelude::*;
use pyo3::types::IntoPyDict;
use pyo3::types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyLong, PyString};
use pyo3::wrap_pyfunction;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;

// Lazy initialize client with null value to be set later
// because class methods cannot be async which is
// a limitation of Pyo3 / Python interop due to lifetime
// of async functions being potentially longer lived than
// the object we're defining methods on
static mut CLIENT: Lazy<Option<Arc<Mutex<Client>>>> = Lazy::new(|| None);

#[pymodule]
fn genius_core_client(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(new_with_oauth2_token, m)?)?;
    m.add_function(wrap_pyfunction!(make_swid, m)?)?;
    m.add_class::<PyClient>()?;
    m.add_class::<PyHSMLEntity>()?;

    let auth_module = PyModule::new(_py, "auth")?;
    let auth_utils_module = PyModule::new(_py, "utils")?;
    auth_utils_module.add_function(wrap_pyfunction!(retrieve_auth_token_client_credentials, m)?)?;
    auth_module.add_submodule(auth_utils_module)?;
    m.add_submodule(auth_module)?;
    Ok(())
}

#[pymodule]
fn auth_utils(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(retrieve_auth_token_client_credentials, m)?)?;
    Ok(())
}

#[pyfunction]
pub fn make_swid(class: String) -> String {
    utils::make_swid(&class)
}

#[pyfunction]
pub fn new_with_oauth2_token(
    py: Python,
    protocol: String,
    host: String,
    port: String,
    token: String,
    timeout: Option<u64>,
    retries: Option<u32>,
) -> PyResult<&PyAny> {
    pyo3_asyncio::tokio::future_into_py(py, async move {
        let timeout_and_retries = TimeoutAndRetries {
            timeout: tokio::time::Duration::from_secs(timeout.unwrap_or(30)),
            retries: retries.unwrap_or(3),
        };
        let result = Client::new_with_oauth2_token(
            crate::client::Protocol::from(protocol.as_str()),
            host,
            port,
            token,
            Some(timeout_and_retries),
        )
        .await;
        match result {
            Ok(client) => {
                // Initialize client because class methods cannot be async
                // which is a limitation of Pyo3 / Python interop due to lifetime
                // of async functions being potentially longer lived than
                // the object we're defining methods on
                unsafe {
                    *CLIENT = Some(Arc::new(Mutex::new(client)));
                }
                // Pass back to give a convenient and familiar OOP-like interface
                // that people are familiar with and that way
                // there is no confusion when using static functions
                // whether the client has been initialized because the very fact
                // of the object being returned generally means successful
                // client creation
                Ok(PyClient {
                    inner: unsafe { CLIENT.as_ref().unwrap().clone() },
                    inference: PyInference {},
                })
            }
            Err(err) => Err(PyErr::new::<pyo3::exceptions::PyException, _>(format!(
                "{}",
                err
            ))),
        }
    })
}

#[pyclass]
#[derive(Clone)]
pub struct PyInference {}

#[pymethods]
impl PyInference {
    pub fn get_probability(
        &self,
        py: Python,
        variables: Vec<String>,
        evidence: Option<&PyDict>,
    ) -> PyResult<PyObject> {
        let evidence_map: Option<HashMap<String, ObservationValue>> = evidence.map(|dict| {
            dict.into_iter()
                .map(|(key, val)| {
                    let key: String = key.extract().unwrap();
                    let val_dict: &PyDict = val.extract().expect("Failed to extract PyDict");
                    let element = val_dict
                        .get_item("element")
                        .unwrap()
                        .unwrap()
                        .extract::<String>()
                        .ok();
                    let distribution = val_dict
                        .get_item("distribution")
                        .unwrap()
                        .unwrap()
                        .extract::<Vec<f64>>()
                        .ok();
                    let none = val_dict
                        .get_item("none")
                        .unwrap()
                        .unwrap()
                        .extract::<bool>()
                        .ok();

                    let val = match (element, distribution, none) {
                        (Some(e), _, _) => ObservationValue::Element(e),
                        (_, Some(d), _) => ObservationValue::Distribution(d),
                        (_, _, Some(_)) => ObservationValue::None,
                        _ => panic!("Invalid type"),
                    };

                    (key, val)
                })
                .collect()
        });

        let result = pyo3_asyncio::tokio::future_into_py(py, async move {
            let mut client = unsafe { CLIENT.as_ref().unwrap().lock().await };
            let result = get_probability_rs(&mut client, variables, evidence_map).await;
            Python::with_gil(|py| {
                match result {
                    Ok(result) => {
                        // Convert HashMap<String, Vec<f64>> into pyo3 Python object
                        let dict = PyDict::new(py);
                        for (key, val) in result {
                            let py_list = PyList::new(py, &val);
                            dict.set_item(key, py_list).unwrap();
                        }
                        Ok(dict.to_object(py))
                    }
                    Err(err) => Err(PyErr::new::<pyo3::exceptions::PyException, _>(format!(
                        "{:#?}",
                        err
                    ))),
                }
            })
        });
        Ok(result.unwrap().to_object(py))
    }

    pub fn clear_observations(
        &self,
        py: Python,
        variables: Option<Vec<String>>,
    ) -> PyResult<PyObject> {
        let result = pyo3_asyncio::tokio::future_into_py(py, async move {
            let mut client = unsafe { CLIENT.as_ref().unwrap().lock().await };
            match clear_observations_rs(&mut client, variables).await {
                Ok(result) => {
                    // Convert the result into strings here, inside the async block
                    let result: Vec<String> =
                        result.into_iter().map(|value| value.to_string()).collect();
                    Ok(result)
                }
                Err(err) => Err(PyErr::new::<pyo3::exceptions::PyException, _>(format!(
                    "{:#?}",
                    err
                ))),
            }
        });
        Ok(result.unwrap().to_object(py))
    }
}

#[pyclass]
pub struct PyClient {
    pub inner: Arc<Mutex<Client>>,
    inference: PyInference,
}

#[pymethods]
impl PyClient {
    #[staticmethod]
    fn query(py: Python, query: String) -> PyResult<&PyAny> {
        pyo3_asyncio::tokio::future_into_py(py, async move {
            let mut client = unsafe { CLIENT.as_ref().unwrap().lock().await };
            let result = client.query(query).await;
            match result {
                Ok(value) => {
                    // Convert the HSML Entity into a Python Dictionary
                    // to be usable from Python without have to
                    // pass in a generic parameter to return an abstracted
                    // HSML Entity type, which would not allow it to
                    // be typed as any shape of data structure, and only
                    // the strict type of the generic parameter passed in.
                    let value_str = serde_json::to_string(&value).unwrap();
                    Python::with_gil(|py| {
                        let py_value_str = PyString::new(py, &value_str);
                        let json_module = PyModule::import(py, "json")?;
                        let parsed = json_module.getattr("loads")?.call1((py_value_str,))?;
                        Ok(parsed.to_object(py))
                    })
                }
                Err(err) => {
                    let err_msg = format!("{}", err);
                    Python::with_gil(|_py| {
                        let py_err = PyErr::new::<pyo3::exceptions::PyException, _>(err_msg);
                        Err(py_err)
                    })
                }
            }
        })
    }

    #[getter]
    fn get_inference(&self) -> PyResult<PyInference> {
        Ok(self.inference.clone())
    }
}

#[pyfunction]
pub fn retrieve_auth_token_client_credentials(
    client_id: String,
    client_secret: String,
    token_url: String,
    audience: Option<String>,
    scope: Option<String>,
) -> PyResult<PyObject> {
    Python::with_gil(|py| {
        let result = tokio::runtime::Runtime::new().unwrap().block_on(
            retrieve_auth_token_client_credentials_rs(
                client_id,
                client_secret,
                token_url,
                audience,
                scope,
            ),
        );

        match result {
            Ok(token_response) => {
                let dict = [("access_token", token_response.access_token)].into_py_dict(py);
                Ok(dict.to_object(py))
            }
            Err(err) => Err(PyErr::new::<pyo3::exceptions::PyException, _>(format!(
                "{}",
                err
            ))),
        }
    })
}

#[pyclass]
pub struct PyHSMLEntity {
    pub inner: HSMLEntity,
}

#[pymethods]
impl PyHSMLEntity {
    #[new]
    fn new(kwargs: Option<&PyDict>) -> Self {
        let mut entity = HSMLEntity::new(String::from(""));
        if let Some(kwargs) = kwargs {
            for (key, val) in kwargs {
                match key.to_string().as_str() {
                    "swid" => entity.swid = val.extract().unwrap(),
                    "__archived" => entity.__archived = val.extract().unwrap(),
                    "schema" => entity.schema = val.extract().unwrap(),
                    "name" => entity.name = val.extract().unwrap(),
                    "source_swid" => {
                        // Make sure entity schema type is link if contains source_swid
                        entity.schema =
                            vec![ENTITY_SCHEMA_SWID.to_string(), LINK_SCHEMA_SWID.to_string()];
                        if let Ok(py_list) = val.downcast::<PyList>() {
                            let vec: Vec<Value> = py_list
                                .into_iter()
                                .map(|item| {
                                    // Recursively convert each item in the list to a Value
                                    // You may need to handle different types here as well
                                    Value::String(item.extract::<String>().unwrap())
                                })
                                .collect();
                            entity.source_swid = Some(Value::Array(vec));
                        } else {
                            panic!("Invalid type")
                        }
                    }
                    "destination_swid" => {
                        // Make sure entity schema type is link if contains destination_swid
                        entity.schema =
                            vec![ENTITY_SCHEMA_SWID.to_string(), LINK_SCHEMA_SWID.to_string()];
                        if let Ok(py_list) = val.downcast::<PyList>() {
                            let vec: Vec<Value> = py_list
                                .into_iter()
                                .map(|item| {
                                    // Recursively convert each item in the list to a Value
                                    // You may need to handle different types here as well
                                    Value::String(item.extract::<String>().unwrap())
                                })
                                .collect();
                            entity.destination_swid = Some(Value::Array(vec));
                        } else {
                            panic!("Invalid type")
                        }
                    }
                    _ => {
                        let value = if let Ok(py_str) = val.downcast::<PyString>() {
                            Value::String(py_str.to_string())
                        } else if let Ok(py_bool) = val.downcast::<PyBool>() {
                            Value::Bool(py_bool.is_true())
                        } else if let Ok(py_int) = val.downcast::<PyInt>() {
                            Value::Number(py_int.extract::<i64>().unwrap().into())
                        } else if let Ok(py_int) = val.downcast::<PyLong>() {
                            Value::Number(py_int.extract::<i64>().unwrap().into())
                        } else if let Ok(py_float) = val.downcast::<PyFloat>() {
                            Value::Number(
                                serde_json::Number::from_f64(py_float.extract::<f64>().unwrap())
                                    .unwrap(),
                            )
                        } else if let Ok(py_list) = val.downcast::<PyList>() {
                            let vec: Vec<Value> = py_list
                                .into_iter()
                                .map(|item| {
                                    // Recursively convert each item in the list to a Value
                                    // You may need to handle different types here as well
                                    Value::String(item.extract::<String>().unwrap())
                                })
                                .collect();
                            Value::Array(vec)
                        } else {
                            panic!("Invalid type")
                        };
                        entity.extra_fields.insert(key.to_string(), value);
                    }
                }
            }
        }
        PyHSMLEntity { inner: entity }
    }

    // Out of scope:
    // Implement this if we want a convenience function
    // to instantiate a new link entity
    // fn new_link(kwargs: Option<&PyDict>) -> Self {
    //     let entity = HSMLEntity::new_link(String::from(""), &[], &[], None);
    //     PyHSMLEntity { inner: }
    // }

    // Out of scope:
    // Implement this if you want to give users a way to print
    // the entity to a dictionary
    // #[text_signature = "(self)"]
    // fn to_dict(&self, py: Python) -> PyResult<&PyDict> {
    //     let dict = PyDict::new(py);
    //     dict.set_item("swid", self.inner.swid.clone())?;
    //     // Add other fields of the HSMLEntity struct here
    //     // dict.set_item("field_name", self.inner.field_name.clone())?;
    //     Ok(dict)
    // }

    #[getter]
    fn get_swid(&self) -> PyResult<String> {
        Ok(self.inner.swid.clone())
    }

    #[setter]
    fn set_swid(&mut self, swid: String) {
        self.inner.swid = swid;
    }

    #[getter]
    fn get_destination_swid(&self) -> PyResult<Py<PyAny>> {
        Python::with_gil(|py| {
            let list = PyList::empty(py);
            for item in self
                .inner
                .destination_swid
                .clone()
                .unwrap()
                .as_array()
                .unwrap()
            {
                list.append(PyString::new(py, item.as_str().unwrap()))
                    .unwrap();
            }
            Ok(list.to_object(py))
        })
    }

    #[setter]
    fn set_destination_swid(&mut self, destination_swid: &PyList) {
        let vec: Vec<Value> = destination_swid
            .into_iter()
            .map(|item| {
                // Recursively convert each item in the list to a Value
                // You may need to handle different types here as well
                Value::String(item.extract::<String>().unwrap())
            })
            .collect();
        self.inner.destination_swid = Some(Value::Array(vec));
    }
}