consciousness/src/config_writer.rs

// config_writer.rs — Surgical edits to ~/.consciousness/config.json5
//
// Uses json-five's round-trip parser to mutate specific fields while
// preserving the surrounding comments, whitespace, and formatting.

use std::path::Path;

use anyhow::{anyhow, Context as _, Result};
use json_five::rt::parser::{
    from_str, JSONKeyValuePair, JSONObjectContext, JSONValue, KeyValuePairContext,
};

use crate::config::config_path;

/// Read the config, apply `mutate` to the root JSONValue, write it back atomically.
fn edit_config<F: FnOnce(&mut JSONValue) -> Result<()>>(mutate: F) -> Result<()> {
    let path = config_path();
    let src = std::fs::read_to_string(&path)
        .with_context(|| format!("read {}", path.display()))?;

    let mut text = from_str(&src)
        .map_err(|e| anyhow!("parse {}: {}", path.display(), e))?;
    mutate(&mut text.value)?;

    write_atomic(&path, &text.to_string())
}

fn write_atomic(path: &Path, content: &str) -> Result<()> {
    let parent = path.parent()
        .ok_or_else(|| anyhow!("config path has no parent: {}", path.display()))?;
    let tmp = parent.join(format!(
        ".{}.tmp",
        path.file_name().unwrap_or_default().to_string_lossy(),
    ));
    std::fs::write(&tmp, content)
        .with_context(|| format!("write {}", tmp.display()))?;
    std::fs::rename(&tmp, path)
        .with_context(|| format!("rename {} -> {}", tmp.display(), path.display()))?;
    Ok(())
}

/// Match a key JSONValue against a string name. JSON5 allows keys to be
/// unquoted identifiers or single/double-quoted strings.
fn key_matches(key: &JSONValue, name: &str) -> bool {
    match key {
        JSONValue::Identifier(s)
        | JSONValue::DoubleQuotedString(s)
        | JSONValue::SingleQuotedString(s) => s == name,
        _ => false,
    }
}

/// Find (or create) a child object under `parent`, returning a mutable borrow
/// of its key_value_pairs vector.
fn get_or_create_object<'a>(
    parent: &'a mut JSONValue,
    section: &str,
) -> Result<&'a mut Vec<JSONKeyValuePair>> {
    let pairs = match parent {
        JSONValue::JSONObject { key_value_pairs, .. } => key_value_pairs,
        _ => return Err(anyhow!("config root is not an object")),
    };

    // Separate the lookup from the mutable borrow we return — needed to
    // satisfy the borrow checker when we create a new entry.
    let idx = pairs.iter().position(|kvp| key_matches(&kvp.key, section));

    let idx = match idx {
        Some(i) => i,
        None => {
            pairs.push(JSONKeyValuePair {
                key: JSONValue::Identifier(section.to_string()),
                value: JSONValue::JSONObject {
                    key_value_pairs: Vec::new(),
                    context: Some(JSONObjectContext {
                        wsc: (String::new(),),
                    }),
                },
                context: Some(KeyValuePairContext {
                    wsc: (
                        String::from("\n\n    "), // whitespace before ':'
                        String::from(" "),         // whitespace after ':'
                        String::new(),             // whitespace after value
                        Some(String::new()),       // whitespace after trailing comma
                    ),
                }),
            });
            pairs.len() - 1
        }
    };

    match &mut pairs[idx].value {
        JSONValue::JSONObject { key_value_pairs, .. } => Ok(key_value_pairs),
        _ => Err(anyhow!("config key '{}' is not an object", section)),
    }
}

/// Set `section.key` to a literal scalar value (e.g., "1e-7", "42", "true").
/// The literal is parsed as JSON5 so we preserve its source-form on round-trip.
pub fn set_scalar(section: &str, key: &str, literal: &str) -> Result<()> {
    let value = parse_scalar_literal(literal)?;
    edit_config(|root| {
        let pairs = get_or_create_object(root, section)?;

        if let Some(kvp) = pairs.iter_mut().find(|k| key_matches(&k.key, key)) {
            kvp.value = value;
            return Ok(());
        }

        pairs.push(JSONKeyValuePair {
            key: JSONValue::Identifier(key.to_string()),
            value,
            context: Some(KeyValuePairContext {
                wsc: (
                    String::from("\n        "),
                    String::from(" "),
                    String::new(),
                    Some(String::new()),
                ),
            }),
        });
        Ok(())
    })
}

/// Parse a scalar literal by round-tripping it through json-five. Keeps us
/// consistent with whatever scalars the library considers valid (hex,
/// exponents, Infinity, etc.).
fn parse_scalar_literal(literal: &str) -> Result<JSONValue> {
    let text = from_str(literal)
        .map_err(|e| anyhow!("parse literal {:?}: {}", literal, e))?;
    match text.value {
        JSONValue::JSONObject { .. } | JSONValue::JSONArray { .. } => {
            Err(anyhow!("set_scalar only accepts scalar literals, got {:?}", literal))
        }
        v => Ok(v),
    }
}

/// Convenience: set `learn.threshold` to the given f64.
pub fn set_learn_threshold(value: f64) -> Result<()> {
    // {:e} gives the minimal scientific notation that preserves the value.
    set_scalar("learn", "threshold", &format!("{:e}", value))
}

#[cfg(test)]
mod tests {
    use super::*;

