//! Unit tests of the reference adaptation policies (http-delivery Appendix F): the
//! integer log, each policy at points worked by hand, the budget rule, layer depth, and
//! an entry order that keeps the pick the same on the range path and the whole path.

use super::abr::*;
use super::head::{parse_head, Head, HeadParse};
use super::head_proof::HeadProof;
use super::origin;
use super::plan::FetchPolicy;
use super::session::{drive, FetchSession, ResponseMeta, Selector, VerifiedBy, Want};
use super::{content_hash, Fragment, RangeRefusal};
use crate::container::{Rendition, RenditionMeta, Skip, WaiMulti};
use crate::energy_class::EnergyProvenance;

fn pattern(len: usize, a: usize, b: usize) -> Vec<u8> {
    (0..len).map(|i| ((a * i + b) % 256) as u8).collect()
}

fn head_of(env: &[u8]) -> Head {
    match parse_head(env, env.len() as u64).unwrap() {
        HeadParse::Head(h) => h,
        HeadParse::Need(n) => panic!("needs {n}"),
    }
}

fn state(buffer_ms: u32, bps: u64, options: &[u64]) -> AbrState {
    AbrState {
        buffer_ms,
        // One second at `bps`.
        samples: if bps == 0 { Vec::new() } else { vec![Sample { at_ms: 1000, bytes: bps / 8, ms: 1000 }] },
        last: None,
        option_bytes: options.to_vec(),
        option_joules_micro: vec![None; options.len()],
        joule_budget_micro: None,
        segment_ms: 1000,
    }
}

const OPTIONS: [u64; 3] = [10_000, 100_000, 1_000_000];

#[test]
fn ilog2_q16_points() {
    for (a, b, want) in [(1u64, 1u64, 0u64), (1, 2, 0), (2, 1, 65_536), (4, 1, 131_072), (3, 2, 38_336), (3, 1, 103_872), (17, 16, 5_732),
                         (1000, 1, 653_108), (u64::MAX, 1, 63 * 65_536 + 65_535), (5, 0, 152_170)] {
        assert_eq!(ilog2_q16(a, b), want, "log2({a}/{b})");
    }
    // Within 48/65536 of log2 over a sweep (linear interpolation over sixteenths of an
    // octave), never above it, and never decreasing in a.
    let mut prev = 0;
    for a in 1..5000u64 {
        let got = ilog2_q16(a, 1);
        assert!(got >= prev);
        prev = got;
        let exact = (a as f64).log2() * 65536.0;
        let gap = exact - (got as f64);
        assert!(gap > -1.0 && gap < 48.0, "a={a}: {got} against {exact}");
    }
}

/// The throughput policy at points worked by hand: options of 80 kb/s, 800 kb/s and
/// 8 Mb/s, a 900 ‰ margin.
#[test]
fn throughput_points() {
    let p = Throughput { last_n: 3, safety_permille: 900 };
    assert_eq!(throughput_bps(&state(0, 10_000_000, &OPTIONS).samples, 3), 10_000_000);
    assert_eq!(p.choose(&state(0, 10_000_000, &OPTIONS)), 2, "9 Mb/s fits 8");
    assert_eq!(p.choose(&state(0, 8_000_000, &OPTIONS)), 1, "7.2 Mb/s does not fit 8");
    assert_eq!(p.choose(&state(0, 1_000_000, &OPTIONS)), 1);
    assert_eq!(p.choose(&state(0, 80_000, &OPTIONS)), 0, "below the lowest: the lowest");
    assert_eq!(p.choose(&state(0, 0, &OPTIONS)), 0, "no sample");
    assert_eq!(p.choose(&state(0, 10_000_000, &[])), 0);
    // Only the last n samples count; a sample of 0 ms counts as 1 ms.
    let mut s = state(0, 0, &OPTIONS);
    s.samples = vec![Sample { at_ms: 0, bytes: 10, ms: 1000 }, Sample { at_ms: 1, bytes: 1_250_000, ms: 1000 }];
    assert_eq!(Throughput { last_n: 1, safety_permille: 1000 }.choose(&s), 2);
    assert_eq!(Throughput { last_n: 2, safety_permille: 1000 }.choose(&s), 1);
    assert_eq!(throughput_bps(&[Sample { at_ms: 0, bytes: 1, ms: 0 }], 0), 8000);
}

