oth_rvault 0.4.0

use serde_yaml::Value;

use crate::types::{new_key_path, ProcessHandling, TypedValue, ValueType};

fn typed_value_from_yaml(val: &Value) -> Option<(TypedValue, ValueType)> {
    match val {
        Value::Bool(b) => Some((TypedValue::Bool(*b), ValueType::Bool)),
        Value::String(s) => Some((TypedValue::String(s.clone()), ValueType::String)),
        Value::Number(n) => {
            if let Some(i) = n.as_i64() {
                // serde_yaml reports is_i64() for numbers parsed from integer literals
                if n.is_i64() {
                    return Some((TypedValue::Integer(i), ValueType::Integer));
                }
            }
            n.as_f64()
                .map(|f| (TypedValue::Number(f), ValueType::Number))
        }
        Value::Null => Some((TypedValue::Null, ValueType::Null)),
        // Ansible's `!vault` tagged values and other custom tags – extract the inner string.
        Value::Tagged(tagged) => match &tagged.value {
            Value::String(s) => Some((TypedValue::String(s.clone()), ValueType::String)),
            inner => typed_value_from_yaml(inner),
        },
        _ => None, // Mapping/Sequence handled by recursion
    }
}

fn typed_value_to_yaml(typed: TypedValue) -> Value {
    match typed {
        TypedValue::Bool(b) => Value::Bool(b),
        TypedValue::Integer(i) => Value::Number(serde_yaml::Number::from(i)),
        TypedValue::Number(f) => Value::Number(serde_yaml::Number::from(f)),
        TypedValue::String(s) => Value::String(s),
        TypedValue::Null => Value::Null,
    }
}

type Processor<'a> = &'a mut dyn FnMut(
    TypedValue,
    ValueType,
    &str,
) -> Result<(TypedValue, ValueType, ProcessHandling), Box<dyn std::error::Error>>;

fn traverse(
    val: &mut Value,
    key_path: &str,
    processor: Processor<'_>,
    filter_keys: &[String],
) -> Result<bool, Box<dyn std::error::Error>> {
    // Sub-node path: if this node matches a filter key and is Mapping/Sequence, offer the
    // whole node to the processor as a serialized YAML string.
    // - Process: replace node with the processor's result (encrypt case).
    // - Skip: fall through to normal per-leaf recursion (decrypt on a node whose items
    //   are individually encrypted, not as a single blob).
    // - Cancel: abort.
    if !filter_keys.is_empty()
        && filter_keys.iter().any(|k| k.as_str() == key_path)
        && matches!(val, Value::Mapping(_) | Value::Sequence(_))
    {
        let serialized = format!("__subnode__\n{}", serde_yaml::to_string(val)?);
        match processor(TypedValue::String(serialized), ValueType::String, key_path) {
            Ok((new_typed, _, ProcessHandling::Process)) => {
                *val = typed_value_to_yaml(new_typed);
                return Ok(true);
            }
            Ok((_, _, ProcessHandling::Cancel)) => return Ok(false),
            Ok((_, _, ProcessHandling::Skip)) => {
                // fall through to normal recursion below
            }
            Err(e) => {
                eprintln!("error processing key '{}': {}", key_path, e);
                return Ok(false);
            }
        }
    }

    match val {
        Value::Mapping(map) => {
            // Collect keys first so we can get_mut on each without borrow conflicts.
            let keys: Vec<Value> = map.keys().cloned().collect();
            for yaml_key in &keys {
                let key_str = match yaml_key {
                    Value::String(s) => s.clone(),
                    Value::Number(n) => n.to_string(),
                    Value::Bool(b) => b.to_string(),
                    _ => continue, // null/mapping/sequence keys are not supported
                };
                let new_path = new_key_path(key_path, &key_str);
                if let Some(child) = map.get_mut(yaml_key) {
                    if !traverse(child, &new_path, processor, filter_keys)? {
                        return Ok(false);
                    }
                }
            }
        }
        Value::Sequence(seq) => {
            for item in seq.iter_mut() {
                if !traverse(item, key_path, processor, filter_keys)? {
                    return Ok(false);
                }
            }
        }
        scalar => {
            if !filter_keys.is_empty()
                && !filter_keys.iter().any(|k| k.as_str() == key_path)
            {
                return Ok(true);
            }
            if let Some((typed, vt)) = typed_value_from_yaml(scalar) {
                match processor(typed, vt, key_path) {
                    Ok((new_typed, _, ProcessHandling::Process)) => {
                        // Sub-node decryption: the marker prefix "__subnode__\n" identifies
                        // values that were encrypted as a whole YAML structure. Strip the
                        // prefix and restore the structured node.
                        let mut restored = false;
                        if let TypedValue::String(ref s) = new_typed {
                            if let Some(inner) = s.strip_prefix("__subnode__\n") {
                                if let Ok(parsed) = serde_yaml::from_str::<Value>(inner) {
                                    if matches!(parsed, Value::Mapping(_) | Value::Sequence(_)) {
                                        *scalar = parsed;
                                        restored = true;
                                    }
                                }
                            }
                        }
                        if !restored {
                            *scalar = typed_value_to_yaml(new_typed);
                        }
                    }
                    Ok((_, _, ProcessHandling::Cancel)) => return Ok(false),
                    Ok((_, _, ProcessHandling::Skip)) => {}
                    Err(e) => {
                        eprintln!("error processing key '{}': {}", key_path, e);
                        return Ok(false);
                    }
                }
            }
        }
    }
    Ok(true)
}

