py_mcai_worker_sdk 1.1.3

use crate::helpers::{get_parameters, py_err_to_string};
use mcai_worker_sdk::prelude::*;
use pyo3::{types::*, Python};
use schemars::{
  gen::SchemaGenerator,
  schema::{InstanceType, ObjectValidation, Schema, SchemaObject},
  JsonSchema,
};
use serde_json::Value;
use std::collections::{BTreeMap, HashMap};

#[derive(Debug, Deserialize)]
pub struct PythonWorkerParameters {
  #[serde(flatten)]
  pub(crate) parameters: HashMap<String, Value>,
}

impl JsonSchema for PythonWorkerParameters {
  fn schema_name() -> String {
    "PythonWorkerParameters".to_string()
  }

  fn json_schema(_gen: &mut SchemaGenerator) -> Schema {
    let parameters = get_parameters();

    let schema_parameters: BTreeMap<String, Schema> = parameters
      .iter()
      .map(|parameter| {
        let parameter_type = &parameter.kind[0];

        let instance_type = if parameter.required {
          get_instance_type_from_parameter_type(parameter_type).into()
        } else {
          vec![
            get_instance_type_from_parameter_type(parameter_type),
            InstanceType::Null,
          ]
          .into()
        };

        let instance_type = Some(instance_type);

        let object = SchemaObject {
          instance_type,
          ..Default::default()
        };

        (parameter.identifier.clone(), object.into())
      })
      .collect();

    let schema = SchemaObject {
      instance_type: Some(InstanceType::Object.into()),
      object: Some(Box::new(ObjectValidation {
        properties: schema_parameters,
        ..Default::default()
      })),
      ..Default::default()
    };

    schema.into()
  }
}

pub(crate) fn get_instance_type_from_parameter_type(
  parameter_type: &WorkerParameterType,
) -> InstanceType {
  match parameter_type {
    WorkerParameterType::String => InstanceType::String,
    WorkerParameterType::ArrayOfStrings => InstanceType::Array,
    WorkerParameterType::Boolean => InstanceType::Boolean,
    WorkerParameterType::Credential => InstanceType::String,
    WorkerParameterType::Integer => InstanceType::Integer,
    WorkerParameterType::Requirements => InstanceType::Object,
  }
}

pub fn build_parameters(parameters: PythonWorkerParameters, py: Python) -> Result<&PyDict> {
  let list_of_parameters = PyDict::new(py);
  for (key, value) in parameters.parameters {
    serde_json_to_pyo3_value(&key, &value, list_of_parameters, py).map_err(|e| {
      MessageError::ParameterValueError(format!(
        "Cannot build parameters: {}",
        py_err_to_string(py, e)
      ))
    })?;
  }
  Ok(list_of_parameters)
}

fn serde_json_to_pyo3_value(
  key: &str,
  value: &Value,
  result: &PyDict,
  py: Python,
) -> pyo3::PyResult<()> {
  match value {
    Value::Null => {}
    Value::Bool(boolean) => result.set_item(key, boolean)?,
    Value::Number(number) => result.set_item(key, number.as_u64())?,
    Value::String(content) => result.set_item(key, content)?,
    Value::Array(values) => {
      let list = PyList::empty(py);
      for value in values {
        add_value_to_py_list(value, list, py)?;
      }

      result.set_item(key, list)?;
    }
    Value::Object(map) => {
      let object = PyDict::new(py);
      for (key, value) in map.iter() {
        serde_json_to_pyo3_value(key, value, object, py)?;
      }
      result.set_item(key, object)?;
    }
  }
  Ok(())
}

fn add_value_to_py_list(value: &Value, list: &PyList, py: Python) -> pyo3::PyResult<()> {
  match value {
    Value::String(string) => {
      list.append(string)?;
    }
    Value::Null => {}
    Value::Bool(boolean) => {
      list.append(boolean)?;
    }
    Value::Number(number) => {
      list.append(number.as_u64())?;
    }
    Value::Array(values) => {
      let sub_list = PyList::empty(py);
      for value in values {
        add_value_to_py_list(value, sub_list, py)?;
      }
      list.append(sub_list)?;
    }
    Value::Object(map) => {
      let object = PyDict::new(py);
      for (key, value) in map.iter() {
        serde_json_to_pyo3_value(key, value, object, py)?;
      }
      list.append(object)?;
    }
  }
  Ok(())
}