#[test]
fn buffer_based_points() {
    let p = BufferBased { reservoir_ms: 5000, cushion_ms: 10_000 };
    assert_eq!(p.choose(&state(5000, 0, &OPTIONS)), 0);
    assert_eq!(p.choose(&state(15_000, 0, &OPTIONS)), 2);
    // 80 kb/s + (8 Mb/s − 80 kb/s)·5/10 = 4.04 Mb/s: the 800 kb/s option.
    assert_eq!(p.choose(&state(10_000, 0, &OPTIONS)), 1);
    // 80 kb/s + 7.92 Mb/s · 100/10000 = 159.2 kb/s: only the lowest fits.
    assert_eq!(p.choose(&state(5_100, 0, &OPTIONS)), 0);
    assert_eq!(BufferBased { reservoir_ms: 5000, cushion_ms: 0 }.choose(&state(5001, 0, &OPTIONS)), 2);
}

/// Utility over buffer: the choice never falls as the buffer grows, an empty buffer
/// takes the lowest option and a buffer past every threshold the highest.
#[test]
fn bola_like_shape() {
    let p = BolaLike { v_milli: 1000, max_buffer_ms: 20_000 };
    let opts = [100_000u64, 200_000, 400_000, 800_000];
    let mut prev = 0;
    for b in (0..=30_000).step_by(250) {
        let c = p.choose(&state(b, 0, &opts));
        assert!(c >= prev, "buffer {b}: {c} after {prev}");
        prev = c;
    }
    assert_eq!(p.choose(&state(0, 0, &opts)), 0);
    assert_eq!(p.choose(&state(30_000, 0, &opts)), 3);
    // A maximum buffer of one segment or less leaves V at 0: the lowest option until the
    // buffer is past it, the scores then all equal.
    assert_eq!(BolaLike { v_milli: 1000, max_buffer_ms: 1000 }.choose(&state(0, 0, &opts)), 0);
    assert_eq!(BolaLike { v_milli: 1000, max_buffer_ms: 500 }.choose(&state(9000, 0, &opts)), 0);
}

#[test]
fn mpc_shape() {
    let p = Mpc { horizon: 3, safety_permille: 1000 };
    assert_eq!(p.choose(&state(0, 0, &OPTIONS)), 0, "no prediction");
    assert_eq!(p.choose(&state(10_000, 100_000_000, &OPTIONS)), 2);
    assert_eq!(p.choose(&state(0, 100_000, &OPTIONS)), 0);
    let mut prev = 0;
    for bps in (100_000..20_000_000).step_by(250_000) {
        let c = p.choose(&state(4000, bps, &OPTIONS));
        assert!(c >= prev, "{bps}: {c} after {prev}");
        prev = c;
    }
    assert_eq!(prev, 2);
    // Many options shorten the horizon; the choice is still made.
    let many: Vec<u64> = (1..=40).map(|k| k * 10_000).collect();
    assert!(Mpc { horizon: 5, safety_permille: 900 }.choose(&state(8000, 3_000_000, &many)) < 40);
}