/// Read a YAML file, apply `processor` to every matching scalar value, and return the result in memory.
/// `filter_keys` is a list of dot-notation key paths to restrict processing to; empty means all values.
pub fn process_in_memory(
    input: &str,
    filter_keys: &[String],
    processor: Processor<'_>,
) -> Result<Value, Box<dyn std::error::Error>> {
    let content = std::fs::read_to_string(input)
        .map_err(|e| format!("error reading '{}': {}", input, e))?;
    let mut root: Value = serde_yaml::from_str(&content)
        .map_err(|e| format!("error parsing YAML from '{}': {}", input, e))?;
    if !traverse(&mut root, "", processor, filter_keys)? {
        return Err("processing canceled".into());
    }
    Ok(root)
}

/// Read a YAML file, apply `processor` to every matching scalar value, and write the result.
/// `filter_keys` is a list of dot-notation key paths to restrict processing to; empty means all values.
/// Vault-encrypted strings (with embedded newlines) are written as YAML literal block scalars.
pub fn process_file(
    input: &str,
    output: &str,
    filter_keys: &[String],
    processor: Processor<'_>,
) -> Result<(), Box<dyn std::error::Error>> {
    let content = std::fs::read_to_string(input)
        .map_err(|e| format!("error reading '{}': {}", input, e))?;
    let mut root: Value = serde_yaml::from_str(&content)
        .map_err(|e| format!("error parsing YAML from '{}': {}", input, e))?;

    if !traverse(&mut root, "", processor, filter_keys)? {
        return Err("processing canceled".into());
    }

    let yaml_str = serde_yaml::to_string(&root)?;
    if output == "stdout" {
        print!("{}", yaml_str);
    } else {
        std::fs::write(output, &yaml_str)
            .map_err(|e| format!("error writing '{}': {}", output, e))?;
    }
    Ok(())
}

/// Return the dot-notation paths of all vault-encrypted values in a YAML file.
pub fn get_encrypted_keys(input: &str) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let content = std::fs::read_to_string(input)?;
    let root: Value = serde_yaml::from_str(&content)?;
    let mut encrypted = Vec::new();
    collect_encrypted_yaml(&root, "", &mut encrypted);
    Ok(encrypted)
}

fn collect_encrypted_yaml(val: &Value, key_path: &str, out: &mut Vec<String>) {
    match val {
        Value::Mapping(map) => {
            for (key, child) in map.iter() {
                let key_str = match key {
                    Value::String(s) => s.clone(),
                    Value::Number(n) => n.to_string(),
                    Value::Bool(b) => b.to_string(),
                    _ => continue,
                };
                collect_encrypted_yaml(child, &new_key_path(key_path, &key_str), out);
            }
        }
        Value::Sequence(seq) => {
            for item in seq {
                collect_encrypted_yaml(item, key_path, out);
            }
        }
        Value::String(s) if s.contains("$ANSIBLE_VAULT;") => {
            out.push(key_path.to_string());
        }
        Value::Tagged(t) => {
            if let Value::String(s) = &t.value {
                if s.contains("$ANSIBLE_VAULT;") {
                    out.push(key_path.to_string());
                }
            }
        }
        _ => {}
    }
}

/// Collect dot-notation key paths of all nodes exactly at `level` depth.
/// Nodes in branches that don't reach `level` are omitted.
pub fn collect_paths_at_level(input: &str, level: u32) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let content = std::fs::read_to_string(input)
        .map_err(|e| format!("error reading '{}': {}", input, e))?;
    let root: Value = serde_yaml::from_str(&content)
        .map_err(|e| format!("error parsing YAML from '{}': {}", input, e))?;
    let mut paths = Vec::new();
    collect_at_level(&root, "", level, &mut paths);
    Ok(paths)
}

fn collect_at_level(val: &Value, key_path: &str, remaining: u32, out: &mut Vec<String>) {
    if remaining == 0 {
        out.push(key_path.to_string());
        return;
    }
    if let Value::Mapping(map) = val {
        for (yaml_key, child) in map.iter() {
            let key_str = match yaml_key {
                Value::String(s) => s.clone(),
                Value::Number(n) => n.to_string(),
                Value::Bool(b) => b.to_string(),
                _ => continue,
            };
            collect_at_level(child, &new_key_path(key_path, &key_str), remaining - 1, out);
        }
    }
    // Sequences and scalars above the target level are not at the target depth: skip.
}