//! Budget computation: estimates USD spend from agent log files.

use std::path::Path;

/// Compute in-flight budget from per-message usage data in agent log files.
///
/// Sums cost across **all** sessions for the given story+agent combination.
/// For completed sessions (containing a `result` event), uses the accurate
/// `total_cost_usd` from Claude Code.  For running sessions (no `result`
/// event yet), estimates cost from per-turn token counts in `assistant`
/// events.
///
/// This cumulative total is useful for cost analysis (metric tool) but is
/// NOT used by the watchdog for enforcement — the watchdog uses
/// [`compute_budget_from_single_log`] on the current session only.
pub(crate) fn compute_budget_from_logs(
    project_root: &Path,
    story_id: &str,
    agent_name: &str,
) -> f64 {
    let log_files =
        crate::agent_log::list_story_log_files(project_root, story_id, Some(agent_name));
    let mut total_cost = 0.0;

    for path in &log_files {
        total_cost += compute_budget_from_single_log(path);
    }

    total_cost
}

/// Compute budget from a single log file.
///
/// If the log contains a `result` event with `total_cost_usd` (completed
/// session), returns that.  Otherwise estimates cost from per-turn token
/// usage in `assistant` events (running session).
pub(crate) fn compute_budget_from_single_log(path: &Path) -> f64 {
    use crate::agents::TokenUsage;

    let entries = match crate::agent_log::read_log(path) {
        Ok(e) => e,
        Err(_) => return 0.0,
    };

    // Check for a result event with accurate total_cost_usd (completed session).
    let result_cost = entries.iter().rev().find_map(|entry| {
        if entry.event.get("type").and_then(|v| v.as_str()) == Some("agent_json")
            && let Some(data) = entry.event.get("data")
            && data.get("type").and_then(|v| v.as_str()) == Some("result")
        {
            data.get("total_cost_usd").and_then(|v| v.as_f64())
        } else {
            None
        }
    });

    if let Some(cost) = result_cost {
        return cost;
    }

    // Running session — estimate from per-message token usage.
    let mut cost = 0.0;
    for entry in &entries {
        if entry.event.get("type").and_then(|v| v.as_str()) == Some("agent_json")
            && let Some(data) = entry.event.get("data")
            && data.get("type").and_then(|v| v.as_str()) == Some("assistant")
            && let Some(message) = data.get("message")
            && let Some(usage) = message.get("usage")
        {
            let model = message
                .get("model")
                .and_then(|v| v.as_str())
                .map(crate::agents::AgentModel::from_api_str);
            let token_usage = TokenUsage {
                input_tokens: usage
                    .get("input_tokens")
                    .and_then(|v| v.as_u64())
                    .unwrap_or(0),
                output_tokens: usage
                    .get("output_tokens")
                    .and_then(|v| v.as_u64())
                    .unwrap_or(0),
                cache_creation_input_tokens: usage
                    .get("cache_creation_input_tokens")
                    .and_then(|v| v.as_u64())
                    .unwrap_or(0),
                cache_read_input_tokens: usage
                    .get("cache_read_input_tokens")
                    .and_then(|v| v.as_u64())
                    .unwrap_or(0),
                total_cost_usd: 0.0,
            };
            cost += token_usage.estimate_cost_usd(model.as_ref());
        }
    }
    cost
}
