//! Task command — one-shot implementation context for a task description.

use anyhow::Result;
use colored::Colorize;

use cqs::{task, task_to_json, Embedder};

/// Waterfall budget weight for the scout section (file groups, chunk roles).
const WATERFALL_SCOUT: f64 = 0.15;
/// Waterfall budget weight for the code section (gathered chunks with content).
const WATERFALL_CODE: f64 = 0.50;
/// Waterfall budget weight for the impact section (risk scores + tests).
const WATERFALL_IMPACT: f64 = 0.15;
/// Waterfall budget weight for the placement section (where to add).
const WATERFALL_PLACEMENT: f64 = 0.10;
// Notes section takes whatever budget remains (no explicit constant needed).

pub(crate) fn cmd_task(
    cli: &crate::cli::Cli,
    description: &str,
    limit: usize,
    json: bool,
    max_tokens: Option<usize>,
) -> Result<()> {
    let _span = tracing::info_span!("cmd_task", ?max_tokens).entered();
    let (store, root, _) = crate::cli::open_project_store_readonly()?;
    let embedder = Embedder::new(cli.model_config().clone())?;
    let limit = limit.clamp(1, 10);

    let result = task(&store, &embedder, description, &root, limit)?;

    if let Some(budget) = max_tokens {
        output_with_budget(&result, &root, &embedder, budget, json)?;
    } else if json {
        let output = task_to_json(&result);
        println!("{}", serde_json::to_string_pretty(&output)?);
    } else {
        output_text(&result, &root);
    }

    Ok(())
}

/// Waterfall token budgeting: allocate budget across sections, surplus flows forward.
fn output_with_budget(
    result: &cqs::TaskResult,
    root: &std::path::Path,
    embedder: &Embedder,
    budget: usize,
    json: bool,
) -> Result<()> {
    let overhead = if json {
        super::JSON_OVERHEAD_PER_RESULT
    } else {
        0
    };
    let packed = waterfall_pack(result, embedder, budget, overhead);

    if json {
        output_json_budgeted(result, &packed)?;
    } else {
        output_text_budgeted(result, root, &packed);
    }

    Ok(())
}

/// Packed section indices from waterfall budgeting.
pub(crate) struct PackedSections {
    scout: Vec<usize>,
    code: Vec<usize>,
    risk: Vec<usize>,
    tests: Vec<usize>,
    placement: Vec<usize>,
    notes: Vec<usize>,
    pub(crate) total_used: usize,
    pub(crate) budget: usize,
}

/// Pack a section: count tokens for texts, run index_pack, return (indices, used).
///
/// Extracts the repeated pattern of count_tokens_batch + index_pack used by
/// each waterfall section. CQ-26.
fn pack_section(
    embedder: &Embedder,
    texts: &[&str],
    section_budget: usize,
    overhead: usize,
    score_fn: impl Fn(usize) -> f32,
) -> (Vec<usize>, usize) {
    let counts = super::count_tokens_batch(embedder, texts);
    super::index_pack(&counts, section_budget, overhead, score_fn)
}

