kibble 0.1.0

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::BTreeMap;
use std::path::Path;

#[derive(Deserialize)]
pub struct SoulConfig {
    pub name: String,
    pub version: String,
    #[serde(default)] pub description: Option<String>,
    #[serde(default)] pub persona: Persona,
    #[serde(default)] pub sampling: Sampling,
    #[serde(default, rename = "tool")] pub tools: Vec<Tool>,
    #[serde(default)] pub template: Option<Template>,
}

#[derive(Deserialize, Default)]
pub struct Persona {
    #[serde(default)] pub system_prompt: Option<String>,
    #[serde(default)] pub default: Option<String>,
    #[serde(default)] pub prompts: BTreeMap<String, String>,
}

#[derive(Deserialize, Default)]
pub struct Sampling {
    pub temperature: Option<f64>,
    pub top_p: Option<f64>,
    pub top_k: Option<u64>,
    pub repetition_penalty: Option<f64>,
    pub max_new_tokens: Option<u64>,
}

#[derive(Deserialize)]
pub struct Tool {
    pub name: String,
    pub description: String,
    #[serde(default)] pub parameters: Option<String>,
}

fn default_marker() -> String { "{{SOUL_DEFAULT_SYSTEM}}".to_string() }

#[derive(Deserialize)]
pub struct Template {
    pub base: String,
    #[serde(default = "default_marker")] pub marker: String,
}

pub fn parse_soul(path: &Path) -> std::io::Result<SoulConfig> {
    let text = std::fs::read_to_string(path)?;
    toml::from_str(&text).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, format!("soul.toml: {e}")))
}

pub fn resolve_default_prompt(s: &SoulConfig) -> Result<String, String> {
    let key = s.persona.default.as_deref().unwrap_or("system_prompt");
    let prompt = if key == "system_prompt" {
        s.persona.system_prompt.clone()
    } else {
        s.persona.prompts.get(key).cloned()
    };
    match prompt {
        Some(p) if !p.trim().is_empty() => Ok(p),
        _ => Err(format!("no default prompt: [persona].default = '{key}' resolves to nothing")),
    }
}

pub fn generation_config_json(s: &Sampling) -> String {
    let mut m = serde_json::Map::new();
    if let Some(t) = s.temperature {
        m.insert("temperature".into(), serde_json::json!(t));
        m.insert("do_sample".into(), serde_json::json!(t > 0.0));
    }
    if let Some(v) = s.top_p { m.insert("top_p".into(), serde_json::json!(v)); }
    if let Some(v) = s.top_k { m.insert("top_k".into(), serde_json::json!(v)); }
    if let Some(v) = s.repetition_penalty { m.insert("repetition_penalty".into(), serde_json::json!(v)); }
    if let Some(v) = s.max_new_tokens { m.insert("max_new_tokens".into(), serde_json::json!(v)); }
    serde_json::to_string_pretty(&serde_json::Value::Object(m)).unwrap_or_else(|_| "{}".into())
}

pub fn tools_json(tools: &[Tool]) -> Result<String, String> {
    let mut arr = Vec::new();
    for t in tools {
        let params: serde_json::Value = match &t.parameters {
            Some(p) => serde_json::from_str(p).map_err(|e| format!("tool '{}' parameters: {e}", t.name))?,
            None => serde_json::json!({"type":"object","properties":{}}),
        };
        arr.push(serde_json::json!({
            "type": "function",
            "function": { "name": t.name, "description": t.description, "parameters": params }
        }));
    }
    serde_json::to_string_pretty(&serde_json::Value::Array(arr)).map_err(|e| e.to_string())
}

/// The escaped *body* of a JSON string (no surrounding quotes) — safe to drop inside a double-quoted
/// literal in a jinja template.
pub fn json_escape_body(s: &str) -> String {
    let quoted = serde_json::to_string(s).unwrap_or_else(|_| "\"\"".into());
    quoted[1..quoted.len() - 1].to_string()
}

pub fn apply_template(base: &str, marker: &str, prompt: &str) -> Result<String, String> {
    if !base.contains(marker) {
        return Err(format!("template marker '{marker}' not found in base template"));
    }
    Ok(base.replace(marker, &json_escape_body(prompt)))
}

pub fn bundle_sha256(files: &[(String, Vec<u8>)]) -> String {
    let mut sorted: Vec<&(String, Vec<u8>)> = files.iter().collect();
    sorted.sort_by(|a, b| a.0.cmp(&b.0));
    let mut h = Sha256::new();
    for (name, bytes) in sorted {
        h.update(name.as_bytes());
        h.update([0u8]);
        h.update(bytes);
        h.update([0u8]);
    }
    format!("{:x}", h.finalize())
}

