#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <linux/i2c-dev.h>
#include <sys/ioctl.h>
#include <math.h>

#define I2C_DEV_PATH     "/dev/i2c-10"
#define I2C_SLAVE_ADDR   0x43
#define MAX_BUF_LEN      4128

static uint16_t requestID = 1;

/**
 * Convert a hex string (e.g., "010203") to a byte array
 */
void hexStringToBytes(const char* hexStr, uint8_t* outBuf, int* outLen) {
	size_t len = strlen(hexStr);
	*outLen = 0;

	for (size_t i = 0; i < len && i + 1 < len; i += 2) {
		sscanf(&hexStr[i], "%2hhx", &outBuf[(*outLen)++]);
	}

	printf("Parsed %d bytes: ", *outLen);
	for (int i = 0; i < *outLen; i++) {
		printf("0x%02X ", outBuf[i]);
	}
	printf("\n");
}

/**
 * Send a register read command(i2c write)
 * Keep polling until response of register read ready (i2c read of 1 byte)
 * Get the response of register read (i2c read of n bytes)
 */
int readRegister(int fd, uint32_t address, uint16_t length, uint8_t *outBuf, int outBufLen) {
	uint8_t write_buf[20] = {0};
	uint8_t resp_len = 0;
	int ret;

	// CCD Header
	write_buf[0] = 0x00;
	write_buf[1] = 0x40;
	write_buf[2] = 0x00;
	write_buf[3] = 0x08;

	write_buf[4] = 0x0C;  // SCD length = 12
	write_buf[5] = 0x00;

	write_buf[6] = requestID & 0xFF;
	write_buf[7] = (requestID >> 8) & 0xFF;
	requestID++;

	// 64-bit address (only 32 LSB used)
	write_buf[8]  = address & 0xFF;
	write_buf[9]  = (address >> 8) & 0xFF;
	write_buf[10] = (address >> 16) & 0xFF;
	write_buf[11] = (address >> 24) & 0xFF;
	memset(&write_buf[12], 0, 4); // upper 4 bytes = 0

	// Reserved bytes
	write_buf[16] = 0;
	write_buf[17] = 0;
	// Requested length (little-endian)
	write_buf[18] = length & 0xFF;
	write_buf[19] = (length >> 8) & 0xFF;

	if (write(fd, write_buf, sizeof(write_buf)) != sizeof(write_buf)) {
		perror("I2C write failed");
		return -1;
	}


	// Poll for response
	do {
		usleep(100000);  // 100 ms
		ret = read(fd, &resp_len, 1);
		if (ret < 0) {
			perror("Poll read failed");
			return -1;
		}
	} while (resp_len == 0);

	if (resp_len > outBufLen) {
		fprintf(stderr, "Response too long (%d > %d)\n", resp_len, outBufLen);
		return -1;
	}

	ret = read(fd, outBuf, resp_len);
	if (ret < 0) {
		perror("Read failed");
		return -1;
	}

	return ret;
}

/**
 * Send a register write command(i2c write)
 * Keep polling until acknowledgement of register write ready (i2c read of 1 byte)
 * Get the acknowledgement of register write (i2c read of n bytes)
 */
int writeRegister(int fd, uint32_t address, const uint8_t* data, int data_len, uint8_t *outBuf, int outBufLen) {
	if ((data_len + 16) > MAX_BUF_LEN) {
		fprintf(stderr, "Data too long\n");
		return -1;
	}

	uint8_t write_buf[MAX_BUF_LEN];
	int cmd_len = 16 + data_len; // CCD (8) + addr (8) + data
	uint8_t resp_len = 0;
	int ret = -1; 

	// CCD Header
	write_buf[0] = 0x00;
	write_buf[1] = 0x40;
	write_buf[2] = 0x02;
	write_buf[3] = 0x08;

	write_buf[4] = (data_len + 8) & 0xFF;
	write_buf[5] = ((data_len + 8) >> 8) & 0xFF;

	write_buf[6] = requestID & 0xFF;
	write_buf[7] = (requestID >> 8) & 0xFF;
	requestID++;

	// 64-bit address
	write_buf[8]  = address & 0xFF;
	write_buf[9]  = (address >> 8) & 0xFF;
	write_buf[10] = (address >> 16) & 0xFF;
	write_buf[11] = (address >> 24) & 0xFF;
	memset(&write_buf[12], 0, 4);  // Upper 4 bytes

	memcpy(&write_buf[16], data, data_len);

	if (write(fd, write_buf, cmd_len) != cmd_len) {
		perror("I2C write failed");
		return -1;
	}
	printf("Written %d bytes of data\r\n", cmd_len);

#if 0 // To print the data written
	for(int i = 0; i < cmd_len; i++) {
		printf("%d - %x ",i, write_buf[i]);
	}
	printf("\r\n");
#endif

	// Poll for response
	do {
		usleep(1000000); // Do not reduce this delay; to prevent isr from running too much
		ret = read(fd, &resp_len, 1);
		if (ret < 0) {
			perror("Poll read failed");
			return -1;
		}
	} while (resp_len == 0);

	if (resp_len > outBufLen) {
		fprintf(stderr, "Response too long (%d > %d)\n", resp_len, outBufLen);
		return -1;
	}

	ret = read(fd, outBuf, resp_len);
	if (ret < 0) {
		perror("Read failed");
		return -1;
	}
	return ret;
}