/// Compute waterfall token budgeting across all task sections.
///
/// Shared between CLI `cqs task --tokens` and batch `task --tokens`.
/// `overhead_per_item` should be `JSON_OVERHEAD_PER_RESULT` for JSON, 0 for text.
pub(crate) fn waterfall_pack(
    result: &cqs::TaskResult,
    embedder: &Embedder,
    budget: usize,
    overhead_per_item: usize,
) -> PackedSections {
    let _span = tracing::info_span!("waterfall_budget", budget).entered();
    let mut remaining = budget;

    // 1. Scout section — pack file groups by relevance
    let scout_budget = ((budget as f64 * WATERFALL_SCOUT) as usize).min(remaining);
    let group_texts: Vec<String> = result
        .scout
        .file_groups
        .iter()
        .map(|g| {
            g.chunks
                .iter()
                .map(|c| c.signature.as_str())
                .collect::<Vec<_>>()
                .join("\n")
        })
        .collect();
    let group_refs: Vec<&str> = group_texts.iter().map(|s| s.as_str()).collect();
    let (scout_indices, scout_used) = pack_section(
        embedder,
        &group_refs,
        scout_budget,
        overhead_per_item,
        |i| result.scout.file_groups[i].relevance_score,
    );
    // Charge actual usage to remaining — overshoot from first-item guarantee
    // must reduce downstream budgets to prevent total overshoot
    remaining = remaining.saturating_sub(scout_used);

    // 2. Code section (+ surplus) — pack gathered chunks by score
    let code_budget = ((budget as f64 * WATERFALL_CODE) as usize
        + scout_budget.saturating_sub(scout_used))
    .min(remaining);
    let code_refs: Vec<&str> = result.code.iter().map(|c| c.content.as_str()).collect();
    let (code_indices, code_used) =
        pack_section(embedder, &code_refs, code_budget, overhead_per_item, |i| {
            result.code[i].score
        });
    remaining = remaining.saturating_sub(code_used);

    // 3. Impact section (+ surplus) — risk by score, tests by depth
    let impact_budget = ((budget as f64 * WATERFALL_IMPACT) as usize
        + code_budget.saturating_sub(code_used))
    .min(remaining);
    let risk_texts: Vec<String> = result
        .risk
        .iter()
        .map(|fr| {
            format!(
                "{}: {:?} score:{:.1} callers:{} cov:{:.0}%",
                fr.name,
                fr.risk.risk_level,
                fr.risk.score,
                fr.risk.caller_count,
                fr.risk.test_ratio * 100.0
            )
        })
        .collect();
    let risk_refs: Vec<&str> = risk_texts.iter().map(|s| s.as_str()).collect();
    let (risk_indices, risk_used) = pack_section(
        embedder,
        &risk_refs,
        impact_budget,
        overhead_per_item,
        |i| result.risk[i].risk.score,
    );

    let tests_budget = impact_budget.saturating_sub(risk_used);
    let test_texts: Vec<String> = result
        .tests
        .iter()
        .map(|t| {
            format!(
                "{} {}:{} depth:{}",
                t.name,
                t.file.display(),
                t.line,
                t.call_depth
            )
        })
        .collect();
    let test_refs: Vec<&str> = test_texts.iter().map(|s| s.as_str()).collect();
    let (test_indices, tests_used) =
        pack_section(embedder, &test_refs, tests_budget, overhead_per_item, |i| {
            1.0 / (result.tests[i].call_depth as f32 + 1.0)
        });
    remaining = remaining.saturating_sub(risk_used + tests_used);

    // 4. Placement section (+ surplus)
    let placement_budget = ((budget as f64 * WATERFALL_PLACEMENT) as usize
        + impact_budget.saturating_sub(risk_used + tests_used))
    .min(remaining);
    let placement_texts: Vec<String> = result
        .placement
        .iter()
        .map(|s| {
            format!(
                "{}: {} line:{} near:{}",
                s.file.display(),
                s.reason,
                s.insertion_line,
                s.near_function
            )
        })
        .collect();
    let placement_refs: Vec<&str> = placement_texts.iter().map(|s| s.as_str()).collect();
    let (placement_indices, placement_used) = pack_section(
        embedder,
        &placement_refs,
        placement_budget,
        overhead_per_item,
        |i| result.placement[i].score,
    );
    remaining = remaining.saturating_sub(placement_used);

    // 5. Notes section — takes whatever budget remains
    let note_refs: Vec<&str> = result
        .scout
        .relevant_notes
        .iter()
        .map(|n| n.text.as_str())
        .collect();
    let (note_indices, notes_used) =
        pack_section(embedder, &note_refs, remaining, overhead_per_item, |i| {
            result.scout.relevant_notes[i].sentiment.abs()
        });

    let total_used = scout_used + code_used + risk_used + tests_used + placement_used + notes_used;

    tracing::info!(
        total = total_used,
        budget,
        scout = scout_used,
        code = code_used,
        risk = risk_used,
        tests = tests_used,
        placement = placement_used,
        notes = notes_used,
        "Waterfall budget complete"
    );

    PackedSections {
        scout: scout_indices,
        code: code_indices,
        risk: risk_indices,
        tests: test_indices,
        placement: placement_indices,
        notes: note_indices,
        total_used,
        budget,
    }
}

/// Build budgeted JSON for a task result using full waterfall token budgeting.
///
/// Shared between CLI `cqs task --tokens --json` and batch `task --tokens`.
pub(crate) fn task_to_budgeted_json(
    result: &cqs::TaskResult,
    embedder: &Embedder,
    budget: usize,
) -> serde_json::Value {
    let packed = waterfall_pack(result, embedder, budget, super::JSON_OVERHEAD_PER_RESULT);
    budgeted_json(result, &packed)
}

