stepflow-json 0.0.3

use std::collections::HashMap;
use tinyjson::JsonValue;
use stepflow::{prelude::*, step::StepId};
use stepflow::object::{ObjectStore, IdError};
use stepflow::step::Step;
use stepflow::{Session, Error};
use stepflow::data::{VarId, StringVar};
use super::error::StepFlowParseError;
use super::json::{jsonval_get_obj, jsonval_obj_get_vec_str};

const NAME_ROOT_STEP: &str = "$root";
const STEP_INPUTS_KEY: &str = "inputs";
const STEP_OUTPUTS_KEY: &str = "outputs";
const STEP_SUBSTEPS_KEY: &str = "substeps";


fn parse_var_names(step_info: &HashMap<String, JsonValue>, key: &str, var_store: &ObjectStore<Box<dyn Var + Send + Sync>, VarId>)
    -> Result<Vec<VarId>, StepFlowParseError>
{
  let var_names = jsonval_obj_get_vec_str(step_info, key)?;
  var_names
    .iter()
    .map(|name| {
      var_store.id_from_name(name)
        .map(|name| name.clone())
        .ok_or_else(|| StepFlowParseError::UnexpectedValue((*name).to_owned()))
    })
    .collect::<Result<_, _>>()
    .and_then(|r| Ok(r))
}

pub fn parse_steps_json(session: &mut Session, steps_json: &HashMap<String, JsonValue>, allow_implicit_var: bool)
    -> Result<(), StepFlowParseError> 
{
  // Create implicit vars from Steps
  if allow_implicit_var {
    for (step_name, step_info) in steps_json {
      let step_info = jsonval_get_obj(step_info, step_name)?;
      for key in vec![STEP_INPUTS_KEY, STEP_OUTPUTS_KEY] {
        let varnames = jsonval_obj_get_vec_str(step_info, key);
        if let Ok(varnames) = varnames {
          for varname in varnames {
            if session.var_store().id_from_name(varname) == None {
              session.var_store_mut().insert_new_named(varname.clone(), |id| Ok(StringVar::new(id).boxed()))
                .map_err(|e| Error::VarId(e))?;
            }
          }
        }
      }
    }
  }

  // Create Steps
  // steps in 2 passes.
  // 1. register just the steps, no sub-steps since it's possible they'll be registered later
  // 2. once all the steps are registered, assign the child sub-steps
  //let mut stepid_to_substep_names = HashMap::with_capacity(steps_json.len());

  // Pass 1: Register the steps
  for (step_name, step_info) in steps_json {
    let step_info = jsonval_get_obj(step_info, step_name)?;
  
    let input_var_ids = match step_info.contains_key(STEP_INPUTS_KEY) {
      true => Some(parse_var_names(step_info, STEP_INPUTS_KEY, session.var_store())?),
      false => None,
    };
    let output_var_ids = parse_var_names(step_info, STEP_OUTPUTS_KEY, session.var_store())?;

    session.step_store_mut().insert_new_named(step_name.clone(), |step_id| {
        Ok(Step::new(step_id, input_var_ids, output_var_ids))
      })
      .map_err(|e| Error::StepId(e))?;
  }

  // Pass 2: Assign child steps
  for (step_name, step_info) in steps_json {
    let step_info = jsonval_get_obj(step_info, step_name)?;
    if !step_info.contains_key(STEP_SUBSTEPS_KEY) {
      continue;
    }
    let substep_names = jsonval_obj_get_vec_str(step_info, STEP_SUBSTEPS_KEY)?;
    let substep_ids = substep_names.into_iter().map(|substep_name| {
        session.step_store().id_from_name(substep_name)
          .ok_or_else(|| Error::StepId(IdError::NoSuchName(substep_name.clone())))
          .map(|id| id.clone())
      })
      .collect::<Result<Vec<StepId>, _>>()?;

    let step_id = session.step_store()
      .id_from_name(step_name)
      .ok_or_else(|| Error::StepId(stepflow::object::IdError::NoSuchName(step_name.clone())))?
      .clone();
    let step = session.step_store_mut().get_mut(&step_id).ok_or_else(|| Error::StepId(stepflow::object::IdError::IdMissing(step_id.clone())))?;
    for substep_id in substep_ids {
      step.push_substep(substep_id.clone())
    }
  }

  // Set Root Step
  let root_step_id = session.step_store()
      .id_from_name(NAME_ROOT_STEP)
      .ok_or_else(|| Error::StepId(IdError::NoSuchName(NAME_ROOT_STEP.to_owned())))?.clone();
  session.push_root_substep(root_step_id);

  Ok(())
}