/**
 * Send a register write command(i2c write)
 * This function is to be used only to trigger FW update in MCU
 * because it does not poll for the acknowledgent. Polling for acknowledgent 
 * can lead  i2c bus hang if the device reboots during an i2c transaction
 */
int writeRegisterOnly(int fd, uint32_t address, const uint8_t* data, int data_len, uint8_t *outBuf, int outBufLen) {
	if ((data_len + 16) > MAX_BUF_LEN) {
		fprintf(stderr, "Data too long\n");
		return -1;
	}

	uint8_t write_buf[MAX_BUF_LEN];
	int cmd_len = 16 + data_len; // CCD (8) + addr (8) + data
	uint8_t resp_len = 0;
	int ret = -1; 

	// CCD Header
	write_buf[0] = 0x00;
	write_buf[1] = 0x40;
	write_buf[2] = 0x02;
	write_buf[3] = 0x08;

	write_buf[4] = (data_len + 8) & 0xFF;
	write_buf[5] = ((data_len + 8) >> 8) & 0xFF;

	write_buf[6] = requestID & 0xFF;
	write_buf[7] = (requestID >> 8) & 0xFF;
	requestID++;

	// 64-bit address
	write_buf[8]  = address & 0xFF;
	write_buf[9]  = (address >> 8) & 0xFF;
	write_buf[10] = (address >> 16) & 0xFF;
	write_buf[11] = (address >> 24) & 0xFF;
	memset(&write_buf[12], 0, 4);  // Upper 4 bytes

	memcpy(&write_buf[16], data, data_len);

	if (write(fd, write_buf, cmd_len) != cmd_len) {
		perror("I2C write failed");
		return -1;
	}
	printf("Written %d bytes of data\r\n", cmd_len);
}

/**
 * Main 
 * Provides CLI to select:
 *  -Register Read
 *  -Register Write
 *  -MCU FW update
 *  -FPGA FW update
 */