/// Build budgeted JSON combining all packed sections into a single output.
fn budgeted_json(result: &cqs::TaskResult, packed: &PackedSections) -> serde_json::Value {
    let mut scout_json = build_scout_json(result, &packed.scout);
    let code_json = build_code_json(result, &packed.code);
    let risk_json = build_risk_json(result, &packed.risk);
    let tests_json = build_tests_json(result, &packed.tests);
    let placement_json = build_placement_json(result, &packed.placement);
    let notes_json = build_notes_json(result, &packed.notes);
    scout_json["relevant_notes"] = serde_json::json!(notes_json);

    serde_json::json!({
        "description": result.description,
        "scout": scout_json,
        "code": code_json,
        "risk": risk_json,
        "tests": tests_json,
        "placement": placement_json,
        "summary": {
            "total_files": result.summary.total_files,
            "total_functions": result.summary.total_functions,
            "modify_targets": result.summary.modify_targets,
            "high_risk_count": result.summary.high_risk_count,
            "test_count": result.summary.test_count,
            "stale_count": result.summary.stale_count,
        },
        "token_count": packed.total_used,
        "token_budget": packed.budget,
    })
}

fn output_json_budgeted(result: &cqs::TaskResult, packed: &PackedSections) -> Result<()> {
    let output = budgeted_json(result, packed);
    println!("{}", serde_json::to_string_pretty(&output)?);
    Ok(())
}

fn build_scout_json(result: &cqs::TaskResult, indices: &[usize]) -> serde_json::Value {
    let groups: Vec<serde_json::Value> = indices
        .iter()
        .map(|&i| {
            let g = &result.scout.file_groups[i];
            let chunks: Vec<serde_json::Value> = g
                .chunks
                .iter()
                .map(|c| {
                    serde_json::json!({
                        "name": c.name,
                        "chunk_type": c.chunk_type.to_string(),
                        "signature": c.signature,
                        "line_start": c.line_start,
                        "role": c.role.as_str(),
                        "caller_count": c.caller_count,
                        "test_count": c.test_count,
                        "search_score": c.search_score,
                    })
                })
                .collect();
            serde_json::json!({
                "file": cqs::normalize_path(&g.file),
                "relevance_score": g.relevance_score,
                "is_stale": g.is_stale,
                "chunks": chunks,
            })
        })
        .collect();
    serde_json::json!({
        "file_groups": groups,
        "summary": {
            "total_files": result.scout.summary.total_files,
            "total_functions": result.scout.summary.total_functions,
            "untested_count": result.scout.summary.untested_count,
            "stale_count": result.scout.summary.stale_count,
        }
    })
}

fn build_code_json(result: &cqs::TaskResult, indices: &[usize]) -> Vec<serde_json::Value> {
    indices
        .iter()
        .filter_map(|&i| match serde_json::to_value(&result.code[i]) {
            Ok(v) => Some(v),
            Err(e) => {
                tracing::warn!(error = %e, chunk = %result.code[i].name, "Failed to serialize chunk");
                None
            }
        })
        .collect()
}

fn build_risk_json(result: &cqs::TaskResult, indices: &[usize]) -> Vec<serde_json::Value> {
    indices
        .iter()
        .filter_map(|&i| serde_json::to_value(&result.risk[i]).ok())
        .collect()
}

/// Converts a subset of test results to JSON format.
///
/// Paths in tests are already relative to the project root.
fn build_tests_json(result: &cqs::TaskResult, indices: &[usize]) -> Vec<serde_json::Value> {
    indices
        .iter()
        .filter_map(|&i| serde_json::to_value(&result.tests[i]).ok())
        .collect()
}

/// Builds a JSON representation of task result placements for specified indices.
///
/// Paths in placement are already relative to the project root.
fn build_placement_json(result: &cqs::TaskResult, indices: &[usize]) -> Vec<serde_json::Value> {
    indices
        .iter()
        .filter_map(|&i| serde_json::to_value(&result.placement[i]).ok())
        .collect()
}

fn build_notes_json(result: &cqs::TaskResult, indices: &[usize]) -> Vec<serde_json::Value> {
    indices
        .iter()
        .map(|&i| {
            let n = &result.scout.relevant_notes[i];
            serde_json::json!({
                "text": n.text,
                "sentiment": n.sentiment,
                "mentions": n.mentions,
            })
        })
        .collect()
}

