mp-ai-module-controller/src/dictionary.rs

use anyhow::{Context, Result, bail};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, BTreeSet};
use std::fs::{self, File};
use std::io::Write;
#[cfg(test)]
use std::io::{BufRead, BufReader};
use std::path::{Path, PathBuf};

use crate::registry::ActionDefinition;

pub const DICTIONARY_SCHEMA_VERSION: &str = "1.0";

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ActionDictionaryRecord {
    pub action_id: String,
    pub aliases: Vec<String>,
    pub description: String,
    pub intent_examples: Vec<String>,
    pub payload_schema: serde_json::Value,
    pub executable: crate::registry::ExecutableDefinition,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ActionIndex {
    pub schema_version: String,
    pub dictionary_version: String,
    pub registry_checksum: String,
    pub static_base_checksum: String,
    pub codebook_checksum: String,
    pub actions: BTreeMap<String, ActionDictionaryRecord>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct StaticBaseDictionary {
    pub dictionary_version: String,
    pub keys: BTreeMap<String, String>,
    pub types: BTreeMap<String, String>,
    pub errors: BTreeMap<String, String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct AdaptiveCodebook {
    pub actions: BTreeMap<String, String>,
    pub fields: BTreeMap<String, String>,
    pub modules: BTreeMap<String, String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ModelCodebook {
    pub schema_version: String,
    pub dictionary_version: String,
    pub registry_checksum: String,
    pub static_base_checksum: String,
    pub codebook_checksum: String,
    pub static_base: StaticBaseDictionary,
    pub adaptive: AdaptiveCodebook,
}

#[derive(Debug, Clone, Serialize)]
struct ModelCodebookForChecksum<'a> {
    schema_version: &'a str,
    dictionary_version: &'a str,
    registry_checksum: &'a str,
    static_base_checksum: &'a str,
    static_base: &'a StaticBaseDictionary,
    adaptive: &'a AdaptiveCodebook,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DictionarySummary {
    pub action_count: usize,
    pub jsonl_path: PathBuf,
    pub index_path: PathBuf,
    pub codebook_path: PathBuf,
    pub examples_path: PathBuf,
}

pub fn default_dictionary_dir() -> PathBuf {
    PathBuf::from("dictionary")
}

pub fn generate_dictionary_files(
    actions: &[ActionDefinition],
    out_dir: &Path,
) -> Result<DictionarySummary> {
    fs::create_dir_all(out_dir)
        .with_context(|| format!("creating dictionary directory {}", out_dir.display()))?;

    let static_base = load_static_base_dictionary(out_dir)?;
    validate_static_base_dictionary(&static_base)?;

    let records = sorted_records(actions);
    let registry_checksum = checksum_json(&records)?;
    let static_base_checksum = checksum_json(&static_base)?;
    let adaptive = build_adaptive_codebook(&records);
    let codebook = build_model_codebook(
        static_base,
        adaptive,
        registry_checksum,
        static_base_checksum,
    )?;

    let jsonl_path = out_dir.join("actions.jsonl");
    let index_path = out_dir.join("actions.index.json");
    let codebook_path = out_dir.join("model.codebook.json");
    let examples_path = out_dir.join("model.examples.jsonl");

    write_actions_jsonl(&jsonl_path, &records)?;

    let index = ActionIndex {
        schema_version: DICTIONARY_SCHEMA_VERSION.to_string(),
        dictionary_version: codebook.dictionary_version.clone(),
        registry_checksum: codebook.registry_checksum.clone(),
        static_base_checksum: codebook.static_base_checksum.clone(),
        codebook_checksum: codebook.codebook_checksum.clone(),
        actions: records
            .iter()
            .cloned()
            .map(|record| (record.action_id.clone(), record))
            .collect(),
    };
    write_pretty_json(&index_path, &index)?;
    write_pretty_json(&codebook_path, &codebook)?;
    write_model_examples_jsonl(&examples_path, &codebook)?;

    Ok(DictionarySummary {
        action_count: index.actions.len(),
        jsonl_path,
        index_path,
        codebook_path,
        examples_path,
    })
}

pub fn load_action_index(dictionary_dir: &Path) -> Result<ActionIndex> {
    let path = dictionary_dir.join("actions.index.json");
    let raw = fs::read_to_string(&path).with_context(|| {
        format!(
            "reading {}; run `cargo run -- dictionary generate` first",
            path.display()
        )
    })?;
    let index = serde_json::from_str::<ActionIndex>(&raw)
        .with_context(|| format!("parsing {}", path.display()))?;
    Ok(index)
}

pub fn load_model_codebook(dictionary_dir: &Path) -> Result<ModelCodebook> {
    let path = dictionary_dir.join("model.codebook.json");
    let raw = fs::read_to_string(&path).with_context(|| {
        format!(
            "reading {}; run `cargo run -- dictionary generate` first",
            path.display()
        )
    })?;
    let codebook = serde_json::from_str::<ModelCodebook>(&raw)
        .with_context(|| format!("parsing {}", path.display()))?;
    Ok(codebook)
}

pub fn resolve_action_code<'a>(codebook: &'a ModelCodebook, code: &str) -> Option<&'a str> {
    codebook.adaptive.actions.get(code).map(String::as_str)
}

pub fn resolve_field_code<'a>(codebook: &'a ModelCodebook, code: &str) -> Option<&'a str> {
    codebook.adaptive.fields.get(code).map(String::as_str)
}

