use serde::de::{IgnoredAny, SeqAccess};

use super::JsonTreeBuilder;
use super::SampledArray;

pub(crate) fn sample_stream<'de, A>(
    seq: &mut A,
    builder: &JsonTreeBuilder,
    cap: usize,
) -> Result<SampledArray, A::Error>
where
    A: SeqAccess<'de>,
{
    if cap == 0 {
        let total = drain_len(seq)?;
        return Ok(SampledArray {
            children: Vec::new(),
            indices: Vec::new(),
            total_len: total,
        });
    }
    let k = cap;
    let mut ring_idx: Vec<usize> = vec![0; k];
    let mut ring_children: Vec<usize> = vec![0; k];
    let mut count = 0usize;
    let mut head = 0usize;
    loop {
        let seed = builder.seed();
        match seq.next_element_seed(seed)? {
            Some(child_id) => {
                ring_idx[head] = count;
                ring_children[head] = child_id;
                head = if head + 1 == k { 0 } else { head + 1 };
                count = count.saturating_add(1);
            }
            None => break,
        }
    }
    Ok(materialize_tail(&ring_idx, &ring_children, count, head, k))
}

fn drain_len<'de, A>(seq: &mut A) -> Result<usize, A::Error>
where
    A: SeqAccess<'de>,
{
    let mut total = 0usize;
    while (seq.next_element::<IgnoredAny>()?).is_some() {
        total += 1;
    }
    Ok(total)
}

fn materialize_tail(
    ring_idx: &[usize],
    ring_children: &[usize],
    count: usize,
    head: usize,
    k: usize,
) -> SampledArray {
    let kept = count.min(k);
    if kept == 0 {
        return SampledArray {
            children: Vec::new(),
            indices: Vec::new(),
            total_len: count,
        };
    }
    let start = if count >= k { head } else { 0 };
    let mut children = Vec::with_capacity(kept);
    let mut indices = Vec::with_capacity(kept);
    for i in 0..kept {
        let pos = (start + i) % k;
        indices.push(ring_idx[pos]);
        children.push(ring_children[pos]);
    }
    SampledArray {
        children,
        indices,
        total_len: count,
    }
}

#[cfg(test)]
mod tests {
    use crate::order::PriorityConfig;

    #[test]
    fn tail_sampler_keeps_last_n_indices() {
        let input = b"[0,1,2,3,4,5,6,7,8,9]".to_vec();
        let mut cfg = PriorityConfig::new(usize::MAX, 5);
        cfg.array_sampler = crate::ArraySamplerStrategy::Tail;
        let arena =
            crate::ingest::formats::json::build_json_tree_arena_from_bytes(
                input, &cfg,
            )
            .expect("arena");
        let root = &arena.nodes[arena.root_id];
        assert_eq!(root.children_len, 5, "kept 5");
        let mut orig_indices = Vec::new();
        for i in 0..root.children_len {
            let oi = if root.arr_indices_len > 0 {
                arena.arr_indices[root.arr_indices_start + i]
            } else {
                i
            };
            orig_indices.push(oi);
        }
        assert_eq!(orig_indices, vec![5, 6, 7, 8, 9]);
    }
}