    // In-memory variant of set_scalar — used to test the mutation logic
    // without touching disk.
    fn set_scalar_inline(
        root: &mut JSONValue,
        section: &str,
        key: &str,
        literal: &str,
    ) -> Result<()> {
        let value = parse_scalar_literal(literal)?;
        let pairs = get_or_create_object(root, section)?;
        if let Some(kvp) = pairs.iter_mut().find(|k| key_matches(&k.key, key)) {
            kvp.value = value;
            return Ok(());
        }
        pairs.push(JSONKeyValuePair {
            key: JSONValue::Identifier(key.to_string()),
            value,
            context: Some(KeyValuePairContext {
                wsc: (
                    String::from("\n        "),
                    String::from(" "),
                    String::new(),
                    Some(String::new()),
                ),
            }),
        });
        Ok(())
    }

    fn edit_str<F: FnOnce(&mut JSONValue) -> Result<()>>(src: &str, f: F) -> Result<String> {
        let mut text = from_str(src).map_err(|e| anyhow!("{}", e))?;
        f(&mut text.value)?;
        Ok(text.to_string())
    }

    #[test]
    fn replaces_existing_scalar() {
        let src = r#"{
    // threshold for learning
    learn: {
        threshold: 0.001, // the old value
    },
}"#;
        let out = edit_str(src, |root| {
            set_scalar_inline(root, "learn", "threshold", "1e-7")
        }).unwrap();
        assert!(out.contains("1e-7"), "output: {}", out);
        assert!(out.contains("// threshold for learning"));
        assert!(out.contains("// the old value"));
        assert!(!out.contains("0.001"));
    }

    #[test]
    fn creates_missing_section() {
        let src = r#"{
    // comment
    memory: { user_name: "Kent" },
}"#;
        let out = edit_str(src, |root| {
            set_scalar_inline(root, "learn", "threshold", "1e-7")
        }).unwrap();
        assert!(out.contains("learn"));
        assert!(out.contains("1e-7"));
        assert!(out.contains("// comment"));
        assert!(out.contains(r#"user_name: "Kent""#));
    }

    #[test]
    fn preserves_comments_in_siblings() {
        let src = r#"{
    memory: {
        // sensitive setting
        user_name: "Kent", // name
    },
    learn: {
        threshold: 0.5,
    },
}"#;
        let out = edit_str(src, |root| {
            set_scalar_inline(root, "learn", "threshold", "1e-9")
        }).unwrap();
        assert!(out.contains("// sensitive setting"));
        assert!(out.contains("// name"));
        assert!(out.contains("1e-9"));
        assert!(!out.contains("0.5"));
    }

    #[test]
    fn adds_key_to_existing_empty_section() {
        let src = r#"{
    learn: {},
}"#;
        let out = edit_str(src, |root| {
            set_scalar_inline(root, "learn", "threshold", "42")
        }).unwrap();
        assert!(out.contains("threshold"), "output: {}", out);
        assert!(out.contains("42"));
    }

    #[test]
    fn realistic_config_adds_learn_section() {
        // Mirrors the shape of ~/.consciousness/config.json5 — multiple
        // sections, comments, mixed tab/space indent, trailing commas.
        let src = r#"{
    deepinfra: {
        api_key: "bcachefs-agents-2026",
        base_url: "http://example/v1",
    },

    // Named models
    models: {
        "27b": {
            backend: "deepinfra",
            model_id: "Qwen/Qwen3.5-27B",
        },
    },

    default_model: "27b",

    memory: {
        user_name: "Kent",
        // Active agent types
        agent_types: ["linker", "organize"],
    },

    compaction: {
        hard_threshold_pct: 90,
    },
}"#;
        let out = edit_str(src, |root| {
            set_scalar_inline(root, "learn", "threshold", "1e-7")
        }).unwrap();

        // Core assertions: comments and sibling sections survive.
        assert!(out.contains(r#"api_key: "bcachefs-agents-2026""#));
        assert!(out.contains("// Named models"));
        assert!(out.contains("// Active agent types"));
        assert!(out.contains(r#"user_name: "Kent""#));
        assert!(out.contains("hard_threshold_pct: 90"));

        // New section added.
        assert!(out.contains("learn"));
        assert!(out.contains("1e-7"));

        // Parse result should parse back without error (real json5 parser).
        let reparsed: serde_json::Value = json5::from_str(&out)
            .expect("mutated output must be valid JSON5");
        let threshold = reparsed.pointer("/learn/threshold").expect("learn.threshold exists");
        assert_eq!(threshold.as_f64(), Some(1e-7));
    }

    #[test]
    fn realistic_config_updates_existing_threshold() {
        let src = r#"{
    learn: {
        // The divergence threshold
        threshold: 0.001,
    },
    memory: { user_name: "Kent" },
}"#;
        let out = edit_str(src, |root| {
            set_scalar_inline(root, "learn", "threshold", "5e-8")
        }).unwrap();
        assert!(out.contains("5e-8"));
        assert!(!out.contains("0.001"));
        assert!(out.contains("// The divergence threshold"));

        let reparsed: serde_json::Value = json5::from_str(&out).unwrap();
        assert_eq!(reparsed.pointer("/learn/threshold").and_then(|v| v.as_f64()), Some(5e-8));
    }

    #[test]
    fn roundtrip_stable_without_change() {
        let src = r#"{
    // heading
    a: 1,
    b: { c: 2 }, // inline
}"#;
        let text = from_str(src).unwrap();
        assert_eq!(text.to_string(), src);
    }
}