fn load_static_base_dictionary(dictionary_dir: &Path) -> Result<StaticBaseDictionary> {
    let path = dictionary_dir.join("static-base.json");
    let raw = fs::read_to_string(&path).with_context(|| {
        format!(
            "reading tracked static base dictionary {}; create this file before generating dictionaries",
            path.display()
        )
    })?;
    serde_json::from_str::<StaticBaseDictionary>(&raw)
        .with_context(|| format!("parsing {}", path.display()))
}

fn validate_static_base_dictionary(static_base: &StaticBaseDictionary) -> Result<()> {
    let expected_keys = [("K001", "a"), ("K002", "p"), ("K003", "v"), ("K004", "c")];
    let expected_types = [
        ("T001", "action"),
        ("T002", "field"),
        ("T003", "error"),
        ("T004", "module"),
    ];
    let expected_errors = [
        ("E001", "unknown_action"),
        ("E002", "invalid_payload"),
        ("E003", "checksum_mismatch"),
    ];

    for (code, value) in expected_keys {
        if static_base.keys.get(code).map(String::as_str) != Some(value) {
            bail!("static base dictionary must reserve {code}={value}");
        }
    }
    for (code, value) in expected_types {
        if static_base.types.get(code).map(String::as_str) != Some(value) {
            bail!("static base dictionary must reserve {code}={value}");
        }
    }
    for (code, value) in expected_errors {
        if static_base.errors.get(code).map(String::as_str) != Some(value) {
            bail!("static base dictionary must reserve {code}={value}");
        }
    }
    Ok(())
}

fn sorted_records(actions: &[ActionDefinition]) -> Vec<ActionDictionaryRecord> {
    let mut records = actions
        .iter()
        .cloned()
        .map(ActionDictionaryRecord::from)
        .collect::<Vec<_>>();
    records.sort_by(|left, right| left.action_id.cmp(&right.action_id));
    records
}

fn build_adaptive_codebook(records: &[ActionDictionaryRecord]) -> AdaptiveCodebook {
    let actions = records
        .iter()
        .enumerate()
        .map(|(index, record)| (format_code("A", index), record.action_id.clone()))
        .collect::<BTreeMap<_, _>>();

    let fields = records
        .iter()
        .flat_map(payload_fields)
        .collect::<BTreeSet<_>>()
        .into_iter()
        .enumerate()
        .map(|(index, field)| (format_code("F", index), field))
        .collect::<BTreeMap<_, _>>();

    let modules = records
        .iter()
        .map(|record| record.executable.script.clone())
        .collect::<BTreeSet<_>>()
        .into_iter()
        .enumerate()
        .map(|(index, module)| (format_code("M", index), module))
        .collect::<BTreeMap<_, _>>();

    AdaptiveCodebook {
        actions,
        fields,
        modules,
    }
}