fn output_text_budgeted(result: &cqs::TaskResult, root: &std::path::Path, packed: &PackedSections) {
    print_header(
        &result.description,
        &result.summary,
        packed.total_used,
        packed.budget,
    );
    print_scout_section(result, root, &packed.scout);
    print_code_section_idx(&result.code, root, &packed.code, result.code.len());
    print_impact_section_idx(&result.risk, &result.tests, &packed.risk, &packed.tests);
    print_placement_section_idx(
        &result.placement,
        root,
        &packed.placement,
        result.placement.len(),
    );
    print_notes_section_idx(
        &result.scout.relevant_notes,
        &packed.notes,
        result.scout.relevant_notes.len(),
    );
}

fn output_text(result: &cqs::TaskResult, root: &std::path::Path) {
    let all_scout: Vec<usize> = (0..result.scout.file_groups.len()).collect();
    print_header(&result.description, &result.summary, 0, 0);
    print_scout_section(result, root, &all_scout);

    let all_code: Vec<usize> = (0..result.code.len()).collect();
    print_code_section_idx(&result.code, root, &all_code, result.code.len());

    let all_risk: Vec<usize> = (0..result.risk.len()).collect();
    let all_tests: Vec<usize> = (0..result.tests.len()).collect();
    print_impact_section_idx(&result.risk, &result.tests, &all_risk, &all_tests);

    let all_placement: Vec<usize> = (0..result.placement.len()).collect();
    print_placement_section_idx(
        &result.placement,
        root,
        &all_placement,
        result.placement.len(),
    );

    let all_notes: Vec<usize> = (0..result.scout.relevant_notes.len()).collect();
    print_notes_section_idx(
        &result.scout.relevant_notes,
        &all_notes,
        result.scout.relevant_notes.len(),
    );
}

fn print_header(description: &str, summary: &cqs::TaskSummary, used: usize, budget: usize) {
    let token_label = if budget > 0 {
        format!(" ({} of {} tokens)", used, budget)
    } else {
        String::new()
    };
    println!(
        "{} {}{}",
        "═══ Task:".cyan().bold(),
        description.bold(),
        token_label.dimmed()
    );
    println!(
        "  {} targets | {} files | {} tests | {} high-risk",
        summary.modify_targets.to_string().bold(),
        summary.total_files,
        summary.test_count,
        summary.high_risk_count
    );
}

fn print_scout_section(result: &cqs::TaskResult, root: &std::path::Path, indices: &[usize]) {
    if indices.is_empty() {
        return;
    }
    println!();
    println!("{}", "── Scout ──────────────────────────────".cyan());
    let total = result.scout.file_groups.len();
    for &i in indices {
        let g = &result.scout.file_groups[i];
        let rel = cqs::rel_display(&g.file, root);
        print!(
            "  {} {}",
            rel.bold(),
            format!("({:.2})", g.relevance_score).dimmed()
        );
        if g.is_stale {
            print!(" {}", "[STALE]".yellow().bold());
        }
        println!();
        for c in &g.chunks {
            let role = match c.role {
                cqs::ChunkRole::ModifyTarget => "modify",
                cqs::ChunkRole::TestToUpdate => "test",
                cqs::ChunkRole::Dependency => "dep",
            };
            println!(
                "    {} {} {} {}",
                "▸".dimmed(),
                c.name,
                format!("({})", role).dimmed(),
                format!("callers:{} tests:{}", c.caller_count, c.test_count).dimmed()
            );
        }
    }
    if indices.len() < total {
        println!(
            "  {}",
            format!("({} more files truncated)", total - indices.len()).dimmed()
        );
    }
}

fn print_code_section_idx(
    code: &[cqs::GatheredChunk],
    root: &std::path::Path,
    indices: &[usize],
    total: usize,
) {
    if indices.is_empty() {
        return;
    }
    println!();
    println!("{}", "── Code ───────────────────────────────".cyan());
    for &i in indices {
        let c = &code[i];
        let rel = cqs::rel_display(&c.file, root);
        println!("  {} {}:{}", c.name.bold(), rel, c.line_start);
        if !c.signature.is_empty() {
            println!("    {}", c.signature.dimmed());
        }
        let mut line_count = 0;
        for line in c.content.lines().take(5) {
            println!("    {}", line);
            line_count += 1;
        }
        if line_count == 5 && c.content.lines().nth(5).is_some() {
            println!("    {}", "...".dimmed());
        }
    }
    if indices.len() < total {
        println!(
            "  {}",
            format!("({} more items truncated)", total - indices.len()).dimmed()
        );
    }
}

