#include "batch_reassembly_logic.h"

#include <string.h>

void batch_reassembly_reset(BatchReassemblyState &state) {
  state.active = false;
  state.batch_id = 0;
  state.next_index = 0;
  state.expected_chunks = 0;
  state.total_len = 0;
  state.received_len = 0;
  state.last_rx_ms = 0;
  state.timeout_ms = 0;
}

BatchReassemblyStatus batch_reassembly_push(BatchReassemblyState &state, uint16_t batch_id, uint8_t chunk_index,
                                            uint8_t chunk_count, uint16_t total_len, const uint8_t *chunk_data,
                                            size_t chunk_len, uint32_t now_ms, uint32_t timeout_ms_for_new_batch,
                                            uint16_t max_total_len, uint8_t *buffer, size_t buffer_cap,
                                            uint16_t &out_complete_len) {
  out_complete_len = 0;
  if (!buffer || !chunk_data) {
    batch_reassembly_reset(state);
    return BatchReassemblyStatus::ErrorReset;
  }
  if (chunk_len > 0 && total_len == 0) {
    batch_reassembly_reset(state);
    return BatchReassemblyStatus::ErrorReset;
  }

  bool expired = state.timeout_ms > 0 && (now_ms - state.last_rx_ms > state.timeout_ms);
  if (!state.active || batch_id != state.batch_id || expired) {
    if (chunk_index != 0) {
      batch_reassembly_reset(state);
      return BatchReassemblyStatus::ErrorReset;
    }
    if (total_len == 0 || total_len > max_total_len || chunk_count == 0) {
      batch_reassembly_reset(state);
      return BatchReassemblyStatus::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) {
    batch_reassembly_reset(state);
    return BatchReassemblyStatus::ErrorReset;
  }

  if (state.received_len + chunk_len > state.total_len ||
      state.received_len + chunk_len > max_total_len ||
      state.received_len + chunk_len > buffer_cap) {
    batch_reassembly_reset(state);
    return BatchReassemblyStatus::ErrorReset;
  }

  memcpy(&buffer[state.received_len], chunk_data, chunk_len);
  state.received_len += static_cast<uint16_t>(chunk_len);
  state.next_index++;
  state.last_rx_ms = now_ms;

  if (state.next_index == state.expected_chunks && state.received_len == state.total_len) {
    out_complete_len = state.received_len;
    batch_reassembly_reset(state);
    return BatchReassemblyStatus::Complete;
  }

  return BatchReassemblyStatus::InProgress;
}