fn build_model_codebook(
    static_base: StaticBaseDictionary,
    adaptive: AdaptiveCodebook,
    registry_checksum: String,
    static_base_checksum: String,
) -> Result<ModelCodebook> {
    let mut codebook = ModelCodebook {
        schema_version: DICTIONARY_SCHEMA_VERSION.to_string(),
        dictionary_version: static_base.dictionary_version.clone(),
        registry_checksum,
        static_base_checksum,
        codebook_checksum: String::new(),
        static_base,
        adaptive,
    };

    codebook.codebook_checksum = checksum_codebook_without_self(&codebook)?;
    Ok(codebook)
}

fn checksum_codebook_without_self(codebook: &ModelCodebook) -> Result<String> {
    checksum_json(&ModelCodebookForChecksum {
        schema_version: &codebook.schema_version,
        dictionary_version: &codebook.dictionary_version,
        registry_checksum: &codebook.registry_checksum,
        static_base_checksum: &codebook.static_base_checksum,
        static_base: &codebook.static_base,
        adaptive: &codebook.adaptive,
    })
}

fn checksum_json<T: Serialize>(value: &T) -> Result<String> {
    let encoded = serde_json::to_vec(value)?;
    let digest = Sha256::digest(encoded);
    Ok(format!("{digest:x}"))
}

fn payload_fields(record: &ActionDictionaryRecord) -> Vec<String> {
    record
        .payload_schema
        .get("properties")
        .and_then(serde_json::Value::as_object)
        .map(|properties| properties.keys().cloned().collect())
        .unwrap_or_default()
}

fn format_code(prefix: &str, zero_based_index: usize) -> String {
    format!("{prefix}{:03}", zero_based_index + 1)
}

fn write_actions_jsonl(path: &Path, records: &[ActionDictionaryRecord]) -> Result<()> {
    let mut jsonl = File::create(path).with_context(|| format!("creating {}", path.display()))?;
    for record in records {
        serde_json::to_writer(&mut jsonl, record)?;
        jsonl.write_all(b"\n")?;
    }
    Ok(())
}

fn write_model_examples_jsonl(path: &Path, codebook: &ModelCodebook) -> Result<()> {
    let mut jsonl = File::create(path).with_context(|| format!("creating {}", path.display()))?;
    let examples = [
        serde_json::json!({
            "kind": "verbose_example",
            "input": "log hello from codex",
            "target": {
                "action_id": "logger",
                "payload": {
                    "message": "hello from codex"
                }
            }
        }),
        serde_json::json!({
            "kind": "compact_example",
            "input": "log hello from codex",
            "target": {
                "v": codebook.dictionary_version,
                "c": codebook.codebook_checksum,
                "a": "A001",
                "p": {
                    "F001": "hello from codex"
                }
            }
        }),
        serde_json::json!({
            "kind": "invalid_checksum_example",
            "input": "log hello with a stale dictionary",
            "target": {
                "v": codebook.dictionary_version,
                "c": "0000000000000000000000000000000000000000000000000000000000000000",
                "a": "A001",
                "p": {
                    "F001": "hello"
                }
            },
            "expected_runtime_result": "reject_checksum_mismatch"
        }),
        serde_json::json!({
            "kind": "synthetic_template",
            "input": "{{natural_language_request}}",
            "target": {
                "v": codebook.dictionary_version,
                "c": codebook.codebook_checksum,
                "a": "{{action_code}}",
                "p": {
                    "{{field_code}}": "{{field_value}}"
                }
            }
        }),
    ];

    for example in examples {
        serde_json::to_writer(&mut jsonl, &example)?;
        jsonl.write_all(b"\n")?;
    }
    Ok(())
}