fn print_impact_section_idx(
    risk: &[cqs::FunctionRisk],
    tests: &[cqs::TestInfo],
    risk_idx: &[usize],
    test_idx: &[usize],
) {
    if risk_idx.is_empty() && test_idx.is_empty() {
        return;
    }
    if !risk_idx.is_empty() {
        println!();
        println!("{}", "── Impact ─────────────────────────────".cyan());
        for &i in risk_idx {
            let fr = &risk[i];
            let level = match fr.risk.risk_level {
                cqs::RiskLevel::High => {
                    format!("{:?}", fr.risk.risk_level).red().bold().to_string()
                }
                cqs::RiskLevel::Medium => format!("{:?}", fr.risk.risk_level).yellow().to_string(),
                cqs::RiskLevel::Low => format!("{:?}", fr.risk.risk_level).green().to_string(),
            };
            println!(
                "  {}: {} {}",
                fr.name,
                level,
                format!(
                    "(score: {:.1}, callers: {}, test_ratio: {:.0}%)",
                    fr.risk.score,
                    fr.risk.caller_count,
                    fr.risk.test_ratio * 100.0
                )
                .dimmed()
            );
        }
        if risk_idx.len() < risk.len() {
            println!(
                "  {}",
                format!(
                    "({} more risk entries truncated)",
                    risk.len() - risk_idx.len()
                )
                .dimmed()
            );
        }
    }

    if !test_idx.is_empty() {
        println!();
        println!("{}", "── Tests ──────────────────────────────".cyan());
        for &i in test_idx {
            let t = &tests[i];
            let rel = cqs::rel_display(&t.file, std::path::Path::new(""));
            println!(
                "  {} {}:{} {}",
                t.name,
                rel,
                t.line,
                format!("depth:{}", t.call_depth).dimmed()
            );
        }
        if test_idx.len() < tests.len() {
            println!(
                "  {}",
                format!("({} more tests truncated)", tests.len() - test_idx.len()).dimmed()
            );
        }
    }
}

fn print_placement_section_idx(
    placement: &[cqs::FileSuggestion],
    root: &std::path::Path,
    indices: &[usize],
    total: usize,
) {
    if indices.is_empty() {
        return;
    }
    println!();
    println!("{}", "── Placement ──────────────────────────".cyan());
    for &i in indices {
        let s = &placement[i];
        let rel = cqs::rel_display(&s.file, root);
        println!("  {} — {}", rel.bold(), s.reason.dimmed());
    }
    if indices.len() < total {
        println!(
            "  {}",
            format!("({} more suggestions truncated)", total - indices.len()).dimmed()
        );
    }
}

