#include "protocol.h"

// ============================================================================
// Global Buffers
// ============================================================================

uint8_t g_rxBuffer[MAX_PAYLOAD_SIZE + PACKET_HEADER_SIZE + PACKET_CRC_SIZE];
uint16_t g_rxBufferLen = 0;

// Parsed packet info
static char s_rxTag[4];
static uint16_t s_rxPayloadLen = 0;
static uint16_t s_rxSeq = 0;
static uint16_t s_rxPayloadOffset = 0;

// Sequence counters (per-tag would be ideal, but global is simpler)
static uint16_t s_txSeq = 0;

// Receive state machine
enum RxState { RX_SYNC0, RX_SYNC1, RX_HEADER, RX_PAYLOAD, RX_CRC };
static RxState s_rxState = RX_SYNC0;
static uint16_t s_rxIdx = 0;
static uint16_t s_rxExpectedLen = 0;

// ============================================================================
// CRC16-CCITT (polynomial 0x1021, init 0xFFFF)
// ============================================================================

uint16_t crc16Update(uint16_t crc, const uint8_t* data, uint16_t len) {
  for (uint16_t i = 0; i < len; i++) {
    crc ^= (uint16_t)data[i] << 8;
    for (uint8_t j = 0; j < 8; j++) {
      if (crc & 0x8000) {
        crc = (crc << 1) ^ 0x1021;
      } else {
        crc <<= 1;
      }
    }
  }
  return crc;
}

uint16_t crc16Compute(const uint8_t* data, uint16_t len) {
  return crc16Update(0xFFFF, data, len);
}

// ============================================================================
// Packet Sending
// ============================================================================

void sendPacket(const char tag[4], const uint8_t* payload, uint16_t len) {
  // Build header
  uint8_t header[PACKET_HEADER_SIZE];
  header[0] = SYNC0;
  header[1] = SYNC1;
  header[2] = tag[0];
  header[3] = tag[1];
  header[4] = tag[2];
  header[5] = tag[3];
  header[6] = len & 0xFF;
  header[7] = (len >> 8) & 0xFF;
  header[8] = s_txSeq & 0xFF;
  header[9] = (s_txSeq >> 8) & 0xFF;
  
  // Compute CRC over tag + len + seq + payload
  uint16_t crc = 0xFFFF;
  crc = crc16Update(crc, header + 2, 8);  // tag(4) + len(2) + seq(2)
  if (len > 0 && payload != nullptr) {
    crc = crc16Update(crc, payload, len);
  }
  
  // Send
  Serial.write(header, PACKET_HEADER_SIZE);
  if (len > 0 && payload != nullptr) {
    Serial.write(payload, len);
  }
  Serial.write(crc & 0xFF);
  Serial.write((crc >> 8) & 0xFF);
  
  s_txSeq++;
}

void sendMessage(const String& msg) {
  sendPacket(Tag::MSGE, (const uint8_t*)msg.c_str(), msg.length());
}

void sendAck(const char originalTag[4]) {
  uint8_t payload[4];
  memcpy(payload, originalTag, 4);
  sendPacket(Tag::ACK, payload, 4);
}

void sendNack(const char originalTag[4], const String& reason) {
  uint8_t payload[4 + 64];
  memcpy(payload, originalTag, 4);
  uint16_t len = 4;
  
  if (reason.length() > 0) {
    uint16_t reasonLen = min((uint16_t)reason.length(), (uint16_t)60);
    memcpy(payload + 4, reason.c_str(), reasonLen);
    len += reasonLen;
  }
  
  sendPacket(Tag::NACK, payload, len);
}

// ============================================================================
// Packet Receiving
// ============================================================================

bool receivePacket() {
  while (Serial.available()) {
    uint8_t b = Serial.read();
    
    switch (s_rxState) {
      case RX_SYNC0:
        if (b == SYNC0) {
          g_rxBuffer[0] = b;
          s_rxState = RX_SYNC1;
        }
        break;
        
      case RX_SYNC1:
        if (b == SYNC1) {
          g_rxBuffer[1] = b;
          s_rxIdx = 2;
          s_rxState = RX_HEADER;
        } else if (b == SYNC0) {
          // Stay in SYNC1 state, might be repeated sync
        } else {
          s_rxState = RX_SYNC0;
        }
        break;
        
      case RX_HEADER:
        g_rxBuffer[s_rxIdx++] = b;
        if (s_rxIdx >= PACKET_HEADER_SIZE) {
          // Parse header
          memcpy(s_rxTag, g_rxBuffer + 2, 4);
          s_rxPayloadLen = g_rxBuffer[6] | (g_rxBuffer[7] << 8);
          s_rxSeq = g_rxBuffer[8] | (g_rxBuffer[9] << 8);
          s_rxPayloadOffset = PACKET_HEADER_SIZE;
          
          // Validate length
          if (s_rxPayloadLen > MAX_PAYLOAD_SIZE) {
            Serial.println("Packet too large");
            s_rxState = RX_SYNC0;
            break;
          }
          
          s_rxExpectedLen = PACKET_HEADER_SIZE + s_rxPayloadLen + PACKET_CRC_SIZE;
          
          if (s_rxPayloadLen == 0) {
            s_rxState = RX_CRC;
          } else {
            s_rxState = RX_PAYLOAD;
          }
        }
        break;
        
      case RX_PAYLOAD:
        g_rxBuffer[s_rxIdx++] = b;
        if (s_rxIdx >= PACKET_HEADER_SIZE + s_rxPayloadLen) {
          s_rxState = RX_CRC;
        }
        break;
        
      case RX_CRC:
        g_rxBuffer[s_rxIdx++] = b;
        if (s_rxIdx >= s_rxExpectedLen) {
          // Verify CRC
          uint16_t receivedCrc = g_rxBuffer[s_rxIdx - 2] | (g_rxBuffer[s_rxIdx - 1] << 8);
          uint16_t computedCrc = crc16Compute(g_rxBuffer + 2, s_rxPayloadLen + 8);
          
          s_rxState = RX_SYNC0;
          g_rxBufferLen = s_rxIdx;
          
          if (receivedCrc == computedCrc) {
            return true;  // Valid packet ready
          } else {
            Serial.println("CRC mismatch");
          }
        }
        break;
    }
  }
  
  return false;
}

const char* getReceivedTag() {
  return s_rxTag;
}

const uint8_t* getReceivedPayload() {
  return g_rxBuffer + s_rxPayloadOffset;
}

uint16_t getReceivedPayloadLen() {
  return s_rxPayloadLen;
}

uint16_t getReceivedSeq() {
  return s_rxSeq;
}

bool tagMatches(const char* received, const char expected[4]) {
  return memcmp(received, expected, 4) == 0;
}