use super::*;
use dllm_core::policy::{GateMode, Policy};
use dllm_core::serving::{Item, ItemResult};
use std::sync::atomic::Ordering::Relaxed;

#[test]
fn cached_prefix_equals_cold_and_skips_rows() {
    let prefix = toks(1, 200);
    let items: Vec<Item> = (0..4).map(|q| text_item(&format!("q{q}"), &prefix, &toks(10 + q, 20), 3)).collect();
    let (cold, _) = engine(32768, 2048, 512, |c| c.prefix_cache = false);
    let (warm, ctl) = engine(32768, 2048, 512, |_| {});
    let want = cold.run(&items, &[], &think_pol()).unwrap();
    let r1 = warm.run(&items, &[], &think_pol()).unwrap();
    let rows1 = warm.stats().prefill_rows;
    let r2 = warm.run(&items, &[], &think_pol()).unwrap();
    let rows2 = warm.stats().prefill_rows - rows1;
    for ((a, b), c) in want.iter().zip(&r1).zip(&r2) {
        assert_eq!(a.p0, b.p0);
        assert_eq!(a.p0, c.p0);
        assert_eq!(a.think, b.think);
        assert_eq!(a.think, c.think);
        assert_eq!((a.trace_tokens, a.out_tokens), (c.trace_tokens, c.out_tokens));
    }
    // 4 identical-prefix questions: first run computes the shared rows once, second run only the partial tails (<= 16 + 2 rows per branch)
    assert!(rows1 < 4 * 230, "shared prefix must be computed once, rows {rows1}");
    assert!(rows2 <= 4 * 34, "cached rerun computes only the tail, rows {rows2}");
    assert!(warm.stats().hit_rate() > 0.4);
    let _ = ctl;
}

#[test]
fn image_identity_does_not_alias_in_cache() {
    let mut imgs = vec![];
    for k in 0..3u8 {
        imgs.push(vec![7u8, k, 1, 2, 3]); // same n_soft, different content
    }
    let prefix: Vec<u32> = {
        let mut p = toks(5, 20);
        p.insert(3, 258880);
        p
    };
    let item = text_item("q", &prefix, &toks(77, 40), 2);
    let (shared, _) = engine(32768, 2048, 512, |_| {});
    let mut got = vec![];
    for im in &imgs {
        got.push(shared.run(&[item.clone()], &[im.clone()], &think_pol()).unwrap().remove(0));
    }
    for (k, im) in imgs.iter().enumerate() {
        let (fresh, _) = engine(32768, 2048, 512, |c| c.prefix_cache = false);
        let want = fresh.run(&[item.clone()], &[im.clone()], &think_pol()).unwrap().remove(0);
        assert_eq!(want.p0, got[k].p0, "image {k} aliased in the prefix cache");
    }
    assert_ne!(got[0].p0, got[1].p0);
    // same image again => cache hit and same result; the encoded soft tokens are reused
    let again = shared.run(&[item.clone()], &[imgs[1].clone()], &think_pol()).unwrap().remove(0);
    assert_eq!(again.p0, got[1].p0);
    assert!(shared.stats().soft_cache_hits + shared.stats().cache_hit_pages > 0);
}

#[test]
fn image_block_straddling_chunk_and_shared_boundary() {
    // long image (n_soft = 40) + short state: shared boundary falls inside the image block for some lengths
    for pre in 0..20u32 {
        let mut prefix = toks(9, 3 + pre as usize);
        prefix.push(258880);
        prefix.extend(toks(4, 5));
        let items: Vec<Item> = (0..3).map(|q| text_item(&format!("q{q}"), &prefix, &toks(30 + q, 25 + q as usize), 2)).collect();
        let img = vec![39u8, 5, 6, 7];
        let (a, _) = engine(32768, 2048, 512, |_| {});
        let (b, _) = engine(32768, 2048, 48, |c| { c.prefix_cache = false; c.max_prefill_rows = 48 });
        let ra = a.run(&items, &[img.clone()], &think_pol()).unwrap();
        let rb = b.run(&items, &[img.clone()], &think_pol()).unwrap();
        for (x, y) in ra.iter().zip(&rb) {
            assert_eq!((x.p0.clone(), x.think.clone()), (y.p0.clone(), y.think.clone()), "pre {pre}");
        }
    }
}