fn print_notes_section_idx(notes: &[cqs::store::NoteSummary], indices: &[usize], total: usize) {
    if indices.is_empty() {
        return;
    }
    println!();
    println!("{}", "── Notes ──────────────────────────────".cyan());
    for &i in indices {
        let n = &notes[i];
        let sentiment = if n.sentiment < 0.0 {
            format!("[{:.1}]", n.sentiment).red().to_string()
        } else if n.sentiment > 0.0 {
            format!("[+{:.1}]", n.sentiment).green().to_string()
        } else {
            "[0.0]".dimmed().to_string()
        };
        let text = if n.text.len() > 80 {
            format!("{}...", &n.text[..n.text.floor_char_boundary(77)])
        } else {
            n.text.clone()
        };
        println!("  {} {}", sentiment, text.dimmed());
    }
    if indices.len() < total {
        println!(
            "  {}",
            format!("({} more notes truncated)", total - indices.len()).dimmed()
        );
    }
}

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

    #[test]
    fn test_waterfall_allocation_percentages() {
        // Notes takes the remainder, so the explicit weights must sum to ≤1.0
        let total = WATERFALL_SCOUT + WATERFALL_CODE + WATERFALL_IMPACT + WATERFALL_PLACEMENT;
        assert!(
            total <= 1.0 && total >= 0.9,
            "Explicit budget weights must leave a small remainder for notes, got {total}"
        );
    }

    #[test]
    fn test_waterfall_section_budgets() {
        let budget: usize = 1000;
        let scout = (budget as f64 * WATERFALL_SCOUT) as usize;
        let code = (budget as f64 * WATERFALL_CODE) as usize;
        let impact = (budget as f64 * WATERFALL_IMPACT) as usize;
        let placement = (budget as f64 * WATERFALL_PLACEMENT) as usize;
        let notes = budget - scout - code - impact - placement;
        assert_eq!(scout + code + impact + placement + notes, budget);
    }

    #[test]
    fn test_index_pack_empty() {
        let (indices, used) = index_pack(&[], 100, 0, |_| 1.0);
        assert!(indices.is_empty());
        assert_eq!(used, 0);
    }

    #[test]
    fn test_index_pack_all_fit() {
        let counts = vec![10, 20, 30];
        let (indices, used) = index_pack(&counts, 100, 0, |_| 1.0);
        assert_eq!(indices, vec![0, 1, 2]);
        assert_eq!(used, 60);
    }

    #[test]
    fn test_index_pack_budget_forces_selection() {
        let counts = vec![10, 10, 10, 10, 10];
        // Scores: 0=1.0, 1=5.0, 2=3.0, 3=4.0, 4=2.0
        // Budget 30 fits 3 items → picks indices 1, 3, 2 (by score), sorted → [1, 2, 3]
        let (indices, used) = index_pack(&counts, 30, 0, |i| match i {
            0 => 1.0,
            1 => 5.0,
            2 => 3.0,
            3 => 4.0,
            4 => 2.0,
            _ => 0.0,
        });
        assert_eq!(indices.len(), 3);
        assert_eq!(used, 30);
        assert!(indices.contains(&1));
        assert!(indices.contains(&2));
        assert!(indices.contains(&3));
    }

    #[test]
    fn test_index_pack_preserves_order() {
        let counts = vec![10, 10, 10];
        // Budget fits 2 → picks highest score items, returned in original order
        let (indices, _) = index_pack(&counts, 20, 0, |i| match i {
            0 => 1.0,
            1 => 3.0,
            2 => 2.0,
            _ => 0.0,
        });
        assert_eq!(indices, vec![1, 2]); // original order, not score order
    }

    #[test]
    fn test_index_pack_always_includes_one() {
        let counts = vec![100]; // over budget
        let (indices, used) = index_pack(&counts, 10, 0, |_| 1.0);
        assert_eq!(indices, vec![0]);
        assert_eq!(used, 100);
    }

    #[test]
    fn test_index_pack_with_overhead() {
        let counts = vec![10, 10, 10];
        // With overhead 35, each item costs 45. Budget 100 fits 2 (90), not 3 (135)
        let (indices, used) = index_pack(&counts, 100, 35, |_| 1.0);
        assert_eq!(indices.len(), 2);
        assert_eq!(used, 90);
    }

    // TC-8: index_pack with zero budget returns nothing
    #[test]
    fn test_index_pack_zero_budget() {
        let counts = vec![10, 20, 30];
        let (indices, used) = index_pack(&counts, 0, 0, |_| 1.0);
        assert!(indices.is_empty());
        assert_eq!(used, 0);
    }

    // TC-11: Waterfall surplus forwarding — verify unused budget flows to next section
    #[test]
    fn test_waterfall_surplus_forwarding() {
        let budget: usize = 1000;
        let weights = [
            WATERFALL_SCOUT,
            WATERFALL_CODE,
            WATERFALL_IMPACT,
            WATERFALL_PLACEMENT,
        ];
        let base_budgets: Vec<usize> = weights
            .iter()
            .map(|w| (budget as f64 * w) as usize)
            .collect();

        // Scenario: scout uses only 50 of its 150 budget → 100 surplus flows to code
        let scout_budget = base_budgets[0]; // 150
        let scout_used = 50;
        let code_budget_with_surplus =
            (base_budgets[1] + scout_budget.saturating_sub(scout_used)).min(budget - scout_used);
        // Code gets 500 base + 100 surplus = 600 (capped by remaining = 950)
        assert_eq!(code_budget_with_surplus, 600);

        // Scenario: code uses all 600 → 0 surplus to impact
        let code_used = 600;
        let impact_budget_with_surplus = (base_budgets[2]
            + code_budget_with_surplus.saturating_sub(code_used))
        .min(budget - scout_used - code_used);
        // Impact gets 150 base + 0 surplus = 150 (remaining = 350)
        assert_eq!(impact_budget_with_surplus, 150);

        // Scenario: impact uses only 30 → 120 surplus flows to placement
        let impact_used = 30;
        let placement_budget_with_surplus = (base_budgets[3]
            + impact_budget_with_surplus.saturating_sub(impact_used))
        .min(budget - scout_used - code_used - impact_used);
        // Placement gets 100 base + 120 surplus = 220 (remaining = 320)
        assert_eq!(placement_budget_with_surplus, 220);

        // Notes gets remaining
        let placement_used = 80;
        let notes_budget = budget - scout_used - code_used - impact_used - placement_used;
        assert_eq!(notes_budget, 240);
    }
}
