//! Zero-allocation byte extent calculator for serialized q IPC values.
//!
//! Given a byte slice and an offset pointing to the start of a serialized q
//! value, [`value_byte_extent`] returns how many bytes that value occupies
//! without allocating memory or constructing a [`Value`].  This is used by
//! the parallel column decoder to split a table's column data into
//! independent sub-slices before dispatching them to worker threads.

use crate::error::CoreError;
use crate::error::CoreResult;
use crate::protocol::Primitive;
use crate::protocol::Shape;
use crate::protocol::TypeCode;

/// Returns the byte extent of a serialized q value starting at `bytes[offset..]`.
///
/// The function reads only type codes, attributes, and lengths — it never
/// allocates or constructs a `Value`.  For fixed-width vectors this is O(1);
/// for symbol vectors and nested structures it scans forward.
pub fn value_byte_extent(bytes: &[u8], offset: usize) -> CoreResult<usize> {
    if offset >= bytes.len() {
        return Err(CoreError::InvalidStructure(format!(
            "extent: offset {offset} beyond buffer length {}",
            bytes.len()
        )));
    }

    let type_code = TypeCode::try_from(bytes[offset] as i8)?;
    let shape = type_code.shape();

    match shape {
        Shape::Atom => atom_extent(bytes, offset, type_code),
        Shape::Vector => vector_extent(bytes, offset, type_code),
        Shape::List => list_extent(bytes, offset),
        Shape::Dictionary => dictionary_extent(bytes, offset),
        Shape::Table => table_extent(bytes, offset),
        Shape::UnaryPrimitive => {
            // type byte + opcode byte
            check_available(bytes, offset, 2)?;
            Ok(2)
        }
        Shape::Error => {
            check_available(bytes, offset, 1)?;
            let data_start = offset + 1;
            let pos = bytes[data_start..]
                .iter()
                .position(|&b| b == 0)
                .ok_or_else(|| {
                    CoreError::InvalidStructure(format!(
                        "extent: unterminated error string at offset {offset}"
                    ))
                })?;
            Ok(1 + pos + 1)
        }
    }
}

/// Checks that at least `need` bytes are available from `offset`.
#[inline]
fn check_available(bytes: &[u8], offset: usize, need: usize) -> CoreResult<()> {
    if offset + need > bytes.len() {
        Err(CoreError::InvalidStructure(format!(
            "extent: need {need} bytes at offset {offset}, but buffer length is {}",
            bytes.len()
        )))
    } else {
        Ok(())
    }
}

/// Reads an i32 length field at `bytes[offset..offset+4]` (little-endian).
#[inline]
fn read_len(bytes: &[u8], offset: usize) -> CoreResult<usize> {
    check_available(bytes, offset, 4)?;
    let len = i32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap());
    if len < 0 {
        return Err(CoreError::InvalidStructure(format!(
            "extent: negative length {len} at offset {offset}"
        )));
    }
    Ok(len as usize)
}

fn atom_extent(bytes: &[u8], offset: usize, type_code: TypeCode) -> CoreResult<usize> {
    // 1 byte for type code + data bytes
    let primitive = type_code
        .primitive()
        .ok_or(CoreError::InvalidTypeCode(type_code as i8))?;

    if let Some(width) = primitive.width() {
        check_available(bytes, offset, 1 + width)?;
        Ok(1 + width)
    } else {
        // Symbol atom: scan for null terminator
        debug_assert_eq!(primitive, Primitive::Symbol);
        let data_start = offset + 1;
        let pos = bytes[data_start..]
            .iter()
            .position(|&b| b == 0)
            .ok_or_else(|| {
                CoreError::InvalidStructure(format!(
                    "extent: unterminated symbol atom at offset {offset}"
                ))
            })?;
        // type byte + symbol bytes + null terminator
        Ok(1 + pos + 1)
    }
}