/// A budget only lowers the choice, to the highest option at or below it with a labelled
/// figure within it; with none, option 0 and the budget unmet. It is not applied only
/// when no option has a labelled figure, so an unlabelled choice does not escape it, and
/// it never raises the choice, not to a labelled option within it nor past it.
#[test]
fn abr_budget_holds_labelled_options() {
    let mut s = state(0, 100_000_000, &OPTIONS);
    s.joule_budget_micro = Some(100);
    let p = Throughput { last_n: 1, safety_permille: 900 };
    let d = |s: &AbrState| decide(&p, s).unwrap();
    assert_eq!(d(&s), (2, Budget::NotApplied));
    s.option_joules_micro = vec![Some(10), Some(50), None];
    assert_eq!(d(&s), (1, Budget::Lowered), "the unlabelled choice is lowered to the highest labelled option within the budget");
    s.option_joules_micro = vec![Some(10), Some(50), Some(500)];
    assert_eq!(d(&s), (1, Budget::Lowered));
    s.option_joules_micro = vec![Some(10), None, Some(500)];
    assert_eq!(d(&s), (0, Budget::Lowered), "the unlabelled option is passed over");
    s.option_joules_micro = vec![Some(200), None, Some(500)];
    assert_eq!(d(&s), (0, Budget::Unmet), "none within: option 0, the budget unmet");
    s.option_joules_micro = vec![None, Some(300), Some(200)];
    assert_eq!(d(&s), (0, Budget::Unmet), "none within: option 0, not the least figure");
    s.option_joules_micro = vec![Some(10), Some(50), Some(100)];
    assert_eq!(d(&s), (2, Budget::Within));
    // A choice of option 0 at low throughput is never raised: not to labelled options
    // within the budget above it, nor to an option past the budget.
    let mut low = state(0, 10_000, &OPTIONS);
    low.joule_budget_micro = Some(100);
    low.option_joules_micro = vec![None, Some(50), Some(80)];
    assert_eq!(d(&low), (0, Budget::Unmet));
    low.joule_budget_micro = Some(1000);
    low.option_joules_micro = vec![None, None, Some(7000)];
    assert_eq!(d(&low), (0, Budget::Unmet), "not raised to the largest option, past the budget");
    low.option_joules_micro = vec![None, Some(900), None];
    assert_eq!(d(&low), (0, Budget::Unmet), "not raised to the one option within the budget");
    low.joule_budget_micro = None;
    assert_eq!(d(&low), (0, Budget::NotSet));
    // Whatever the figures and the budget, the choice under it is never above the
    // policy's.
    let mut seed = 0x9e37_79b9_7f4a_7c15u64;
    let mut next = || { seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); seed >> 33 };
    for _ in 0..2000 {
        let mut st = state(0, 1_000 + next() % 20_000_000, &OPTIONS);
        st.option_joules_micro = (0..3).map(|_| if next() % 3 == 0 { None } else { Some(next() % 1000) }).collect();
        st.joule_budget_micro = Some(next() % 1000);
        let policy = p.choose(&st);
        let (k, b) = d(&st);
        assert!(k <= policy, "{st:?}: {k} above the policy's {policy}");
        assert!(b != Budget::Unmet || k == 0);
    }
    s.joule_budget_micro = None;
    assert_eq!(d(&s), (2, Budget::NotSet));
    // Only a figure with its class is a labelled figure.
    assert_eq!(labelled_figure(500, None), None);
    assert_eq!(labelled_figure(500, Some(&EnergyProvenance::ModelBased)), Some(500));
    assert_eq!(labelled_figure(0, Some(&EnergyProvenance::ModelBased)), None);
    assert_eq!(labelled_figure(0, None), None);
}

/// Layer depth over a base and two layers: the bytes through each depth, the budget on
/// the sum of labelled figures, and a budget that cannot be applied.
#[test]
fn depth_over_layers() {
    let layers = [1825u64, 12_711, 18_367];
    let p = Throughput { last_n: 1, safety_permille: 1000 };
    // Through depth 2: 32,903 B a second, 263,224 b/s.
    assert_eq!(depth(&layers, &p, &state(0, 263_224, &[])).unwrap(), (2, Budget::NotSet));
    assert_eq!(depth(&layers, &p, &state(0, 263_223, &[])).unwrap(), (1, Budget::NotSet));
    assert_eq!(depth(&layers, &p, &state(0, 14_600, &[])).unwrap(), (0, Budget::NotSet));
    let mut s = state(0, 10_000_000, &[]);
    s.joule_budget_micro = Some(30);
    s.option_joules_micro = vec![Some(5), Some(10), Some(20)];
    assert_eq!(depth(&layers, &p, &s).unwrap(), (1, Budget::Lowered), "5 + 10 + 20 is past 30");
    s.option_joules_micro = vec![Some(5), None, Some(20)];
    assert_eq!(depth(&layers, &p, &s).unwrap(), (0, Budget::Lowered), "an unlabelled layer leaves the sum through it unknown, so the base alone");
    s.option_joules_micro = vec![None, None, None];
    assert_eq!(depth(&layers, &p, &s).unwrap(), (2, Budget::NotApplied));
    assert_eq!(depth(&[], &p, &s).unwrap().0, 0);
}