fn write_pretty_json<T: Serialize>(path: &Path, value: &T) -> Result<()> {
    let pretty = serde_json::to_string_pretty(value)?;
    fs::write(path, format!("{pretty}\n")).with_context(|| format!("writing {}", path.display()))
}

#[cfg(test)]
fn read_jsonl_records(path: &Path) -> Result<Vec<ActionDictionaryRecord>> {
    let file = File::open(path).with_context(|| format!("opening {}", path.display()))?;
    let reader = BufReader::new(file);
    let mut records = Vec::new();

    for line in reader.lines() {
        let line = line?;
        if line.trim().is_empty() {
            continue;
        }
        records.push(serde_json::from_str::<ActionDictionaryRecord>(&line)?);
    }

    Ok(records)
}

impl From<ActionDefinition> for ActionDictionaryRecord {
    fn from(action: ActionDefinition) -> Self {
        Self {
            action_id: action.action_id,
            aliases: action.aliases,
            description: action.description,
            intent_examples: action.intent_examples,
            payload_schema: action.payload_schema,
            executable: action.executable,
        }
    }
}

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

    #[test]
    fn static_base_dictionary_reserves_expected_codes() {
        let temp = temp_dictionary_dir();
        let static_base = load_static_base_dictionary(temp.path()).unwrap();
        validate_static_base_dictionary(&static_base).unwrap();
        assert_eq!(static_base.keys["K001"], "a");
        assert_eq!(static_base.types["T002"], "field");
        assert_eq!(static_base.errors["E003"], "checksum_mismatch");
    }

    #[test]
    fn adaptive_layer_maps_logger_and_message_deterministically() {
        let records = sorted_records(&registered_actions());
        let adaptive = build_adaptive_codebook(&records);
        assert_eq!(adaptive.actions["A001"], "logger");
        assert_eq!(adaptive.fields["F001"], "message");
        assert_eq!(adaptive.modules["M001"], "logger");
    }

    #[test]
    fn checksums_are_stable_for_unchanged_inputs() {
        let temp = temp_dictionary_dir();
        let first = generate_dictionary_files(&registered_actions(), temp.path()).unwrap();
        let first_codebook = load_model_codebook(temp.path()).unwrap();
        let second = generate_dictionary_files(&registered_actions(), temp.path()).unwrap();
        let second_codebook = load_model_codebook(temp.path()).unwrap();

        assert_eq!(first.action_count, second.action_count);
        assert_eq!(
            first_codebook.registry_checksum,
            second_codebook.registry_checksum
        );
        assert_eq!(
            first_codebook.static_base_checksum,
            second_codebook.static_base_checksum
        );
        assert_eq!(
            first_codebook.codebook_checksum,
            second_codebook.codebook_checksum
        );
    }

    #[test]
    fn generates_all_dictionary_outputs() {
        let temp = temp_dictionary_dir();
        let summary = generate_dictionary_files(&registered_actions(), temp.path()).unwrap();

        assert_eq!(summary.action_count, 1);
        assert!(summary.jsonl_path.exists());
        assert!(summary.index_path.exists());
        assert!(summary.codebook_path.exists());
        assert!(summary.examples_path.exists());

        let records = read_jsonl_records(&summary.jsonl_path).unwrap();
        assert_eq!(records.len(), 1);
        assert_eq!(records[0].action_id, "logger");

        let index = load_action_index(temp.path()).unwrap();
        assert!(index.actions.contains_key("logger"));
        assert!(!index.actions.contains_key("log"));

        let codebook = load_model_codebook(temp.path()).unwrap();
        assert_eq!(codebook.adaptive.actions["A001"], "logger");
        assert_eq!(codebook.adaptive.fields["F001"], "message");
        assert!(!codebook.codebook_checksum.is_empty());
    }

    fn temp_dictionary_dir() -> tempfile::TempDir {
        let temp = tempfile::tempdir().expect("temp dir");
        fs::write(
            temp.path().join("static-base.json"),
            include_str!("../dictionary/static-base.json"),
        )
        .unwrap();
        temp
    }
}