int main() {
	int fd;
	uint32_t address = 0;
	uint16_t length;
	char mode;
	char hexInput[512];
	uint8_t data[MAX_BUF_LEN];
	int data_len = 0;

	uint8_t response[MAX_BUF_LEN];
	int response_len = 0;

	// Open I2C bus
	fd = open(I2C_DEV_PATH, O_RDWR);
	if (fd < 0) {
		perror("Failed to open I2C device");
		return 1;
	}

	// Set slave address
	if (ioctl(fd, I2C_SLAVE, I2C_SLAVE_ADDR) < 0) {
		perror("Failed to set I2C slave address");
		close(fd);
		return 1;
	}

	printf("Select operation: read[r] / write[w] / firmware update[f] / FPGA update[g]: ");
	scanf(" %c", &mode);

	if (mode == 'r' || mode == 'R' || mode == 'w' || mode == 'W') {
		printf("Enter 32-bit register address: ");
		scanf("%x", &address);
	}

	if (mode == 'r' || mode == 'R') {
		printf("Enter number of bytes to read: ");
		scanf("%hu", &length);

		response_len = readRegister(fd, address, length, response, sizeof(response));
		if (response_len > 0) {
			printf("Received %d bytes: ", response_len);
			for (int i = 0; i < response_len; ++i)
				printf("0x%02X ", response[i]);
			printf("\n");
		}

	} else if (mode == 'w' || mode == 'W') {
		printf("Enter data to write (2 to 4 bytes, hex without spaces, e.g., '1A2B'): ");
		scanf("%s", hexInput);

		hexStringToBytes(hexInput, data, &data_len);

		response_len = writeRegister(fd, address, data, data_len, response, sizeof(response));
		if (response_len > 0) {
			printf("Response (%d bytes): ", response_len);
			for (int i = 0; i < response_len; ++i)
				printf("0x%02X ", response[i]);
			printf("\n");
		}

	} else if (mode == 'f' || mode == 'F') {
		char firmwareFile[256];
		printf("Enter firmware file path: ");
		scanf("%s", firmwareFile);

		FILE *fw = fopen(firmwareFile, "rb");
		if (!fw) {
			perror("Failed to open firmware file");
			close(fd);
			return 1;
		}

		fseek(fw, 0, SEEK_END);
		size_t fileSize = ftell(fw);
		fseek(fw, 0, SEEK_SET);

		size_t minLength = fileSize + 2; // +2 for CRC
		size_t sectors = (size_t)ceil((double)minLength / 4096.0);
		size_t totalLength = sectors * 4096;

		uint8_t *firmware = calloc(totalLength, sizeof(uint8_t));
		if (!firmware) {
			fprintf(stderr, "Memory allocation failed\n");
			fclose(fw);
			close(fd);
			return 1;
		}

		fread(firmware, 1, fileSize, fw);
		fclose(fw);

		// Compute CRC-16
		uint16_t crc = 0xFFFF;
		uint16_t polynomial = 0x8005; // CRC-16-IBM polynomial
		for (size_t i = 0; i < totalLength - 2; i++) {
			crc ^= (uint16_t)firmware[i] << 8;
			for (int j = 0; j < 8; j++) {
				if (crc & 0x8000)
					crc = (crc << 1) ^ polynomial;
				else
					crc <<= 1;
			}
		}

		firmware[totalLength - 2] = crc & 0xFF;
		firmware[totalLength - 1] = (crc >> 8) & 0xFF;

		printf("Writing firmware (%zu bytes) to 0x64000000...\n", totalLength);

		size_t offset = 0;
		uint32_t addr = 0x64000000;
		while (offset < totalLength) {
			if (writeRegister(fd, addr, &firmware[offset], 4096, response, sizeof(response)) < 0) {
				fprintf(stderr, "Failed to write firmware block at offset 0x%X\n", offset);
				free(firmware);
				close(fd);
				return 1;
			}
			offset += 4096;
			addr += 4096;  // Increment by 4K
		}

		// Send sector count to trigger update
		uint8_t n = (uint8_t)sectors;
		writeRegisterOnly(fd, 0x2078, &n, 1, response, sizeof(response));
		printf("Firmware update triggered with %d sectors\r\n", n);
		free(firmware);

	} else if (mode == 'g' || mode == 'G') {
		char fpgaFile[256];
		printf("Enter FPGA firmware file path: ");
		scanf("%s", fpgaFile);

		FILE *fpga = fopen(fpgaFile, "rb");
		if (!fpga) {
			perror("Failed to open FPGA firmware file");
			close(fd);
			return 1;
		}

		fseek(fpga, 0, SEEK_END);
		size_t fileSize = ftell(fpga);
		fseek(fpga, 0, SEEK_SET);

		printf("Writing FPGA firmware (%zu bytes) to 0x70000000...\n", fileSize);

		uint8_t *buffer = malloc(4096);
		if (!buffer) {
			fprintf(stderr, "Memory allocation failed\n");
			fclose(fpga);
			close(fd);
			return 1;
		}

		size_t offset = 0;
		uint32_t addr = 0x70000000;
		while (offset < fileSize) {
			size_t chunk = fread(buffer, 1, 4096, fpga);
			if (chunk == 0) break;

			if (writeRegister(fd, addr, buffer, chunk, response, sizeof(response)) < 0) {
				fprintf(stderr, "Failed to write FPGA firmware at 0x%08X\n", addr);
				free(buffer);
				fclose(fpga);
				close(fd);
				return 1;
			}
			offset += chunk;
			addr += 4096;
		}
		free(buffer);
		fclose(fpga);
		printf("FPGA firmware update complete (%zu bytes written).\n", fileSize);
	} else {
		fprintf(stderr, "Invalid operation. Use 'r', 'w', 'f', or 'g'.\n");
	}

	close(fd);
	return 0;
}