#[derive(Serialize)]
pub struct SoulManifest {
    pub name: String,
    pub version: String,
    pub description: Option<String>,
    pub default_prompt: String,
    pub files: Vec<String>,
    pub sha256: String,
}

pub fn run_soul_build(soul_path: &Path, out: &Path) -> std::io::Result<SoulManifest> {
    let soul = parse_soul(soul_path)?;
    let ie = |e: String| std::io::Error::new(std::io::ErrorKind::InvalidData, e);
    let prompt = resolve_default_prompt(&soul).map_err(ie)?;
    let tools = tools_json(&soul.tools).map_err(ie)?;

    let mut files: Vec<(String, Vec<u8>)> = vec![
        ("generation_config.json".into(), generation_config_json(&soul.sampling).into_bytes()),
        ("system_prompt.txt".into(), prompt.clone().into_bytes()),
        ("tools.json".into(), tools.into_bytes()),
    ];
    if let Some(t) = &soul.template {
        let soul_dir = soul_path.parent().unwrap_or(Path::new("."));
        let base = std::fs::read_to_string(soul_dir.join(&t.base))?;
        let rendered = apply_template(&base, &t.marker, &prompt).map_err(ie)?;
        files.push(("chat_template.jinja".into(), rendered.into_bytes()));
    }

    let sha = bundle_sha256(&files);
    let mut names: Vec<String> = files.iter().map(|(n, _)| n.clone()).collect();
    names.push("manifest.json".to_string());
    names.sort();

    let manifest = SoulManifest {
        name: soul.name.clone(),
        version: soul.version.clone(),
        description: soul.description.clone(),
        default_prompt: prompt,
        files: names,
        sha256: sha,
    };

    std::fs::create_dir_all(out)?;
    for (name, bytes) in &files {
        std::fs::write(out.join(name), bytes)?;
    }
    std::fs::write(out.join("manifest.json"),
        serde_json::to_string_pretty(&manifest).map_err(|e| ie(e.to_string()))?)?;
    Ok(manifest)
}

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

    fn sample_toml() -> &'static str {
        r#"