#[test]
fn bad_image_is_an_error() {
    let (e, _) = engine(32768, 256, 256, |_| {});
    let item = text_item("q", &{ let mut p = toks(2, 10); p.push(258880); p }, &toks(3, 5), 2);
    let err = e.run(&[item.clone()], &[b"BAD image".to_vec()], &think_pol()).unwrap_err();
    assert!(format!("{err:#}").contains("cannot decode image"), "{err:#}");
    // engine still works
    assert!(e.run(&[item], &[vec![3u8, 1, 1]], &think_pol()).is_ok());
    // placeholder / image count mismatch
    let plain = text_item("q", &toks(2, 10), &toks(3, 5), 2);
    assert!(e.run(&[plain], &[vec![3u8, 1, 1]], &think_pol()).is_err());
}

#[test]
fn kv_exhaustion_is_an_error_not_a_panic_and_recovers() {
    let (e, _) = engine(32768, 8, 512, |_| {});
    let big = text_item("big", &toks(1, 400), &toks(2, 10), 2);
    let err = e.run(&[big], &[], &think_pol()).unwrap_err();
    assert!(format!("{err:#}").contains("KV pool exhausted"), "{err:#}");
    // larger than the whole pool: reported as a capacity limit (server: 422 "maximum context length"), not a transient shortage
    assert!(format!("{err:#}").contains("KV cache capacity exceeded"), "{err:#}");
    let small = text_item("s", &toks(1, 20), &toks(2, 5), 2);
    assert!(e.run(&[small], &[], &Policy { tau: 0.0, budget_ms: 0, ..Policy::default() }).is_ok());
    let st = e.stats();
    assert_eq!(st.free_pages + st.tree_pages, st.num_pages);
}

#[test]
fn decode_exhaustion_midway_fails_only_that_request() {
    // a never-ending thought (cap huge) outgrows a 40-page pool: Err, pages released, a later request succeeds
    let (e, _) = engine(1 << 20, 40, 512, |_| {});
    // find an item whose thought never ends under tau=1.01 (lim = 100000) by trying toks seeds
    let mut found = false;
    for s in 0..40u32 {
        let it = text_item("n", &toks(s, 30), &toks(s + 100, 8), 2);
        let pol = Policy { tau: 1.01, cap: 4000, budget_ms: 0, ..Policy::default() };
        match e.run(&[it], &[], &pol) {
            Err(err) => {
                assert!(format!("{err:#}").contains("KV pool exhausted"), "{err:#}");
                // the prompt fits the pool; running out while thinking is not the admission capacity error
                assert!(!format!("{err:#}").contains("KV cache capacity exceeded"), "{err:#}");
                found = true;
                break;
            }
            Ok(_) => {}
        }
    }
    assert!(found, "no endless thought found among seeds");
    let st = e.stats();
    assert_eq!(st.free_pages + st.tree_pages, st.num_pages, "pages leaked after failure");
    assert!(e.run(&[text_item("s", &toks(1, 20), &toks(2, 5), 2)], &[], &Policy { tau: 0.0, budget_ms: 0, ..Policy::default() }).is_ok());
}

#[test]
fn eviction_under_tiny_pool_keeps_answers() {
    let (big, _) = engine(32768, 4096, 512, |_| {});
    let (tiny, _) = engine(32768, 60, 512, |_| {});
    for r in 0..60u32 {
        let prefix = toks(r % 25, 150);
        let items: Vec<Item> = (0..2).map(|q| text_item(&format!("q{q}"), &prefix, &toks(r * 3 + q, 12), 3)).collect();
        let a = big.run(&items, &[], &think_pol()).unwrap();
        let b = tiny.run(&items, &[], &think_pol()).unwrap();
        for (x, y) in a.iter().zip(&b) {
            assert_eq!((&x.p0, &x.think, x.trace_tokens), (&y.p0, &y.think, y.trace_tokens));
        }
    }
    let st = tiny.stats();
    assert!(st.cache_evicted_pages > 0, "expected evictions");
    assert_eq!(st.free_pages + st.tree_pages, st.num_pages);
}

