//! Bounded, anonymous timing estimates and two-worker queue planning.
//! Costs retain only numeric aggregates. The waiting queue holds current record
//! IDs in memory; neither it nor timing estimates are saved or logged.
use std::collections::HashMap;
use std::time::Instant;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Kind {
    Reactive,
    Idle,
    Diary,
}
impl Kind {
    fn index(self) -> usize {
        match self {
            Self::Reactive => 0,
            Self::Idle => 1,
            Self::Diary => 2,
        }
    }
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Work {
    pub kind: Kind,
    band: usize,
}
impl Work {
    pub fn new(kind: Kind, estimated_tokens: usize) -> Self {
        Self {
            kind,
            band: if estimated_tokens < 256 {
                0
            } else if estimated_tokens < 640 {
                1
            } else {
                2
            },
        }
    }
    fn index(self) -> usize {
        self.kind.index() * 3 + self.band
    }
}
#[derive(Clone, Copy, Default)]
struct Sample {
    count: u64,
    mean_ms: f64,
}
#[derive(Clone, Default)]
pub struct Costs {
    samples: [Sample; 9],
}
impl Costs {
    pub fn observe(&mut self, work: Work, ms: f64) {
        if !ms.is_finite() || ms <= 0.0 {
            return;
        }
        let sample = &mut self.samples[work.index()];
        // A plain mean for the first five samples, so one unusual first turn
        // does not dominate; then an exponentially weighted mean (20% new).
        let weight = (1.0 / (sample.count as f64 + 1.0)).max(0.2);
        sample.mean_ms += (ms - sample.mean_ms) * weight;
        sample.count = sample.count.saturating_add(1);
    }
    pub fn estimate(&self, work: Work) -> Option<f64> {
        let exact = self.samples[work.index()];
        if exact.count > 0 {
            return Some(exact.mean_ms);
        }
        // A different prompt band is deliberately not borrowed: CPU prefill
        // scales much more sharply than CUDA. Unseen work retains affinity.
        None
    }
    pub fn sample_count(&self, work: Work) -> u64 {
        self.samples[work.index()].count
    }
}

/// Per-lane estimates for planning. When only one lane has run this shape,
/// the other lane's estimate is that cost times the least favourable speed
/// ratio seen on shapes both lanes have run. This lets an unseen GPU shape
/// wait for a GPU that is clearly faster instead of always spilling to an
/// idle CPU. Without such a shared shape, unseen work keeps its affinity.
pub fn estimates(costs: [&Costs; 2], work: Work) -> [Option<f64>; 2] {
    let ratio = |lane: usize| {
        costs[lane]
            .samples
            .iter()
            .zip(&costs[1 - lane].samples)
            .filter(|(a, b)| a.count > 0 && b.count > 0)
            .map(|(a, b)| a.mean_ms / b.mean_ms)
            .reduce(f64::max)
    };
    match [costs[0].estimate(work), costs[1].estimate(work)] {
        [None, Some(gpu)] => [ratio(0).map(|r| gpu * r), Some(gpu)],
        [Some(cpu), None] => [Some(cpu), ratio(1).map(|r| cpu * r)],
        exact => exact,
    }
}

pub struct Waiting {
    pub epoch: u64,
    pub context: u64,
    pub work: Work,
    pub since: Instant,
}
#[derive(Default)]
pub struct Queue {
    pub waiting: HashMap<String, Waiting>,
}
impl Queue {
    pub fn age_ms(&mut self, id: &str, epoch: u64, context: u64, work: Work, now: Instant) -> f64 {
        let entry = self.waiting.entry(id.into()).or_insert(Waiting {
            epoch,
            context,
            work,
            since: now,
        });
        if entry.epoch != epoch || entry.work.kind != work.kind {
            *entry = Waiting {
                epoch,
                context,
                work,
                since: now,
            };
        }
        entry.context = context;
        entry.work = work;
        now.duration_since(entry.since).as_secs_f64() * 1000.0
    }
}
#[derive(Clone, Copy)]
pub struct Request {
    pub work: Work,
    pub age_ms: f64,
}
pub const BACKGROUND_MAX_WAIT_MS: f64 = 30000.0;

pub fn priority(request: Request) -> u8 {
    if request.work.kind != Kind::Reactive && request.age_ms >= BACKGROUND_MAX_WAIT_MS {
        0
    } else if request.work.kind == Kind::Reactive {
        1
    } else {
        2
    }
}

// Returns a lane for every request, in the supplied priority/FIFO order.
// A lane can be busy: its planned request waits instead of spilling blindly.
pub fn plan(
    costs: [&Costs; 2],
    remaining: [f64; 2],
    available: [bool; 2],
    free: [bool; 2],
    requests: &[Request],
) -> Vec<usize> {
    let mut finish = remaining;
    let mut cold_free = free;
    requests
        .iter()
        .map(|request| {
            let times = estimates(costs, request.work);
            let preferred = if request.work.kind == Kind::Reactive {
                1
            } else {
                0
            };
            let lane = if !available[0] {
                1
            } else if !available[1] {
                0
            } else if let [Some(cpu), Some(gpu)] = times {
                // Prefer a clear improvement before sending work to a slower
                // lane. The 10% margin limits oscillation from timing noise.
                let cpu_end = finish[0] + cpu;
                let gpu_end = finish[1] + gpu;
                if cpu_end * 1.1 < gpu_end { 0 } else { 1 }
            } else if cold_free[preferred] {
                preferred
            } else if cold_free[1 - preferred] {
                1 - preferred
            } else {
                preferred
            };
            cold_free[lane] = false;
            finish[lane] += times[lane].unwrap_or(remaining[lane].max(1.0));
            lane
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    fn request() -> Request {
        Request {
            work: Work::new(Kind::Reactive, 100),
            age_ms: 0.0,
        }
    }
    fn close(value: Option<f64>, expected: f64) -> bool {
        value.is_some_and(|v| (v - expected).abs() < 1e-9)
    }
    fn costs() -> [Costs; 2] {
        let mut c = [Costs::default(), Costs::default()];
        c[0].observe(request().work, 140.0);
        c[1].observe(request().work, 30.0);
        c
    }
    #[test]
    fn small_queue_waits_for_gpu_and_large_queue_uses_both() {
        let c = costs();
        assert_eq!(
            plan(
                [&c[0], &c[1]],
                [0.0; 2],
                [true; 2],
                [true; 2],
                &vec![request(); 4]
            ),
            vec![1; 4]
        );
        let large = plan(
            [&c[0], &c[1]],
            [0.0; 2],
            [true; 2],
            [true; 2],
            &vec![request(); 16],
        );
        assert!(large.contains(&0) && large.contains(&1));
        assert_eq!(
            plan(
                [&c[0], &c[1]],
                [0.0, 20.0],
                [true; 2],
                [true, false],
                &vec![request(); 2]
            ),
            vec![1; 2]
        );
        assert_eq!(
            plan(
                [&c[0], &c[1]],
                [0.0, f64::INFINITY],
                [true; 2],
                [true, false],
                &[request()]
            ),
            vec![0],
            "a busy unseen GPU shape cannot promise an immediate finish"
        );
    }
    #[test]
    fn unknown_shapes_keep_affinity_and_unavailable_gpu_uses_cpu() {
        let c = costs();
        let requests = [
            request(),
            Request {
                work: Work::new(Kind::Diary, 700),
                age_ms: 0.0,
            },
        ];
        assert_eq!(
            plan([&c[0], &c[1]], [0.0; 2], [true; 2], [true; 2], &requests),
            vec![1, 0]
        );
        assert_eq!(
            plan(
                [&c[0], &c[1]],
                [0.0; 2],
                [true, false],
                [true, false],
                &requests
            ),
            vec![0, 0]
        );
    }
    #[test]
    fn estimates_are_finite_bounded_and_separate_by_work_shape() {
        let mut c = Costs::default();
        let work = request().work;
        c.observe(work, 100.0);
        c.observe(work, 200.0);
        c.observe(work, f64::NAN);
        c.observe(work, -1.0);
        assert_eq!(c.estimate(work), Some(150.0));
        assert_eq!(c.sample_count(work), 2);
        assert_eq!(c.estimate(Work::new(Kind::Reactive, 900)), None);
        assert_eq!(c.estimate(Work::new(Kind::Idle, 100)), None);
    }
    #[test]
    fn an_unusual_first_sample_is_averaged_out_quickly() {
        // Measured on the GPU: a two-call clock turn, then single-call turns.
        let mut c = Costs::default();
        for ms in [64.0, 30.0, 28.0, 28.0, 33.0] {
            c.observe(request().work, ms);
        }
        assert!(close(c.estimate(request().work), 36.6));
        c.observe(request().work, 31.0);
        assert!(close(c.estimate(request().work), 35.48));
    }
    #[test]
    fn a_shape_seen_on_one_lane_is_estimated_from_shared_shapes() {
        let long = Work::new(Kind::Reactive, 900);
        let mut c = costs();
        c[0].observe(long, 1070.0);
        // Measured 4-pet mixed burst: the GPU was busy and had never run the
        // long prompt, so the idle CPU took it for about one second.
        let busy_gpu = [0.0, 30.0];
        let long_request = Request {
            work: long,
            age_ms: 0.0,
        };
        assert!(close(
            estimates([&c[0], &c[1]], long)[1],
            1070.0 * 30.0 / 140.0
        ));
        assert_eq!(
            plan(
                [&c[0], &c[1]],
                busy_gpu,
                [true; 2],
                [true, false],
                &[long_request]
            ),
            vec![1]
        );
        // The CPU side uses the least favourable CPU/GPU ratio.
        let mut d = costs();
        d[1].observe(long, 100.0);
        assert!(close(
            estimates([&d[0], &d[1]], long)[0],
            100.0 * 140.0 / 30.0
        ));
        assert_eq!(
            plan(
                [&d[0], &d[1]],
                [0.0; 2],
                [true; 2],
                [true; 2],
                &[long_request; 4]
            ),
            vec![1; 4]
        );
        assert!(
            plan(
                [&d[0], &d[1]],
                [0.0; 2],
                [true; 2],
                [true; 2],
                &[long_request; 16]
            )
            .contains(&0)
        );
        // Without a shape seen on both lanes, unseen work keeps its affinity.
        let mut e = [Costs::default(), Costs::default()];
        e[0].observe(long, 1070.0);
        assert_eq!(estimates([&e[0], &e[1]], long), [Some(1070.0), None]);
    }
    #[test]
    fn background_ages_ahead_of_new_reactive_work_and_context_edits_keep_its_place() {
        let work = Work::new(Kind::Diary, 100);
        let start = Instant::now();
        let mut queue = Queue::default();
        assert_eq!(queue.age_ms("test-record", 1, 1, work, start), 0.0);
        let later = start + std::time::Duration::from_secs(31);
        let age = queue.age_ms("test-record", 1, 2, Work::new(Kind::Diary, 900), later);
        assert_eq!(age, 31000.0);
        assert!(priority(Request { work, age_ms: age }) < priority(request()));
        assert_eq!(
            queue.age_ms("test-record", 2, 3, work, later),
            0.0,
            "a new executor starts a new queue wait"
        );
    }
}