fn vector_extent(bytes: &[u8], offset: usize, type_code: TypeCode) -> CoreResult<usize> {
    // Header: 1 (type) + 1 (attribute) + 4 (length) = 6 bytes
    const HEADER: usize = 6;
    check_available(bytes, offset, HEADER)?;
    let length = read_len(bytes, offset + 2)?;

    let primitive = type_code
        .primitive()
        .ok_or(CoreError::InvalidTypeCode(type_code as i8))?;

    if let Some(width) = primitive.width() {
        let data_bytes = length
            .checked_mul(width)
            .ok_or(CoreError::LengthOverflow(length))?;
        check_available(bytes, offset, HEADER + data_bytes)?;
        Ok(HEADER + data_bytes)
    } else {
        // Symbol vector: scan through `length` null-terminated strings
        debug_assert_eq!(primitive, Primitive::Symbol);
        let mut scan = offset + HEADER;
        for _ in 0..length {
            let pos = bytes[scan..].iter().position(|&b| b == 0).ok_or_else(|| {
                CoreError::InvalidStructure(format!(
                    "extent: unterminated symbol in vector at offset {scan}"
                ))
            })?;
            scan += pos + 1; // skip past the null terminator
        }
        Ok(scan - offset)
    }
}

fn list_extent(bytes: &[u8], offset: usize) -> CoreResult<usize> {
    // Header: 1 (type) + 1 (attribute) + 4 (length) = 6 bytes
    const HEADER: usize = 6;
    check_available(bytes, offset, HEADER)?;
    let length = read_len(bytes, offset + 2)?;

    let mut scan = offset + HEADER;
    for _ in 0..length {
        let child_extent = value_byte_extent(bytes, scan)?;
        scan += child_extent;
    }
    Ok(scan - offset)
}

fn dictionary_extent(bytes: &[u8], offset: usize) -> CoreResult<usize> {
    // 1 byte for type code (99 or 127), then keys value, then values value
    check_available(bytes, offset, 1)?;
    let keys_extent = value_byte_extent(bytes, offset + 1)?;
    let values_extent = value_byte_extent(bytes, offset + 1 + keys_extent)?;
    Ok(1 + keys_extent + values_extent)
}