/// Appendix F's bounds: a state of more than 256 options, or a depth over more than 256
/// objects or whose bytes or figure pass 2^64 − 1, is refused; within the bounds every
/// policy chooses at the extremes of its types without overflow.
#[test]
fn abr_bounds_and_extremes() {
    let p = Throughput { last_n: 1, safety_permille: 900 };
    let code = |r: Result<(usize, Budget), AbrError>| r.unwrap_err().code();
    assert_eq!(code(decide(&p, &state(0, 1000, &[1; MAX_OPTIONS + 1]))), "abr_state_out_of_bounds");
    decide(&p, &state(0, 1000, &[1; MAX_OPTIONS])).unwrap();
    assert_eq!(code(depth(&[1; MAX_OPTIONS + 1], &p, &state(0, 1000, &[]))), "abr_state_out_of_bounds");
    assert_eq!(code(depth(&[u64::MAX, 1], &p, &state(0, 1000, &[]))), "abr_state_out_of_bounds");
    let mut s = state(0, 1000, &[]);
    s.option_joules_micro = vec![Some(u64::MAX), Some(1)];
    assert_eq!(code(depth(&[1, 1], &p, &s)), "abr_state_out_of_bounds");
    // The extremes of every type.
    let big = AbrState {
        buffer_ms: u32::MAX,
        samples: (0..300).map(|i| Sample { at_ms: i, bytes: u64::MAX, ms: if i % 2 == 0 { 0 } else { u64::MAX } }).collect(),
        last: Some(1),
        option_bytes: vec![1, u64::MAX / 3, u64::MAX - 1, u64::MAX],
        option_joules_micro: vec![Some(u64::MAX), None, Some(1), Some(u64::MAX - 1)],
        joule_budget_micro: Some(u64::MAX - 1),
        segment_ms: 1,
    };
    let policies: Vec<Box<dyn Policy>> = vec![
        Box::new(Throughput { last_n: 255, safety_permille: u16::MAX }),
        Box::new(BufferBased { reservoir_ms: 0, cushion_ms: u32::MAX }),
        Box::new(BolaLike { v_milli: u32::MAX, max_buffer_ms: u32::MAX }),
        Box::new(Mpc { horizon: 255, safety_permille: u16::MAX }),
    ];
    for q in &policies {
        for seg in [0, 1, u32::MAX] {
            let st = AbrState { segment_ms: seg, ..big.clone() };
            let (k, _) = decide(q.as_ref(), &st).unwrap();
            assert!(k < 4, "{}", q.id());
        }
    }
    assert_eq!(throughput_bps(&big.samples, 255), (255u128 * u128::from(u64::MAX) * 8000) / (128 * u128::from(u64::MAX) + 127));
    assert_eq!(option_bps(u64::MAX, 1), u128::from(u64::MAX) * 8000);
}

fn appendix_d() -> Vec<u8> {
    WaiMulti::new("video", "replicate", serde_json::Value::Null, vec![
        Rendition::media("wai.video.av1", "av1", pattern(3000, 7, 3)).with_id("low").with_payload_digest(),
        Rendition::media("wai.video.av1", "av1", pattern(50_000, 13, 5)).with_id("high").with_payload_digest(),
        Rendition::companion("wai.meta.hdr10plus", pattern(200, 31, 1), "wai.video.av1").with_id("hdr")
            .with_payload_digest().with_companion_of_id("high"),
    ]).to_bytes().unwrap()
}

/// Three capabilities, two sizes of one, in table order av1-large, vp9, av1-small.
fn mixed() -> Vec<u8> {
    let mut wrong = Rendition::media("wai.video.av1", "av1", pattern(9000, 3, 1)).with_id("big").with_payload_digest();
    wrong.meta.sha256 = Some(crate::pin::sha256_hex(b"not these bytes"));
    WaiMulti::new("video", "replicate", serde_json::Value::Null, vec![
        wrong,
        Rendition::media("wai.video.vp9", "vp9", pattern(4000, 5, 1)).with_id("v").with_payload_digest(),
        Rendition::media("wai.video.av1", "av1", pattern(2000, 7, 2)).with_id("small").with_payload_digest(),
    ]).to_bytes().unwrap()
}

/// Appendix D under a throughput policy: low throughput keeps `low` first, high puts
/// `high` first; the companion stays after the options; an ineligible capability is not
/// an option.
#[test]
fn order_entries_appendix_d() {
    let h = head_of(&appendix_d());
    let all = |_: &str| true;
    let p = Throughput { last_n: 1, safety_permille: 900 };
    // `high` is 50,000 B a second: 400 kb/s.
    assert_eq!(order_entries(&h, &all, &p, &state(0, 100_000, &[])).unwrap(), vec![0, 1, 2]);
    assert_eq!(order_entries(&h, &all, &p, &state(0, 10_000_000, &[])).unwrap(), vec![1, 0, 2]);
    let only_vp9 = |c: &str| c == "wai.video.vp9";
    assert_eq!(order_entries(&h, &only_vp9, &p, &state(0, 10_000_000, &[])).unwrap(), vec![0, 1, 2]);
    let m = head_of(&mixed());
    // av1 options by size: small (2), big (0); vp9 alone.
    assert_eq!(order_entries(&m, &all, &p, &state(0, 10_000_000, &[])).unwrap(), vec![0, 2, 1]);
    assert_eq!(order_entries(&m, &all, &p, &state(0, 10, &[])).unwrap(), vec![2, 0, 1]);
}

