yog 0.0.1

//! Whole-tree budget-spend fold (DESIGN §5.1 #16; ARCH §6 budgets).
//!
//! Yog is a pure reader (§3.5). Budget *spent* is derived, never stored:
//! sum every brazen `Usage` event's token counters across **every attempt
//! segment of every step** under a root agent and its entire hyphenated
//! descent — `steps/<root>/` plus every `steps/<root>-*/` (the one shared
//! `steps/` subtree; budgets are whole-tree consumables, ARCH §6, §2.2).
//! This mirrors lernie's own `budget::spend` derivation exactly — a failed
//! or superseded attempt still burned tokens and real money (§4.4, §6) —
//! so the figure yog shows is the figure that exhausts `max_total_tokens`.
//!
//! Forgiving by construction: a missing `steps/` tree, an unreadable step
//! dir, a missing `response.json`, or a malformed / forward-compat event
//! line each contributes zero — the fold never panics on a partial or
//! mid-stream tree.

use std::fs;
use std::path::Path;

mod render;
pub use render::render;

/// Conv-repo subdir of per-conversation step records (ARCH §2.2).
const STEPS_DIR: &str = "steps";
/// Per-step JSONL of `v=1` events (ARCH §2.3, §4.4).
const RESPONSE_FILE: &str = "response.json";
/// Zero-padded step-sequence width (`001`, `002`, …) per ARCH §2.3.
const STEP_SEQ_WIDTH: usize = 3;

/// Tokens spent by a root agent and its descent, split by the four brazen
/// `Usage` counters. The wire shape is
/// `{"type":"usage","input_tokens":N,"output_tokens":M,
/// "cache_read_tokens":R,"cache_write_tokens":W}` (brazen
/// `canonical::event`) — every counter nullable, a `null`/absent field
/// counting **zero, never fabricated**. [`total_tokens`](BudgetSpend::total_tokens)
/// is what exhausts `max_total_tokens` (ARCH §6: all four summed).
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct BudgetSpend {
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cache_read_tokens: u64,
    pub cache_write_tokens: u64,
}

impl BudgetSpend {
    /// Total tokens billed against the tree's `max_total_tokens` ceiling —
    /// all four counters summed (ARCH §6).
    pub fn total_tokens(&self) -> u64 {
        self.input_tokens + self.output_tokens + self.cache_read_tokens + self.cache_write_tokens
    }

    fn add(&mut self, other: BudgetSpend) {
        self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
        self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
        self.cache_read_tokens = self
            .cache_read_tokens
            .saturating_add(other.cache_read_tokens);
        self.cache_write_tokens = self
            .cache_write_tokens
            .saturating_add(other.cache_write_tokens);
    }
}

/// Fold whole-tree spend for `root_id` under `workspace`: every conv-id
/// dir equal to `root_id` or prefixed `root_id-` (its hyphenated descent,
/// ARCH §2.2). A missing `steps/` tree is zero spend.
pub fn spend(workspace: &Path, root_id: &str) -> BudgetSpend {
    let Ok(entries) = fs::read_dir(workspace.join(STEPS_DIR)) else {
        return BudgetSpend::default();
    };
    let prefix_dash = format!("{root_id}-");
    let mut total = BudgetSpend::default();
    for entry in entries.flatten() {
        let raw = entry.file_name();
        let name = raw.to_string_lossy();
        if name != root_id && !name.starts_with(&prefix_dash) {
            continue;
        }
        total.add(conv_spend(&entry.path()));
    }
    total
}

/// Sum the spend of every 3-digit step subdir of one conv-id dir. A
/// conv-id entry that is not a readable directory contributes zero.
fn conv_spend(conv_dir: &Path) -> BudgetSpend {
    let Ok(entries) = fs::read_dir(conv_dir) else {
        return BudgetSpend::default();
    };
    let mut total = BudgetSpend::default();
    for entry in entries.flatten() {
        let raw = entry.file_name();
        let name = raw.to_string_lossy();
        if name.len() == STEP_SEQ_WIDTH && name.bytes().all(|b| b.is_ascii_digit()) {
            total.add(step_spend(&entry.path()));
        }
    }
    total
}

/// Fold every `Usage` line of one step's `response.json` (every segment,
/// ARCH §6). A missing file contributes zero; a malformed or non-`Usage`
/// line is skipped.
fn step_spend(step_dir: &Path) -> BudgetSpend {
    let Ok(bytes) = fs::read(step_dir.join(RESPONSE_FILE)) else {
        return BudgetSpend::default();
    };
    spend_from_bytes(&bytes)
}