#[test]
fn gate_modes() {
    let (e, _) = engine(32768, 2048, 512, |_| {});
    let items: Vec<Item> = (0..12u32).map(|q| text_item(&format!("q{q}"), &toks(1, 30), &toks(q, 9), 3)).collect();
    let n = |g: GateMode, tau: f32| e.run(&items, &[], &Policy { tau, gate: g, budget_ms: 0, cap: 5, ..Policy::default() }).unwrap().iter().filter(|r| r.think.is_some()).count();
    assert_eq!(n(GateMode::Never, 2.0), 0);
    assert_eq!(n(GateMode::Tau, 2.0), 12);
    assert_eq!(n(GateMode::Tau, 0.0), 0);
}

#[test]
fn budget_forces_readout_of_endless_thoughts() {
    // 8 questions whose thoughts would run 100000 tokens: the 150 ms wall-clock budget cuts them and reads them out
    let (e, _) = engine(1 << 20, 1 << 16, 512, |_| {});
    for s in 0..60u32 {
        let probe = text_item("n", &toks(s, 30), &toks(s + 100, 8), 2);
        if e.run(&[probe.clone()], &[], &Policy { tau: 1.01, cap: 200, budget_ms: 0, ..Policy::default() }).unwrap()[0].trace_tokens < 200 {
            continue;
        }
        let items: Vec<Item> = (0..8).map(|q| Item { key: format!("q{q}"), ..probe.clone() }).collect();
        let pol = Policy { tau: 1.01, cap: 1 << 19, budget_ms: 150, ..Policy::default() };
        let t = std::time::Instant::now();
        let r = e.run(&items, &[], &pol).unwrap();
        assert!(e.stats().forced_by_budget >= 8, "forced {}", e.stats().forced_by_budget);
        assert!(r.iter().all(|r| r.think.is_some() && r.trace_tokens > 0 && r.trace_tokens < 100000));
        assert!(t.elapsed().as_millis() < 1000, "the budget must cut the thoughts");
        return;
    }
    panic!("no endless thought among the seeds");
}

#[test]
fn concurrency_equals_serial_with_prefix_sharing() {
    let (e, ctl) = engine(32768, 3000, 512, |_| {});
    let e = std::sync::Arc::new(e);
    let mut reqs = vec![];
    for r in 0..24u32 {
        let prefix = toks(r % 6, 100 + 7 * (r % 6) as usize);
        reqs.push((0..3u32).map(|q| text_item(&format!("q{q}"), &prefix, &toks(r * 5 + q, 10 + q as usize), 3)).collect::<Vec<_>>());
    }
    let serial: Vec<_> = reqs.iter().map(|r| e.run(r, &[], &think_pol()).unwrap()).collect();
    let reqs = std::sync::Arc::new(reqs);
    let hs: Vec<_> = (0..32).map(|t| {
        let (e, reqs) = (e.clone(), reqs.clone());
        std::thread::spawn(move || (0..reqs.len()).map(|i| (i, e.run(&reqs[(i + t) % reqs.len()], &[], &think_pol()).unwrap(), (i + t) % reqs.len())).collect::<Vec<_>>())
    }).collect();
    for h in hs {
        for (_, res, ri) in h.join().unwrap() {
            for (a, b) in res.iter().zip(&serial[ri]) {
                assert_eq!((&a.p0, &a.think, a.trace_tokens, a.out_tokens), (&b.p0, &b.think, b.trace_tokens, b.out_tokens));
            }
        }
    }
    assert!(e.stats().hit_rate() > 0.5);
    let _ = ctl.rows.load(Relaxed);
}