fn fetch_with(env: &[u8], adv: &'static [&'static str], p: Box<dyn Policy>, s: AbrState, frag: Fragment) -> Result<super::session::Fetched, RangeRefusal> {
    let proof = HeadProof::build(env).unwrap().to_bytes();
    let mut sel = Selector::new(move |c| adv.contains(&c));
    sel.order = Some(selector_order(p, s, move |c| adv.contains(&c)).unwrap());
    let mut sess = FetchSession::new(content_hash(env), frag, FetchPolicy::default(), sel);
    drive(&mut sess, |w| {
        let range = match w {
            Want::HeadProof => return Some((ResponseMeta { status: 200, ..Default::default() }, proof.clone())),
            Want::Whole(_) => None,
            Want::Range(r, _) => Some(r.header()),
            Want::Rest(r, _) => Some(format!("bytes={}-", r.first)),
        };
        let a = origin::serve(env.len() as u64, range.as_deref());
        Some((ResponseMeta { status: a.status, content_range: a.content_range, content_encoding: None }, env[a.body.start as usize..a.body.end as usize].to_vec()))
    })
}

fn whole_pick(env: &[u8], adv: &[&str], order: &[usize]) -> Result<usize, Vec<(usize, Skip)>> {
    let w = WaiMulti::from_bytes(env).unwrap();
    let eligible = |_: usize, e: &RenditionMeta| if adv.contains(&e.capability.as_str()) { Ok(()) } else { Err(Skip::NotEligible) };
    let digest_ok = |i: usize, e: &RenditionMeta| e.payload_digest_matches(&w.renditions[i]);
    crate::container::select_entries(&w.manifest, order, eligible, digest_ok,
        |c: &str, _: &crate::container::Prior| Err::<(), _>(crate::container::PinRefusal::NotHeld { capability: c.into() }))
        .map(|(i, _, _)| i)
}

/// The order a policy gives is the order on both paths: for every policy, state and
/// sink, the range fetch picks what whole-envelope selection picks in that order. Under
/// high throughput Appendix D's `high` is picked, which (with its companion) is past
/// the threshold, so the object is fetched whole.
#[test]
fn same_pick_under_policies() {
    const ADV: [&[&str]; 4] = [&["wai.video.av1"], &["wai.video.vp9"], &["wai.video.av1", "wai.video.vp9", "wai.meta.hdr10plus"], &[]];
    let policies = || -> Vec<Box<dyn Policy>> {
        vec![
            Box::new(Throughput { last_n: 2, safety_permille: 900 }),
            Box::new(BufferBased { reservoir_ms: 2000, cushion_ms: 8000 }),
            Box::new(BolaLike { v_milli: 900, max_buffer_ms: 12_000 }),
            Box::new(Mpc { horizon: 3, safety_permille: 800 }),
        ]
    };
    for env in [appendix_d(), mixed()] {
        let h = head_of(&env);
        for adv in ADV {
            for (buffer, bps) in [(0u32, 50_000u64), (6000, 900_000), (20_000, 50_000_000)] {
                for (k, p) in policies().into_iter().enumerate() {
                    let s = state(buffer, bps, &[]);
                    let order = order_entries(&h, &|c| adv.contains(&c), p.as_ref(), &s).unwrap();
                    let want = whole_pick(&env, adv, &order);
                    let got = fetch_with(&env, adv, policies().remove(k), s, Fragment { h: Some(1009), n: Some(env.len() as u64), ..Default::default() });
                    match (&want, &got) {
                        (Ok(i), Ok(f)) => assert_eq!(*i, f.index, "{adv:?} {buffer} {bps} {}", p.id()),
                        (Err(_), Err(RangeRefusal::Selection(_))) => {}
                        _ => panic!("paths differ: whole {want:?}, range {got:?}"),
                    }
                }
            }
        }
    }
    let f = fetch_with(&appendix_d(), &["wai.video.av1"], Box::new(Throughput { last_n: 1, safety_permille: 900 }), state(0, 10_000_000, &[]),
        Fragment { h: Some(1009), n: Some(53_794), ..Default::default() }).unwrap();
    assert_eq!((f.index, f.verified_by), (1, VerifiedBy::WholeObject));
}
