kodegraf-cli 0.1.2

use std::process::Command;
use std::time::Instant;

use anyhow::Result;

/// Run an A/B token benchmark: same question with and without Kodegraf.
/// Uses `claude -p` (print mode) with `--output-format json` to capture token counts.
pub fn run(question: &str) -> Result<()> {
    let repo_root = kodegraf_core::config::find_repo_root()?;

    println!("# Kodegraf Token Benchmark\n");
    println!("Question: {}\n", question);

    // ── Run WITHOUT Kodegraf ──
    println!("## Without Kodegraf\n");
    let mcp_path = repo_root.join(".mcp.json");
    let mcp_bak = repo_root.join(".mcp.json.kodegraf-benchmark-bak");

    // Temporarily disable MCP
    if mcp_path.exists() {
        std::fs::rename(&mcp_path, &mcp_bak)?;
    }

    let without = run_claude_query(question, &repo_root);

    // Restore MCP
    if mcp_bak.exists() {
        std::fs::rename(&mcp_bak, &mcp_path)?;
    }

    // ── Run WITH Kodegraf ──
    println!("## With Kodegraf\n");
    let with = run_claude_query(question, &repo_root);

    // ── Compare ──
    println!("\n## Comparison\n");
    match (without, with) {
        (Some(w), Some(k)) => {
            let ratio = if k.total > 0 { w.total as f64 / k.total as f64 } else { 1.0 };
            let saved = w.total.saturating_sub(k.total);
            let pct = if w.total > 0 { saved as f64 / w.total as f64 * 100.0 } else { 0.0 };

            println!("| Metric | Without Kodegraf | With Kodegraf | Savings |");
            println!("|--------|-----------------|---------------|---------|");
            println!("| Input tokens | {} | {} | {} ({:.0}%) |",
                w.input, k.input, w.input.saturating_sub(k.input),
                if w.input > 0 { (w.input - k.input) as f64 / w.input as f64 * 100.0 } else { 0.0 });
            println!("| Output tokens | {} | {} | {} |", w.output, k.output, w.output.saturating_sub(k.output));
            println!("| **Total tokens** | **{}** | **{}** | **{} ({:.0}%)** |", w.total, k.total, saved, pct);
            println!("| **Reduction ratio** | — | — | **{:.1}x** |", ratio);
            println!("| Time | {:.1}s | {:.1}s | |", w.duration_secs, k.duration_secs);
            println!("| Tool calls | {} | {} | |", w.tool_calls, k.tool_calls);
        }
        _ => {
            println!("Could not collect token data from both runs.");
            println!("Make sure `claude` CLI is installed and supports `--output-format json`.");
            println!();
            println!("Manual test: run these commands and compare token counts shown at the end:");
            println!();
            println!("  # Without Kodegraf:");
            println!("  mv .mcp.json .mcp.json.bak");
            println!("  claude -p \"{}\"", question);
            println!("  mv .mcp.json.bak .mcp.json");
            println!();
            println!("  # With Kodegraf:");
            println!("  claude -p \"{}\"", question);
        }
    }

    Ok(())
}

struct RunResult {
    input: usize,
    output: usize,
    total: usize,
    tool_calls: usize,
    duration_secs: f64,
}

fn run_claude_query(question: &str, cwd: &std::path::Path) -> Option<RunResult> {
    let start = Instant::now();

    eprintln!("  Running claude -p \"{}\" ...", &question[..question.len().min(50)]);

    let output = Command::new("claude")
        .args(["-p", question, "--output-format", "json", "--max-turns", "5"])
        .current_dir(cwd)
        .output()
        .ok()?;

    let duration = start.elapsed().as_secs_f64();

    if !output.status.success() {
        eprintln!("  claude exited with status: {}", output.status);
        let stderr = String::from_utf8_lossy(&output.stderr);
        if !stderr.is_empty() {
            eprintln!("  stderr: {}", &stderr[..stderr.len().min(200)]);
        }
        return None;
    }

    let stdout = String::from_utf8_lossy(&output.stdout);

    // Parse JSON output — claude --output-format json returns session data
    let data: serde_json::Value = serde_json::from_str(&stdout).ok()?;

    // Extract token counts from various possible locations in the JSON
    let input = extract_token_count(&data, "input_tokens");
    let output_tokens = extract_token_count(&data, "output_tokens");
    let total = input + output_tokens;

    // Count tool calls
    let tool_calls = count_tool_calls(&stdout);

    println!("  Input tokens:  {}", input);
    println!("  Output tokens: {}", output_tokens);
    println!("  Total tokens:  {}", total);
    println!("  Tool calls:    {}", tool_calls);
    println!("  Duration:      {:.1}s", duration);
    println!();

    Some(RunResult {
        input,
        output: output_tokens,
        total,
        tool_calls,
        duration_secs: duration,
    })
}

fn extract_token_count(data: &serde_json::Value, key: &str) -> usize {
    // Try direct field
    if let Some(v) = data.get(key).and_then(|v| v.as_u64()) {
        return v as usize;
    }
    // Try under "usage"
    if let Some(v) = data.get("usage").and_then(|u| u.get(key)).and_then(|v| v.as_u64()) {
        return v as usize;
    }
    // Try under "result" -> "usage"
    if let Some(v) = data.get("result").and_then(|r| r.get("usage")).and_then(|u| u.get(key)).and_then(|v| v.as_u64()) {
        return v as usize;
    }
    // Estimate from text length if no token data
    0
}

fn count_tool_calls(json_str: &str) -> usize {
    // Count occurrences of tool_use in the output
    json_str.matches("tool_use").count()
        + json_str.matches("\"type\":\"tool\"").count()
        + json_str.matches("kodegraf_").count()
}