mlt-core 0.12.7

use arbitrary::Error::IncorrectFormat;
use arbitrary::{Arbitrary, Result, Unstructured};

use crate::encoder::model::StagedLayer;
use crate::encoder::optimizer::Presence;
use crate::encoder::{EncoderConfig, StagedId, StagedProperty, StagedSharedDict, StagedStrings};

impl Arbitrary<'_> for EncoderConfig {
    fn arbitrary(u: &mut Unstructured<'_>) -> Result<Self> {
        Ok(Self::default()
            .with_wire_version(u.arbitrary()?)
            .with_tessellation(u.arbitrary()?)
            .with_spatial_morton_sort(u.arbitrary()?)
            .with_spatial_hilbert_sort(u.arbitrary()?)
            .with_id_sort(u.arbitrary()?)
            .with_fsst(u.arbitrary()?)
            .with_fastpfor(u.arbitrary()?)
            .with_shared_dict(u.arbitrary()?))
    }
}

impl Arbitrary<'_> for StagedId {
    fn arbitrary(u: &mut Unstructured<'_>) -> Result<Self> {
        // Bound ID count to prevent OOM from unbounded vector generation
        let count = usize::from(u.int_in_range(0..=64u8)?);
        let values: Vec<Option<u64>> = (0..count).map(|_| u.arbitrary()).collect::<Result<_>>()?;
        Ok(Self::from_optional(values))
    }
}

impl Arbitrary<'_> for StagedLayer {
    fn arbitrary(u: &mut Unstructured<'_>) -> Result<Self> {
        // Bound name length to prevent OOM from unbounded string generation
        let len = u.int_in_range(1..=32)?;
        let name: String = (0..len)
            .map(|_| u.arbitrary::<char>())
            .collect::<Result<_>>()?;
        let extent: u32 = u.arbitrary()?;
        // Generate geometry first -- its feature count drives ID and property columns.
        let geometry: crate::decoder::GeometryValues = u.arbitrary()?;
        let fc = geometry.vector_types().len();
        let id = if u.arbitrary::<bool>()? {
            let ids: Vec<Option<u64>> = (0..fc)
                .map(|_| -> Result<_> {
                    if u.arbitrary::<bool>()? {
                        Ok(Some(u.arbitrary::<u64>()?))
                    } else {
                        Ok(None)
                    }
                })
                .collect::<Result<_>>()?;
            StagedId::from_optional(ids)
        } else {
            StagedId::None
        };
        // Bound property count to prevent OOM from unbounded vector generation.
        // Each column must have exactly `fc` values to match the feature count.
        let prop_count = usize::from(u.int_in_range(0..=4u8)?);
        let properties: Vec<StagedProperty> = (0..prop_count)
            .map(|i| {
                let values: Vec<Option<u32>> =
                    (0..fc).map(|_| u.arbitrary()).collect::<Result<_>>()?;
                Ok(StagedProperty::opt_u32(format!("prop{i}"), values))
            })
            .collect::<Result<_>>()?;

        Self::new(name, extent, id, geometry, properties).map_err(|_| IncorrectFormat)
    }
}

impl<'a> Arbitrary<'a> for StagedSharedDict {
    fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
        // Bound item count and string sizes to prevent OOM
        let item_count = usize::from(u.int_in_range(0..=8u8)?);
        let items_raw: Vec<(String, Vec<Option<String>>, Presence)> = (0..item_count)
            .map(|_| {
                let values = generate_strings(u)?;
                let presence = if values.iter().all(Option::is_some) {
                    Presence::AllPresent
                } else {
                    Presence::Mixed
                };
                Ok((bounded_string(u, 32)?, values, presence))
            })
            .collect::<Result<_>>()?;
        if items_raw.is_empty() {
            return Ok(Self {
                prefix: bounded_string(u, 32)?,
                data: String::new(),
                items: Vec::new(),
            });
        }
        let prefix = bounded_string(u, 32)?;
        Self::new(prefix, items_raw).map_err(|_| IncorrectFormat)
    }
}

impl Arbitrary<'_> for StagedProperty {
    fn arbitrary(u: &mut Unstructured<'_>) -> Result<Self> {
        // Bound value count to prevent OOM from unbounded vector generation
        let count = usize::from(u.int_in_range(0..=64u8)?);
        let values: Vec<Option<u32>> = (0..count).map(|_| u.arbitrary()).collect::<Result<_>>()?;
        Ok(Self::opt_u32("prop", values))
    }
}

impl Arbitrary<'_> for StagedStrings {
    fn arbitrary(u: &mut Unstructured<'_>) -> Result<Self> {
        Ok(Self::from_optional(
            bounded_string(u, 32)?,
            generate_strings(u)?,
        ))
    }
}

/// Generate a string with bounded length to prevent OOM from unbounded string generation.
pub fn bounded_string(u: &mut Unstructured<'_>, max_len: u8) -> Result<String> {
    let len = usize::from(u.int_in_range(0..=max_len)?);
    (0..len)
        .map(|_| u.arbitrary::<char>())
        .collect::<Result<_>>()
}

fn generate_strings(u: &mut Unstructured) -> Result<Vec<Option<String>>> {
    // Bound string count and individual string lengths to prevent OOM
    let val_count = usize::from(u.int_in_range(0..=16u8)?);
    let values: Vec<Option<String>> = (0..val_count)
        .map(|_| -> Result<_> {
            if u.arbitrary()? {
                Ok(Some(bounded_string(u, 64)?))
            } else {
                Ok(None)
            }
        })
        .collect::<Result<_>>()?;
    Ok(values)
}