#[test]
fn think_override_and_trace_ids() {
    use dllm_core::policy::ThinkMode;
    let (e, _) = engine(32768, 2048, 512, |_| {});
    let items: Vec<Item> = (0..12u32).map(|q| text_item(&format!("q{q}"), &toks(1, 30), &toks(q, 9), 3)).collect();
    let base = Policy { tau: 0.0, budget_ms: 0, cap: 20, ..Policy::default() }; // server policy that never thinks
    let auto = e.run(&items, &[], &base.with_think(ThinkMode::Auto)).unwrap();
    assert!(auto.iter().all(|r| r.think.is_none() && r.trace_ids.is_empty()));
    let on = e.run(&items, &[], &base.with_think(ThinkMode::On)).unwrap();
    assert!(on.iter().all(|r| r.think.is_some() && r.trace_tokens > 0 && r.trace_ids.is_empty()), "trace ids only on request");
    let on_tr = e.run(&items, &[], &Policy { return_trace: true, ..base.with_think(ThinkMode::On) }).unwrap();
    for (a, b) in on.iter().zip(&on_tr) {
        assert_eq!((&a.p0, &a.think, a.trace_tokens, a.out_tokens), (&b.p0, &b.think, b.trace_tokens, b.out_tokens), "return_trace must not change results");
        assert_eq!(b.trace_ids.len(), b.trace_tokens);
        assert!(b.trace_tokens <= 20);
    }
    let think_all = Policy { tau: 2.0, ..base };
    let off = e.run(&items, &[], &Policy { return_trace: true, ..think_all.with_think(ThinkMode::Off) }).unwrap();
    assert!(off.iter().all(|r| r.think.is_none() && r.trace_tokens == 0 && r.trace_ids.is_empty()));
}