fn table_extent(bytes: &[u8], offset: usize) -> CoreResult<usize> {
    // 1 byte type code + 1 byte attribute + inner dictionary
    check_available(bytes, offset, 2)?;
    let dict_extent = value_byte_extent(bytes, offset + 2)?;
    Ok(2 + dict_extent)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::decode::decode_value;
    use crate::encode::encode_value;
    use crate::frame::Encoding;
    use crate::value::*;

    /// Helper: encode a value, then verify extent equals encoded body length.
    fn assert_extent_matches(value: &Value) {
        let body = encode_value(value, Encoding::LittleEndian).unwrap();
        let extent = value_byte_extent(&body, 0).unwrap();
        assert_eq!(
            extent,
            body.len(),
            "extent mismatch for {value:?}: expected {}, got {extent}",
            body.len()
        );
    }

    // -- Atoms --

    #[test]
    fn extent_boolean_atom() {
        assert_extent_matches(&Value::Atom(Atom::Boolean(true)));
    }

    #[test]
    fn extent_byte_atom() {
        assert_extent_matches(&Value::Atom(Atom::Byte(0x42)));
    }

    #[test]
    fn extent_short_atom() {
        assert_extent_matches(&Value::Atom(Atom::Short(42)));
    }

    #[test]
    fn extent_int_atom() {
        assert_extent_matches(&Value::Atom(Atom::Int(42)));
    }

    #[test]
    fn extent_long_atom() {
        assert_extent_matches(&Value::Atom(Atom::Long(42)));
    }

    #[test]
    fn extent_real_atom() {
        assert_extent_matches(&Value::Atom(Atom::Real(1.5)));
    }

    #[test]
    fn extent_float_atom() {
        assert_extent_matches(&Value::Atom(Atom::Float(1.5)));
    }

    #[test]
    fn extent_char_atom() {
        assert_extent_matches(&Value::Atom(Atom::Char(b'c')));
    }

    #[test]
    fn extent_symbol_atom() {
        assert_extent_matches(&Value::Atom(Atom::Symbol(bytes::Bytes::from_static(
            b"hello",
        ))));
    }

    #[test]
    fn extent_empty_symbol_atom() {
        assert_extent_matches(&Value::Atom(Atom::Symbol(bytes::Bytes::from_static(b""))));
    }

    #[test]
    fn extent_guid_atom() {
        assert_extent_matches(&Value::Atom(Atom::Guid([0u8; 16])));
    }

    #[test]
    fn extent_timestamp_atom() {
        assert_extent_matches(&Value::Atom(Atom::Timestamp(1)));
    }

    #[test]
    fn extent_month_atom() {
        assert_extent_matches(&Value::Atom(Atom::Month(1)));
    }

    #[test]
    fn extent_date_atom() {
        assert_extent_matches(&Value::Atom(Atom::Date(1)));
    }

    #[test]
    fn extent_datetime_atom() {
        assert_extent_matches(&Value::Atom(Atom::Datetime(1.5)));
    }

    #[test]
    fn extent_timespan_atom() {
        assert_extent_matches(&Value::Atom(Atom::Timespan(1)));
    }

    #[test]
    fn extent_minute_atom() {
        assert_extent_matches(&Value::Atom(Atom::Minute(1)));
    }

    #[test]
    fn extent_second_atom() {
        assert_extent_matches(&Value::Atom(Atom::Second(1)));
    }

    #[test]
    fn extent_time_atom() {
        assert_extent_matches(&Value::Atom(Atom::Time(1)));
    }

    // -- Vectors --

    #[test]
    fn extent_int_vector() {
        assert_extent_matches(&Value::Vector(Vector::new(
            crate::protocol::Attribute::None,
            VectorData::from_i32s(&[1, 2, 3]),
        )));
    }

    #[test]
    fn extent_empty_int_vector() {
        assert_extent_matches(&Value::Vector(Vector::new(
            crate::protocol::Attribute::None,
            VectorData::from_i32s(&[]),
        )));
    }

    #[test]
    fn extent_symbol_vector() {
        assert_extent_matches(&Value::Vector(Vector::new(
            crate::protocol::Attribute::None,
            VectorData::Symbol(vec![
                bytes::Bytes::from_static(b"alpha"),
                bytes::Bytes::from_static(b"beta"),
            ]),
        )));
    }

    #[test]
    fn extent_empty_symbol_vector() {
        assert_extent_matches(&Value::Vector(Vector::new(
            crate::protocol::Attribute::None,
            VectorData::Symbol(vec![]),
        )));
    }

    #[test]
    fn extent_boolean_vector() {
        assert_extent_matches(&Value::Vector(Vector::new(
            crate::protocol::Attribute::None,
            VectorData::Boolean(bytes::Bytes::from_static(&[1, 0, 1])),
        )));
    }

    #[test]
    fn extent_guid_vector() {
        assert_extent_matches(&Value::Vector(Vector::new(
            crate::protocol::Attribute::None,
            VectorData::from_guids(&[[0u8; 16], [1u8; 16]]),
        )));
    }

    #[test]
    fn extent_long_vector() {
        assert_extent_matches(&Value::Vector(Vector::new(
            crate::protocol::Attribute::None,
            VectorData::from_i64s(&[1, 2, 3]),
        )));
    }

    #[test]
    fn extent_float_vector() {
        assert_extent_matches(&Value::Vector(Vector::new(
            crate::protocol::Attribute::None,
            VectorData::from_f64s(&[1.0, 2.0]),
        )));
    }

    #[test]
    fn extent_char_vector() {
        assert_extent_matches(&Value::Vector(Vector::new(
            crate::protocol::Attribute::None,
            VectorData::Char(bytes::Bytes::from_static(b"hello")),
        )));
    }

    #[test]
    fn extent_byte_vector() {
        assert_extent_matches(&Value::Vector(Vector::new(
            crate::protocol::Attribute::None,
            VectorData::Byte(bytes::Bytes::from(vec![1, 2, 3])),
        )));
    }

    // -- Composites --

    #[test]
    fn extent_general_list() {
        assert_extent_matches(&Value::List(List::new(
            crate::protocol::Attribute::None,
            vec![
                Value::Atom(Atom::Int(42)),
                Value::Atom(Atom::Symbol(bytes::Bytes::from_static(b"ab"))),
            ],
        )));
    }

    #[test]
    fn extent_empty_list() {
        assert_extent_matches(&Value::List(List::new(
            crate::protocol::Attribute::None,
            vec![],
        )));
    }

    #[test]
    fn extent_dictionary() {
        assert_extent_matches(&Value::Dictionary(Dictionary::new(
            false,
            Value::Vector(Vector::new(
                crate::protocol::Attribute::None,
                VectorData::Symbol(vec![
                    bytes::Bytes::from_static(b"a"),
                    bytes::Bytes::from_static(b"b"),
                ]),
            )),
            Value::Vector(Vector::new(
                crate::protocol::Attribute::None,
                VectorData::from_i32s(&[1, 2]),
            )),
        )));
    }

    #[test]
    fn extent_sorted_dictionary() {
        assert_extent_matches(&Value::Dictionary(Dictionary::new(
            true,
            Value::Vector(Vector::new(
                crate::protocol::Attribute::Sorted,
                VectorData::Symbol(vec![
                    bytes::Bytes::from_static(b"a"),
                    bytes::Bytes::from_static(b"b"),
                ]),
            )),
            Value::Vector(Vector::new(
                crate::protocol::Attribute::None,
                VectorData::from_i32s(&[1, 2]),
            )),
        )));
    }

    #[test]
    fn extent_table() {
        assert_extent_matches(&Value::Table(Table::new(
            crate::protocol::Attribute::None,
            vec![
                bytes::Bytes::from_static(b"sym"),
                bytes::Bytes::from_static(b"px"),
            ],
            vec![
                Value::Vector(Vector::new(
                    crate::protocol::Attribute::None,
                    VectorData::Symbol(vec![
                        bytes::Bytes::from_static(b"a"),
                        bytes::Bytes::from_static(b"b"),
                    ]),
                )),
                Value::Vector(Vector::new(
                    crate::protocol::Attribute::None,
                    VectorData::from_i32s(&[10, 20]),
                )),
            ],
        )));
    }

    #[test]
    fn extent_nested_list() {
        assert_extent_matches(&Value::List(List::new(
            crate::protocol::Attribute::None,
            vec![
                Value::Vector(Vector::new(
                    crate::protocol::Attribute::None,
                    VectorData::from_i32s(&[1, 2, 3]),
                )),
                Value::Vector(Vector::new(
                    crate::protocol::Attribute::None,
                    VectorData::from_i32s(&[4, 5]),
                )),
            ],
        )));
    }

    #[test]
    fn extent_unary_primitive() {
        let value = Value::UnaryPrimitive { opcode: 42 };
        assert_extent_matches(&value);
    }

    /// Verify extent matches for every value encoded in a real roundtrip body.
    #[test]
    fn extent_matches_decode_consumption() {
        // Encode a table, get the body, verify extent == body.len()
        let table = Value::Table(Table::new(
            crate::protocol::Attribute::None,
            vec![
                bytes::Bytes::from_static(b"a"),
                bytes::Bytes::from_static(b"b"),
                bytes::Bytes::from_static(b"c"),
            ],
            vec![
                Value::Vector(Vector::new(
                    crate::protocol::Attribute::None,
                    VectorData::from_i32s(&[1, 2, 3]),
                )),
                Value::Vector(Vector::new(
                    crate::protocol::Attribute::None,
                    VectorData::Symbol(vec![
                        bytes::Bytes::from_static(b"x"),
                        bytes::Bytes::from_static(b"y"),
                        bytes::Bytes::from_static(b"z"),
                    ]),
                )),
                Value::Vector(Vector::new(
                    crate::protocol::Attribute::None,
                    VectorData::from_f64s(&[1.0, 2.0, 3.0]),
                )),
            ],
        ));
        let body = encode_value(&table, Encoding::LittleEndian).unwrap();
        let extent = value_byte_extent(&body, 0).unwrap();
        assert_eq!(extent, body.len());

        // Also verify roundtrip
        let decoded =
            decode_value(bytes::Bytes::from(body.clone()), Encoding::LittleEndian).unwrap();
        assert_eq!(decoded, table);
    }
}
