#include <fcntl.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mmap.h>
#include <unistd.h>

#define EIM_PHYS_BASE 0x08000000
#define EIM_MEM_SIZE 0x00002000 // Need at least 8KB

// Offsets in 16-bit words
#define TX_BUFFER_OFFSET 0x0000
#define RX_BUFFER_OFFSET 0x0800
#define CONTROL_REG_OFFSET 0x1000
#define STATUS_REG_OFFSET 0x1001

int main(int argc, char *argv[]) {
  if (argc < 3) {
    printf("Usage: %s <input_codeword.bin> <output_codeword.bin>\n", argv[0]);
    return 1;
  }

  int fd_mem = open("/dev/mem", O_RDWR | O_SYNC);
  if (fd_mem < 0) {
    perror("Failed to open /dev/mem (Requires root)");
    return 1;
  }

  // Map the EIM physical space
  void *map_base = mmap(NULL, EIM_MEM_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED,
                        fd_mem, EIM_PHYS_BASE);
  if (map_base == MAP_FAILED) {
    perror("mmap failed");
    close(fd_mem);
    return 1;
  }

  volatile uint16_t *eim_bus = (volatile uint16_t *)map_base;

  // 1. Read input codeword
  FILE *f_in = fopen(argv[1], "rb");
  if (!f_in) {
    perror("Failed to open input file");
    return 1;
  }

  uint16_t tx_buffer[1152];
  size_t read_bytes = fread(tx_buffer, 1, 2304, f_in);
  fclose(f_in);
  if (read_bytes != 2304) {
    printf("Warning: Input file was not 2304 bytes. Read %zu\n", read_bytes);
  }

  printf("Writing to FPGA TX Buffer (EIM)...\n");
  for (int i = 0; i < 1152; i++) {
    // Swap bytes if necessary due to Big/Little Endianness mapping between CPU
    // and FPGA wrapper Assuming host is little-endian (ARM default) and wrapper
    // expects Big-Endian layout byte 0 in [15:8]
    uint16_t val = tx_buffer[i];
    uint16_t swapped = (val >> 8) | (val << 8);
    eim_bus[TX_BUFFER_OFFSET + i] = swapped;
  }

  printf("Starting Decoder...\n");
  // Trigger Start Bit (Bit 0 of Control Register)
  eim_bus[CONTROL_REG_OFFSET] = 0x0001;

  // Poll Status Register
  // Bit 0: Busy, Bit 1: Done, Bits 15:8: Error Count
  printf("Polling for completion...\n");
  uint16_t status = 0;
  while (1) {
    status = eim_bus[STATUS_REG_OFFSET];
    // printf("Status: 0x%04X\n", status); // Uncomment for debug
    if (status & 0x0002) {
      break; // Done bit is set
    }
    usleep(100); // Sleep briefly to save CPU
  }

  uint8_t errors = (status >> 8) & 0xFF;
  printf("Decoding Complete. FPGA reported %d bit errors corrected.\n", errors);

  printf("Reading from FPGA RX Buffer (EIM)...\n");
  uint16_t rx_buffer[1152];
  for (int i = 0; i < 1152; i++) {
    uint16_t val = eim_bus[RX_BUFFER_OFFSET + i];
    // Swap back to host endianness
    rx_buffer[i] = (val >> 8) | (val << 8);
  }

  // Write output codeword
  FILE *f_out = fopen(argv[2], "wb");
  if (!f_out) {
    perror("Failed to open output file");
    return 1;
  }
  fwrite(rx_buffer, 1, 2304, f_out);
  fclose(f_out);

  printf("Output written to %s\n", argv[2]);

  // Cleanup
  munmap(map_base, EIM_MEM_SIZE);
  close(fd_mem);

  return 0;
}
