//! CPU affinity scenario implementations.

use super::Ctx;
use super::backdrop::Backdrop;
use super::ops::{
    CgroupDef, CpusetSpec, HoldSpec, Op, Step, execute_scenario, execute_scenario_with,
};
use crate::assert::{Assert, AssertResult};
use crate::workload::AffinityIntent;
use anyhow::Result;
use std::collections::BTreeSet;
use std::time::Duration;

/// Two cgroups with worker affinities randomized four times during the
/// run. Each randomization assigns half the available CPUs to each worker.
///
/// Both cgroups persist on the Backdrop so every `SetAffinity` op in
/// the four randomization Steps finds the same live workers.
pub fn custom_cgroup_affinity_change(ctx: &Ctx) -> Result<AssertResult> {
    let backdrop = Backdrop::new()
        .push_cgroup(CgroupDef::named("cg_0"))
        .push_cgroup(CgroupDef::named("cg_1"));
    let mut steps = vec![Step::new(vec![], ctx.settled_hold(0.2))];

    // Pool the random sample across every CPU in the topology;
    // `Op::SetAffinity` intersects with the cgroup's cpuset at apply
    // time so each cgroup ends up sampling from its actual budget.
    // Sample size is half the topology (clamped to at least 1).
    let pool = ctx.topo.all_cpuset();
    let count = (pool.len() / 2).max(1);
    let intent = AffinityIntent::RandomSubset { from: pool, count };

    for _ in 0..4 {
        steps.push(Step::new(
            vec![
                Op::set_affinity("cg_0", intent.clone()),
                Op::set_affinity("cg_1", intent.clone()),
            ],
            HoldSpec::frac(0.2),
        ));
    }

    execute_scenario(ctx, backdrop, steps)
}

/// Two cgroups with all workers pinned to the same 2-CPU subset.
/// Uses a relaxed 75% spread threshold since concentrated pinning
/// increases work-unit spread under EEVDF.
pub fn custom_cgroup_multicpu_pin(ctx: &Ctx) -> Result<AssertResult> {
    let all = ctx.topo.all_cpus();
    let pin_cpus: BTreeSet<usize> = if all.len() >= 2 {
        all[..2].iter().copied().collect()
    } else {
        all.iter().copied().collect()
    };

    // Pinning all workers to 2 CPUs concentrates load and increases
    // spread under EEVDF; relax the default spread threshold
    // (`spread_threshold_pct`: 35% debug / 15% release).
    let checks = Assert::default_checks().max_spread_pct(75.0);

    // Settle step uses a fixed minimum so `ctx.settle = Duration::ZERO`
    // (the production default) does not produce an instant no-op
    // settle step. Workers spawned during Backdrop setup (cg_0 / cg_1
    // below) need wall-clock time to migrate into their cgroup
    // cpusets before the Op::set_affinity ops fire in the next step;
    // without the 500ms floor a ZERO-duration settle would let the
    // affinity-pin step race the worker dispatch path.
    // `HoldSpec::fixed(ZERO)` itself is VALID per HoldSpec::validate
    // (scenario/ops/types/step.rs) — the floor is for scheduler-warmup
    // correctness, not validation.
    let settle = ctx.settle.max(Duration::from_millis(500));
    let backdrop = Backdrop::new()
        .push_cgroup(CgroupDef::named("cg_0"))
        .push_cgroup(CgroupDef::named("cg_1"));
    let steps = vec![
        Step::new(vec![], HoldSpec::fixed(settle)),
        Step::new(
            vec![
                Op::set_affinity("cg_0", AffinityIntent::Exact(pin_cpus.clone())),
                Op::set_affinity("cg_1", AffinityIntent::Exact(pin_cpus)),
            ],
            HoldSpec::fixed(ctx.duration),
        ),
    ];

    execute_scenario_with(ctx, backdrop, steps, Some(&checks))
}

/// Two cgroups with disjoint cpusets, workers pinned to 2 CPUs within
/// each partition. Checks pinning interacts correctly with cpuset
/// constraints. Uses a relaxed 75% spread threshold.
pub fn custom_cgroup_cpuset_multicpu_pin(ctx: &Ctx) -> Result<AssertResult> {
    let usable = ctx.topo.usable_cpus();
    let mid = usable.len() / 2;
    let a: BTreeSet<usize> = usable[..mid].iter().copied().collect();
    let b: BTreeSet<usize> = usable[mid..].iter().copied().collect();

    let pin_a: BTreeSet<usize> = a.iter().copied().take(2.min(a.len())).collect();
    let pin_b: BTreeSet<usize> = b.iter().copied().take(2.min(b.len())).collect();

    // Pinning workers to 2 CPUs within each cpuset partition
    // concentrates load and increases spread; relax the threshold.
    let checks = Assert::default_checks().max_spread_pct(75.0);

    let backdrop = Backdrop::new().extend_cgroups([
        CgroupDef::named("cg_0").cpuset(CpusetSpec::disjoint(0, 2)),
        CgroupDef::named("cg_1").cpuset(CpusetSpec::disjoint(1, 2)),
    ]);
    let steps = vec![
        Step::new(vec![], HoldSpec::fixed(ctx.settle)),
        Step::new(
            vec![
                Op::set_affinity("cg_0", AffinityIntent::Exact(pin_a)),
                Op::set_affinity("cg_1", AffinityIntent::Exact(pin_b)),
            ],
            HoldSpec::fixed(ctx.duration),
        ),
    ];

    execute_scenario_with(ctx, backdrop, steps, Some(&checks))
}