/// Fold every `Usage` line of a `response.json` payload — every attempt
/// segment (ARCH §6). Pure over the bytes: the per-step unit the Y13 steps
/// inspector reuses so the brazen Usage-event vocabulary lives in exactly
/// one place (single source of truth). A malformed or non-`Usage` line
/// contributes zero.
pub fn spend_from_bytes(bytes: &[u8]) -> BudgetSpend {
    let mut total = BudgetSpend::default();
    for line in bytes.split(|b| *b == b'\n') {
        if let Some(spend) = usage_line(line) {
            total.add(spend);
        }
    }
    total
}

/// The counters of one JSONL event line iff it is a `{"type":"usage",…}`
/// event (brazen's internally-tagged `Event::Usage`). A non-`usage` type,
/// a line that does not parse, or one with no string `type` yields `None`.
fn usage_line(line: &[u8]) -> Option<BudgetSpend> {
    let value: serde_json::Value = serde_json::from_slice(line).ok()?;
    if value.get("type")?.as_str()? != "usage" {
        return None;
    }
    Some(BudgetSpend {
        input_tokens: counter(&value, "input_tokens"),
        output_tokens: counter(&value, "output_tokens"),
        cache_read_tokens: counter(&value, "cache_read_tokens"),
        cache_write_tokens: counter(&value, "cache_write_tokens"),
    })
}

/// One `Usage` counter as `u64`: a `null`, absent, or non-integer field is
/// zero — brazen's Usage-Option contract (a counter a provider never
/// reported is unknown, rendered zero, never fabricated).
fn counter(value: &serde_json::Value, key: &str) -> u64 {
    value
        .get(key)
        .and_then(serde_json::Value::as_u64)
        .unwrap_or(0)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;
    use tempfile::tempdir;

    fn write_step(workspace: &Path, conv: &str, seq: &str, lines: &[&str]) {
        let step = workspace.join(STEPS_DIR).join(conv).join(seq);
        std::fs::create_dir_all(&step).unwrap();
        std::fs::write(step.join(RESPONSE_FILE), lines.join("\n")).unwrap();
    }

    fn steps_dir(workspace: &Path) -> PathBuf {
        workspace.join(STEPS_DIR)
    }

    #[test]
    fn folds_whole_tree_across_segments_and_descent() {
        let dir = tempdir().unwrap();
        let root = "20260717T120000Z-root";
        let child = "20260717T120000Z-root-20260717T120100Z-kid0";
        // Root step 001: two Usage segments (a billed retry, §6), plus
        // non-Usage / malformed lines that must contribute nothing.
        write_step(
            dir.path(),
            root,
            "001",
            &[
                r#"{"type":"usage","input_tokens":10,"output_tokens":5,"cache_read_tokens":2,"cache_write_tokens":1}"#,
                r#"{"type":"content_delta","index":0,"delta":{"text_delta":"hi"}}"#,
                r#"{"type":"usage","input_tokens":3,"output_tokens":null}"#,
                "not json",
                r#"{"no_type":1}"#,
                r#"{"type":123}"#,
            ],
        );
        // A step dir with no response.json (read Err → zero), and a
        // non-step entry (skipped by the 3-digit filter).
        std::fs::create_dir_all(steps_dir(dir.path()).join(root).join("002")).unwrap();
        std::fs::create_dir_all(steps_dir(dir.path()).join(root).join("tools")).unwrap();
        // Descended child conv (steps/<root>-*/), counted whole-tree.
        write_step(
            dir.path(),
            child,
            "001",
            &[
                r#"{"type":"usage","input_tokens":100,"output_tokens":50,"cache_read_tokens":4,"cache_write_tokens":8}"#,
            ],
        );
        // A sibling root that must NOT count toward this tree.
        write_step(
            dir.path(),
            "20260717T120000Z-othr",
            "001",
            &[r#"{"type":"usage","input_tokens":999,"output_tokens":999}"#],
        );
        // A conv-id entry matching the prefix but not a directory: its
        // read fails and contributes zero.
        std::fs::write(steps_dir(dir.path()).join(format!("{root}-stray")), b"x").unwrap();

        let s = spend(dir.path(), root);
        // Root's two segments (the second's `null` output and absent cache
        // counters each count zero) plus the child, whole-tree.
        assert_eq!(s.input_tokens, 10 + 3 + 100);
        assert_eq!(s.output_tokens, 5 + 50);
        assert_eq!(s.cache_read_tokens, 2 + 4);
        assert_eq!(s.cache_write_tokens, 1 + 8);
        assert_eq!(s.total_tokens(), 113 + 55 + 6 + 9);
    }

    #[test]
    fn missing_steps_tree_is_zero_spend() {
        let dir = tempdir().unwrap();
        assert_eq!(
            spend(dir.path(), "20260717T120000Z-root"),
            BudgetSpend::default()
        );
    }
}