/// --think-ms: the per-branch cap is a pure function of the request (prefill rows, images, questions, thinking questions) under a fixed cost
/// model; the same request gives identical answers and trace lengths alone, repeated, chunked over several prefill forwards and under
/// concurrent load (other requests interleaved in the same forwards)
#[test]
fn think_ms_cap_is_deterministic_and_follows_the_cost_model() {
    use dllm_core::policy::{ThinkCost, ThinkShape};
    let cost = ThinkCost::from_json(r#"{"base_ms":20,"prefill_ms_per_tok":0.1,"image_ms":0,"question_ms":1,"readout_ms":5,"readout_ms_per_branch":1,"tok_ms":[[1,2.0],[4,4.0],[16,10.0]],"min_tokens":2}"#).unwrap();
    let pol = Policy { tau: 1.01, cap: 4000, budget_ms: 0, think_ms: 150, think_cost: cost, ..Policy::default() };
    let free = Policy { think_ms: 0, ..pol };
    // requests of 1..4 questions sharing a prefix (all think: tau 1.01)
    let prefix = toks(11, 120);
    let reqs: Vec<Vec<Item>> = (1..=4).map(|n| (0..n).map(|q| text_item(&format!("q{q}"), &prefix, &toks(40 + 7 * q as u32 + n as u32, 30 + 5 * q), 3)).collect()).collect();
    let expect_cap = |items: &[Item]| {
        let ids: Vec<Vec<i32>> = items.iter().map(|it| it.ids.iter().map(|&t| t as i32).collect()).collect();
        let s = ThinkShape { prefill_tokens: crate::engine::prefill_rows(&ids), questions: items.len(), thinking: items.len(), ..Default::default() };
        cost.cap(150.0, &s, 4000)
    };
    let (e, _) = engine(32768, 4096, 512, |c| c.prefix_cache = false);
    let mut want = vec![];
    for items in &reqs {
        let cap = expect_cap(items);
        assert!(cap > 0 && cap < 100, "test request must get a small positive cap, got {cap}");
        let unl = e.run(items, &[], &free).unwrap();
        let r = e.run(items, &[], &pol).unwrap();
        for (a, b) in unl.iter().zip(&r) {
            assert_eq!(a.p0, b.p0);
            assert_eq!(b.trace_tokens, a.trace_tokens.min(cap), "trace = unrestricted thought cut at the cost-model cap {cap}");
        }
        want.push(r);
    }
    // more thinking questions => smaller per-branch cap
    assert!(expect_cap(&reqs[3]) < expect_cap(&reqs[0]));
    let same = |a: &[ItemResult], b: &[ItemResult]| a.iter().zip(b).all(|(x, y)| x.p0 == y.p0 && x.think == y.think && x.trace_tokens == y.trace_tokens);
    // repeated
    for (items, w) in reqs.iter().zip(&want) {
        assert!(same(&e.run(items, &[], &pol).unwrap(), w));
    }
    // prefill split over many small forwards (max_rows 64: gate decisions of one request arrive in different iterations)
    let (small, _) = engine(32768, 4096, 64, |c| c.prefix_cache = false);
    for (items, w) in reqs.iter().zip(&want) {
        assert!(same(&small.run(items, &[], &pol).unwrap(), w), "chunked prefill changed the answer");
    }
    // concurrent load with the radix cache on: 8 threads x 3 rounds of all requests interleaved
    let (busy, _) = engine(32768, 4096, 512, |_| {});
    let busy = std::sync::Arc::new(busy);
    let hs: Vec<_> = (0..8)
        .map(|t| {
            let (busy, reqs, want, pol) = (busy.clone(), reqs.clone(), want.clone(), pol);
            std::thread::spawn(move || {
                for k in 0..3 {
                    let i = (t + k) % reqs.len();
                    assert!(same(&busy.run(&reqs[i], &[], &pol).unwrap(), &want[i]), "load changed the answer");
                }
            })
        })
        .collect();
    for h in hs {
        h.join().unwrap();
    }
}

/// number questions: greedy answer after the end-of-thought token until end-of-turn (or ANSWER_MAX tokens); `think: on` thinks first
#[test]
fn number_questions_generate_the_greedy_answer() {
    use crate::engine::{ANSWER_MAX, EOS, MK, TURN_END};
    use crate::fake::{logits, step, REC0};
    use dllm_core::policy::ThinkMode;
    let greedy = |rec: &mut crate::fake::Rec, excl: &[i32]| -> i32 {
        let l = logits(rec);
        let mut b = -1i32;
        for v in 0..256i32 {
            if excl.contains(&v) {
                continue;
            }
            if b < 0 || l[v as usize] > l[b as usize] {
                b = v;
            }
        }
        b
    };
    let stop = |t: i32| t == MK || t == EOS || t == TURN_END;
    let (e, _) = engine(32768, 2048, 512, |_| {});
    let mut checked = 0;
    for s in 0..12u32 {
        let mut it = text_item("n", &toks(s, 30), &toks(s + 50, 7), 0);
        it.number = true;
        it.label_ids = vec![];
        it.keys = vec![];
        let base = Policy { tau: 0.8, cap: 40, budget_ms: 0, ..Policy::default() };
        for think in [false, true] {
            let pol = if think { base.with_think(ThinkMode::On) } else { base };
            let r = e.run(&[it.clone()], &[], &Policy { return_trace: true, ..pol }).unwrap().remove(0);
            // reference with the fake LM
            let mut rec = REC0;
            for &t in &it.ids {
                rec = step(&rec, t as i32, None, 0);
            }
            let mut thought = vec![];
            if think {
                let mut x = rec;
                let mut t = greedy(&mut x, &[MK, EOS, TURN_END]);
                loop {
                    thought.push(t);
                    x = step(&x, t, None, 0);
                    t = greedy(&mut x, &[]);
                    if stop(t) || thought.len() >= 40 {
                        break;
                    }
                }
                for &t in &thought {
                    rec = step(&rec, t, None, 0);
                }
            }
            rec = step(&rec, MK, None, 0);
            let mut ans = vec![];
            let mut lp = 0.0f64;
            loop {
                let l = logits(&rec);
                let t = greedy(&mut rec, &[]);
                if stop(t) || ans.len() >= ANSWER_MAX {
                    break;
                }
                let m = l[t as usize] as f64;
                lp += m - (m + l.iter().map(|&v| (v as f64 - m).exp()).sum::<f64>().ln());
                ans.push(t as u32);
                rec = step(&rec, t, None, 0);
            }
            assert_eq!(r.answer_ids, ans, "seed {s} think {think}");
            assert!((r.answer_logprob - lp).abs() < 1e-9);
            assert_eq!(r.think.is_some(), think);
            assert!(r.p0.is_empty());
            if think {
                assert_eq!(r.trace_ids, thought.iter().map(|&t| t as u32).collect::<Vec<_>>(), "seed {s}");
            }
            checked += !ans.is_empty() as usize;
        }
    }
    assert!(checked > 4, "fake LM produced too few non-empty answers ({checked})");
    // mixed with choice questions in one request: the choice answers are unchanged
    let mut num = text_item("n", &toks(1, 30), &toks(51, 7), 0);
    num.number = true;
    num.label_ids = vec![];
    num.keys = vec![];
    let ch = text_item("c", &toks(1, 30), &toks(52, 7), 3);
    let alone = e.run(&[ch.clone()], &[], &think_pol()).unwrap();
    let mixed = e.run(&[num, ch], &[], &think_pol()).unwrap();
    assert_eq!((&alone[0].p0, &alone[0].think, alone[0].trace_tokens), (&mixed[1].p0, &mixed[1].think, mixed[1].trace_tokens));
}

/// a decoder that answers several tokens per cycle (as a speculative decoder accepting all its drafts): `n` one-row greedy forwards per branch,
/// KV written up to pos + n - 1 and surplus tokens after a stop / ANSWER_MAX returned like rejected-then-ignored ones
struct Burst(usize);
impl crate::ThinkDecoder for Burst {
    fn lookahead(&self) -> usize {
        self.0 - 1
    }
    fn step(&mut self, m: &mut dyn ModelRunner, reqs: &[crate::ThinkReq]) -> anyhow::Result<Vec<Vec<i32>>> {
        crate::GreedyDecoder.step(m, reqs)
    }
    fn answer_step(&mut self, m: &mut dyn ModelRunner, reqs: &[crate::ThinkReq]) -> anyhow::Result<Vec<Vec<(i32, f64)>>> {
        let mut out = vec![];
        for r in reqs {
            let (mut tok, mut acc) = (r.tok, vec![]);
            for j in 0..self.0 {
                m.forward(&crate::model::RowBatch::decode(&[tok], &[(r.pages, r.pos + j)]), None)?;
                let z = m.logits_rows(&[0])?.remove(0);
                acc.push(dllm_model::spec::greedy_logprob(&z));
                tok = acc[j].0;
            }
            out.push(acc);
        }
        Ok(out)
    }
}

/// number answers decoded several tokens per cycle equal the one-token-per-step answers (ids and log-probability bit for bit), with and without a thought
#[test]
fn multi_token_answer_cycles_equal_one_token_steps() {
    use dllm_core::policy::ThinkMode;
    let run = |n: usize| {
        let m = FakeModel::full_kv(2048, 512);
        let cfg = EngineConfig { max_len: 32768, prep: Some(Arc::new(FakePrep)), ..Default::default() };
        let e = Engine::start_with(cfg, Box::new(m), Some(Box::new(FakeVision { encodes: 0 })), tok(), Box::new(Burst(n))).unwrap();
        let mut out = vec![];
        let num = |key: &str, prefix: &[u32], tail: &[u32]| Item { number: true, label_ids: vec![], keys: vec![], ..text_item(key, prefix, tail, 0) };
        for s in 0..12u32 {
            let base = Policy { tau: 0.8, cap: 40, budget_ms: 0, ..Policy::default() };
            for think in [false, true] {
                let pol = if think { base.with_think(ThinkMode::On) } else { base };
                // two questions in one request: answers of different lengths share answer cycles
                let r = e.run(&[num("n", &toks(s, 30), &toks(s + 50, 7)), num("m", &toks(s, 30), &toks(s + 90, 5))], &[], &pol).unwrap();
                out.push(r.iter().map(|x| (x.answer_ids.clone(), x.answer_logprob.to_bits(), x.out_tokens)).collect::<Vec<_>>());
            }
        }
        out
    };
    let want = run(1);
    assert!(want.iter().flatten().filter(|x| x.0.len() >= 2).count() > 4, "fake LM produced too few multi-token answers");
    for n in [2, 3, 7, 40] {
        assert_eq!(run(n), want, "{n} tokens per cycle");
    }
}