#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ReassemblyState {
    pub active: bool,
    pub batch_id: u16,
    pub next_index: u8,
    pub expected_chunks: u8,
    pub total_len: u16,
    pub received_len: u16,
    pub last_rx_ms: u32,
    pub timeout_ms: u32,
}

impl Default for ReassemblyState {
    fn default() -> Self {
        Self {
            active: false,
            batch_id: 0,
            next_index: 0,
            expected_chunks: 0,
            total_len: 0,
            received_len: 0,
            last_rx_ms: 0,
            timeout_ms: 0,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReassemblyStatus {
    InProgress,
    Complete { complete_len: u16 },
    ErrorReset,
}

pub fn reset_reassembly(state: &mut ReassemblyState) {
    *state = ReassemblyState::default();
}

#[allow(clippy::too_many_arguments)]
pub fn push_chunk(
    state: &mut ReassemblyState,
    batch_id: u16,
    chunk_index: u8,
    chunk_count: u8,
    total_len: u16,
    chunk_data: &[u8],
    now_ms: u32,
    timeout_ms_for_new_batch: u32,
    max_total_len: u16,
    buffer: &mut [u8],
) -> ReassemblyStatus {
    if chunk_data.len() > 0 && total_len == 0 {
        reset_reassembly(state);
        return ReassemblyStatus::ErrorReset;
    }

    let expired = state.timeout_ms > 0
        && now_ms.wrapping_sub(state.last_rx_ms) > state.timeout_ms;

    if !state.active || batch_id != state.batch_id || expired {
        if chunk_index != 0 {
            reset_reassembly(state);
            return ReassemblyStatus::ErrorReset;
        }
        if total_len == 0 || total_len > max_total_len || chunk_count == 0 {
            reset_reassembly(state);
            return ReassemblyStatus::ErrorReset;
        }

        state.active = true;
        state.batch_id = batch_id;
        state.expected_chunks = chunk_count;
        state.total_len = total_len;
        state.received_len = 0;
        state.next_index = 0;
        state.last_rx_ms = now_ms;
        state.timeout_ms = timeout_ms_for_new_batch;
    }

    if !state.active || chunk_index != state.next_index || chunk_count != state.expected_chunks {
        reset_reassembly(state);
        return ReassemblyStatus::ErrorReset;
    }

    let next_received = state.received_len as usize + chunk_data.len();
    if next_received > state.total_len as usize
        || next_received > max_total_len as usize
        || next_received > buffer.len()
    {
        reset_reassembly(state);
        return ReassemblyStatus::ErrorReset;
    }

    let start = state.received_len as usize;
    let end = start + chunk_data.len();
    buffer[start..end].copy_from_slice(chunk_data);

    state.received_len = next_received as u16;
    state.next_index = state.next_index.wrapping_add(1);
    state.last_rx_ms = now_ms;

    if state.next_index == state.expected_chunks && state.received_len == state.total_len {
        let complete_len = state.received_len;
        reset_reassembly(state);
        return ReassemblyStatus::Complete { complete_len };
    }

    ReassemblyStatus::InProgress
}