name = "kibble-assistant"
version = "0.1.0"
description = "test"
[persona]
system_prompt = "You are Kibble."
[persona.prompts]
blog = "Write like a blog."
[sampling]
temperature = 0.6
top_p = 0.95
[[tool]]
name = "web_search"
description = "Search the web."
parameters = "{\"type\":\"object\",\"properties\":{\"q\":{\"type\":\"string\"}}}"
"#
    }

    #[test]
    fn parse_soul_reads_all() {
        let dir = std::env::temp_dir().join(format!("kibble_soul_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir); std::fs::create_dir_all(&dir).unwrap();
        let p = dir.join("soul.toml");
        std::fs::write(&p, sample_toml()).unwrap();
        let s = parse_soul(&p).unwrap();
        assert_eq!(s.name, "kibble-assistant");
        assert_eq!(s.persona.system_prompt.as_deref(), Some("You are Kibble."));
        assert_eq!(s.tools.len(), 1);
        assert_eq!(s.sampling.temperature, Some(0.6));
    }

    #[test]
    fn resolve_default_prompt_cases() {
        let s: SoulConfig = toml::from_str(sample_toml()).unwrap();
        assert_eq!(resolve_default_prompt(&s).unwrap(), "You are Kibble."); // default = system_prompt
        let named: SoulConfig = toml::from_str("name=\"x\"\nversion=\"1\"\n[persona]\ndefault=\"blog\"\n[persona.prompts]\nblog=\"B\"\n").unwrap();
        assert_eq!(resolve_default_prompt(&named).unwrap(), "B");
        let none: SoulConfig = toml::from_str("name=\"x\"\nversion=\"1\"\n").unwrap();
        assert!(resolve_default_prompt(&none).is_err());
    }

    #[test]
    fn generation_config_only_set_fields() {
        let s = Sampling { temperature: Some(0.6), top_p: Some(0.95), top_k: None, repetition_penalty: None, max_new_tokens: None };
        let j: serde_json::Value = serde_json::from_str(&generation_config_json(&s)).unwrap();
        assert_eq!(j["temperature"], 0.6);
        assert_eq!(j["top_p"], 0.95);
        assert_eq!(j["do_sample"], true);
        assert!(j.get("top_k").is_none());
        // temperature 0 → do_sample false
        let s0 = Sampling { temperature: Some(0.0), ..Default::default() };
        let j0: serde_json::Value = serde_json::from_str(&generation_config_json(&s0)).unwrap();
        assert_eq!(j0["do_sample"], false);
    }

    #[test]
    fn tools_json_openai_shape_and_errors() {
        let tools = vec![Tool { name: "web_search".into(), description: "Search.".into(),
            parameters: Some(r#"{"type":"object","properties":{"q":{"type":"string"}}}"#.into()) }];
        let j: serde_json::Value = serde_json::from_str(&tools_json(&tools).unwrap()).unwrap();
        assert_eq!(j[0]["type"], "function");
        assert_eq!(j[0]["function"]["name"], "web_search");
        assert_eq!(j[0]["function"]["parameters"]["type"], "object");
        // omitted parameters → default empty object
        let t2 = vec![Tool { name: "n".into(), description: "d".into(), parameters: None }];
        let j2: serde_json::Value = serde_json::from_str(&tools_json(&t2).unwrap()).unwrap();
        assert_eq!(j2[0]["function"]["parameters"]["type"], "object");
        // invalid JSON → error
        let bad = vec![Tool { name: "n".into(), description: "d".into(), parameters: Some("{not json".into()) }];
        assert!(tools_json(&bad).is_err());
    }

    #[test]
    fn escape_and_template() {
        assert_eq!(json_escape_body("a\"b\\c\nd"), "a\\\"b\\\\c\\nd");
        let out = apply_template("prefix {{SOUL_DEFAULT_SYSTEM}} suffix", "{{SOUL_DEFAULT_SYSTEM}}", "he said \"hi\"").unwrap();
        assert_eq!(out, "prefix he said \\\"hi\\\" suffix");
        assert!(apply_template("no marker here", "{{SOUL_DEFAULT_SYSTEM}}", "x").is_err());
    }

    #[test]
    fn bundle_sha256_deterministic_order_independent() {
        let a = vec![("b.txt".to_string(), b"2".to_vec()), ("a.txt".to_string(), b"1".to_vec())];
        let b = vec![("a.txt".to_string(), b"1".to_vec()), ("b.txt".to_string(), b"2".to_vec())];
        assert_eq!(bundle_sha256(&a), bundle_sha256(&b));       // order-independent
        let c = vec![("a.txt".to_string(), b"X".to_vec()), ("b.txt".to_string(), b"2".to_vec())];
        assert_ne!(bundle_sha256(&a), bundle_sha256(&c));       // content-sensitive
    }

    #[test]
    fn run_soul_build_emits_bundle() {
        let dir = std::env::temp_dir().join(format!("kibble_soulb_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir); std::fs::create_dir_all(&dir).unwrap();
        std::fs::write(dir.join("base.jinja"), "{% set soul = \"{{SOUL_DEFAULT_SYSTEM}}\" %}body").unwrap();
        std::fs::write(dir.join("soul.toml"), r#"
name = "s"
version = "0.1.0"
[persona]
system_prompt = "You are Kibble.\nBe direct."
[sampling]
temperature = 0.6
[[tool]]
name = "web_search"
description = "Search."
[template]
base = "base.jinja"
"#).unwrap();
        let out = dir.join("dist");
        let m = run_soul_build(&dir.join("soul.toml"), &out).unwrap();
        assert_eq!(m.default_prompt, "You are Kibble.\nBe direct.");
        assert!(!m.sha256.is_empty());
        for f in ["generation_config.json", "system_prompt.txt", "tools.json", "chat_template.jinja", "manifest.json"] {
            assert!(out.join(f).is_file(), "missing {f}");
        }
        // the multiline prompt was JSON-escaped into the template (newline → \n, no raw newline in the set line)
        let tmpl = std::fs::read_to_string(out.join("chat_template.jinja")).unwrap();
        assert!(tmpl.contains("You are Kibble.\\nBe direct."));
        assert!(!tmpl.contains("{{SOUL_DEFAULT_SYSTEM}}"));

        // no [template] → no chat_template.jinja
        std::fs::write(dir.join("soul2.toml"), "name=\"s\"\nversion=\"1\"\n[persona]\nsystem_prompt=\"P\"\n").unwrap();
        let out2 = dir.join("dist2");
        run_soul_build(&dir.join("soul2.toml"), &out2).unwrap();
        assert!(!out2.join("chat_template.jinja").exists());
    }
}