// SPDX-License-Identifier: GPL
/*

Copyright (c) 2026 FPGA Ninja, LLC

Authors:
- Alex Forencich

*/

#include "flash.h"

#include <stdio.h>
#include <stdlib.h>

#define reg_read32(reg) (*((volatile uint32_t *)(reg)))
#define reg_write32(reg, val) (*((volatile uint32_t *)(reg))) = (val)

#define CFI_QUERY_ADDR             0x55
#define CFI_QUERY_DATA             0x98
#define CFI_READ_ARRAY             0xFF
#define CFI_READ_ARRAY_ALT         0xF0
#define CFI_ID_0                   0x10
#define CFI_ID_1                   0x11
#define CFI_ID_2                   0x12
#define CFI_PRI_CMD_SET_0          0x13
#define CFI_PRI_CMD_SET_1          0x14
#define CFI_DEVICE_SIZE            0x27
#define CFI_WRITE_BUFFER_SIZE_0    0x2A
#define CFI_WRITE_BUFFER_SIZE_1    0x2B
#define CFI_ERASE_REGION_COUNT     0x2C
#define CFI_ERASE_REGION_1_INFO_0  0x2D
#define CFI_ERASE_REGION_1_INFO_1  0x2E
#define CFI_ERASE_REGION_1_INFO_2  0x2F
#define CFI_ERASE_REGION_1_INFO_3  0x30
#define CFI_ERASE_REGION_2_INFO_0  0x31
#define CFI_ERASE_REGION_2_INFO_1  0x32
#define CFI_ERASE_REGION_2_INFO_2  0x33
#define CFI_ERASE_REGION_2_INFO_3  0x34

#define BPI_INTEL_READ_ARRAY                0xFF
#define BPI_INTEL_READ_STATUS_REG           0x70
#define BPI_INTEL_READ_ID                   0x90
#define BPI_INTEL_CLEAR_STATUS_REG          0x50
#define BPI_INTEL_READ_CONFIG_REG_SETUP     0x60
#define BPI_INTEL_SET_READ_CONFIG_REG       0x03
#define BPI_INTEL_BLOCK_LOCK_SETUP          0x60
#define BPI_INTEL_BLOCK_LOCK                0x01
#define BPI_INTEL_BLOCK_UNLOCK              0xD0
#define BPI_INTEL_BLOCK_ERASE_SETUP         0x20
#define BPI_INTEL_BLOCK_ERASE_CONFIRM       0xD0
#define BPI_INTEL_BUFFERED_PROGRAM_SETUP    0xE8
#define BPI_INTEL_BUFFERED_PROGRAM_CONFIRM  0xD0

#define BPI_AMD_UNLOCK_ADDR_1             0x555
#define BPI_AMD_UNLOCK_DATA_1             0xAA
#define BPI_AMD_UNLOCK_ADDR_2             0x2AA
#define BPI_AMD_UNLOCK_DATA_2             0x55
#define BPI_AMD_UNLOCK_BYPASS_ENTER_ADDR  0x555
#define BPI_AMD_UNLOCK_BYPASS_ENTER_DATA  0x20
#define BPI_AMD_UNLOCK_BYPASS_RESET_1     0x90
#define BPI_AMD_UNLOCK_BYPASS_RESET_2     0x00
#define BPI_AMD_BLOCK_ERASE_SETUP         0x80
#define BPI_AMD_BLOCK_ERASE_CONFIRM       0x30
#define BPI_AMD_BUFFERED_PROGRAM_SETUP    0x25
#define BPI_AMD_BUFFERED_PROGRAM_CONFIRM  0x29
#define BPI_AMD_READ_ARRAY                0xF0

#define BPI_MICRON_READ_ARRAY                0xFF
#define BPI_MICRON_READ_STATUS_REG           0x70
#define BPI_MICRON_READ_ID                   0x90
#define BPI_MICRON_CLEAR_STATUS_REG          0x50
#define BPI_MICRON_READ_CONFIG_REG_SETUP     0x60
#define BPI_MICRON_SET_READ_CONFIG_REG       0x03
#define BPI_MICRON_BLOCK_LOCK_SETUP          0x60
#define BPI_MICRON_BLOCK_LOCK                0x01
#define BPI_MICRON_BLOCK_UNLOCK              0xD0
#define BPI_MICRON_BLOCK_ERASE_SETUP         0x20
#define BPI_MICRON_BLOCK_ERASE_CONFIRM       0xD0
#define BPI_MICRON_BUFFERED_PROGRAM_SETUP    0xE9
#define BPI_MICRON_BUFFERED_PROGRAM_CONFIRM  0xD0

#define FLASH_CE_N      (1 << 0)
#define FLASH_OE_N      (1 << 1)
#define FLASH_WE_N      (1 << 2)
#define FLASH_ADV_N     (1 << 3)
#define FLASH_DQ_OE     (1 << 8)
#define FLASH_REGION_OE (1 << 16)

void bpi_flash_set_addr(struct flash_device *fdev, size_t addr)
{
	reg_if_write32(fdev->reg, fdev->addr_reg_offset, addr);
}

uint16_t bpi_flash_read_cur(struct flash_device *fdev)
{
	uint32_t val;

	reg_if_write32(fdev->reg, fdev->ctrl_reg_offset, FLASH_REGION_OE | FLASH_WE_N);
	reg_if_read32(fdev->reg, fdev->ctrl_reg_offset, &val); // dummy read
	reg_if_read32(fdev->reg, fdev->data_reg_offset, &val);
	reg_if_write32(fdev->reg, fdev->ctrl_reg_offset, FLASH_OE_N | FLASH_WE_N | FLASH_ADV_N);

	return val;
}

uint16_t bpi_flash_read_word(struct flash_device *fdev, size_t addr)
{
	bpi_flash_set_addr(fdev, addr);
	return bpi_flash_read_cur(fdev);
}

void bpi_flash_write_cur(struct flash_device *fdev, uint16_t data)
{
	uint32_t val;

	reg_if_write32(fdev->reg, fdev->data_reg_offset, data);
	reg_if_write32(fdev->reg, fdev->ctrl_reg_offset, FLASH_REGION_OE | FLASH_DQ_OE | FLASH_OE_N);
	reg_if_read32(fdev->reg, fdev->ctrl_reg_offset, &val); // dummy read
	reg_if_write32(fdev->reg, fdev->ctrl_reg_offset, FLASH_OE_N | FLASH_WE_N | FLASH_ADV_N);
}

void bpi_flash_write_word(struct flash_device *fdev, size_t addr, uint16_t data)
{
	bpi_flash_set_addr(fdev, addr);
	bpi_flash_write_cur(fdev, data);
}

void bpi_flash_deselect(struct flash_device *fdev)
{
	bpi_flash_write_word(fdev, 0, CFI_READ_ARRAY);
	reg_if_write32(fdev->reg, fdev->ctrl_reg_offset, FLASH_CE_N | FLASH_OE_N | FLASH_WE_N | FLASH_ADV_N);
}

// Intel flash ops (0x0001)

uint16_t bpi_flash_intel_read_status_register(struct flash_device *fdev)
{
	bpi_flash_write_cur(fdev, BPI_INTEL_READ_STATUS_REG);
	return bpi_flash_read_cur(fdev);
}

void bpi_flash_intel_clear_status_register(struct flash_device *fdev)
{
	bpi_flash_write_cur(fdev, BPI_INTEL_CLEAR_STATUS_REG);
}

void bpi_flash_intel_init(struct flash_device *fdev)
{
	bpi_flash_intel_clear_status_register(fdev);
}

int bpi_flash_intel_sector_erase(struct flash_device *fdev, size_t addr)
{
	bpi_flash_set_addr(fdev, addr);
	bpi_flash_write_cur(fdev, BPI_INTEL_BLOCK_LOCK_SETUP);
	bpi_flash_write_cur(fdev, BPI_INTEL_BLOCK_UNLOCK);

	if (bpi_flash_intel_read_status_register(fdev) & 0x30) {
		fprintf(stderr, "Failed to unlock block\n");
		return -1;
	}

	bpi_flash_write_cur(fdev, BPI_INTEL_BLOCK_ERASE_SETUP);
	bpi_flash_write_cur(fdev, BPI_INTEL_BLOCK_ERASE_CONFIRM);

	while (!(bpi_flash_intel_read_status_register(fdev) & 0x80)) {};

	if (bpi_flash_intel_read_status_register(fdev) & 0x30){
		fprintf(stderr, "Failed to erase block\n");
		return -1;
	}

	return 0;
}

int bpi_flash_intel_buffered_program(struct flash_device *fdev, size_t addr, size_t len, const void *src)
{
	const uint8_t *s = (const uint8_t *)src;

	bpi_flash_set_addr(fdev, addr);
	bpi_flash_write_cur(fdev, BPI_INTEL_BUFFERED_PROGRAM_SETUP);
	bpi_flash_write_cur(fdev, len-1);

	for (size_t i = 0; i < len; i++) {
		bpi_flash_write_word(fdev, addr+i, s[0] | (s[1] << 8));
		s += 2;
	}

	bpi_flash_set_addr(fdev, addr);
	bpi_flash_write_cur(fdev, BPI_INTEL_BUFFERED_PROGRAM_CONFIRM);

	while (!(bpi_flash_intel_read_status_register(fdev) & 0x80)) {};

	if (bpi_flash_intel_read_status_register(fdev) & 0x30) {
		fprintf(stderr, "Failed to write block\n");
		return -1;
	}

	return 0;
}

const struct flash_ops bpi_flash_intel_ops = {
	.init = bpi_flash_intel_init,
	.sector_erase = bpi_flash_intel_sector_erase,
	.buffered_program = bpi_flash_intel_buffered_program
};

// AMD flash ops (0x0002)

void bpi_flash_amd_unlock(struct flash_device *fdev)
{
	bpi_flash_write_word(fdev, BPI_AMD_UNLOCK_ADDR_1, BPI_AMD_UNLOCK_DATA_1);
	bpi_flash_write_word(fdev, BPI_AMD_UNLOCK_ADDR_2, BPI_AMD_UNLOCK_DATA_2);
}

void bpi_flash_amd_write_buffer_abort_reset(struct flash_device *fdev)
{
	bpi_flash_amd_unlock(fdev);
	bpi_flash_write_word(fdev, BPI_AMD_UNLOCK_ADDR_1, BPI_AMD_READ_ARRAY);
}

void bpi_flash_amd_init(struct flash_device *fdev)
{
	// write-to-buffer-abort reset (just in case)
	bpi_flash_amd_write_buffer_abort_reset(fdev);
}

int bpi_flash_amd_wait_for_operation(struct flash_device *fdev, uint16_t stop_mask)
{
	uint16_t read_1, read_2, read_3;

	while (1) {
		read_1 = bpi_flash_read_cur(fdev);
		read_2 = bpi_flash_read_cur(fdev);
		read_3 = bpi_flash_read_cur(fdev);

		if ((read_1 ^ read_2) & (read_2 ^ read_3) & 0x40) {
			if (read_1 & stop_mask) {
				return read_1;
			}
		} else {
			return 0;
		}
	}
}

int bpi_flash_amd_sector_erase(struct flash_device *fdev, size_t addr)
{
	bpi_flash_amd_unlock(fdev);
	bpi_flash_write_word(fdev, BPI_AMD_UNLOCK_ADDR_1, BPI_AMD_BLOCK_ERASE_SETUP);
	bpi_flash_amd_unlock(fdev);
	bpi_flash_write_word(fdev, addr, BPI_AMD_BLOCK_ERASE_CONFIRM);

	while (!(bpi_flash_read_cur(fdev) & 0x08)) {};

	if (bpi_flash_amd_wait_for_operation(fdev, 0x20) & 0x20) {
		// write-to-buffer-abort reset
		bpi_flash_amd_write_buffer_abort_reset(fdev);

		fprintf(stderr, "Failed to erase block\n");
		return -1;
	}

	return 0;
}

int bpi_flash_amd_buffered_program(struct flash_device *fdev, size_t addr, size_t len, const void *src)
{
	const uint8_t *s = (const uint8_t *)src;

	bpi_flash_amd_unlock(fdev);
	bpi_flash_write_word(fdev, addr, BPI_AMD_BUFFERED_PROGRAM_SETUP);
	bpi_flash_write_cur(fdev, len-1);

	for (size_t i = 0; i < len; i++) {
		bpi_flash_write_word(fdev, addr+i, s[0] | (s[1] << 8));
		s += 2;
	}

	bpi_flash_set_addr(fdev, addr);
	bpi_flash_write_cur(fdev, BPI_AMD_BUFFERED_PROGRAM_CONFIRM);

	if (bpi_flash_amd_wait_for_operation(fdev, 0x22) & 0x22) {
		// write-to-buffer-abort reset
		bpi_flash_amd_write_buffer_abort_reset(fdev);

		fprintf(stderr, "Failed to write block\n");
		return -1;
	}

	return 0;
}

const struct flash_ops bpi_flash_amd_ops = {
	.init = bpi_flash_amd_init,
	.sector_erase = bpi_flash_amd_sector_erase,
	.buffered_program = bpi_flash_amd_buffered_program
};

// Micron flash ops (0x0002)

uint16_t bpi_flash_micron_read_status_register(struct flash_device *fdev)
{
	bpi_flash_write_cur(fdev, BPI_MICRON_READ_STATUS_REG);
	return bpi_flash_read_cur(fdev);
}

void bpi_flash_micron_clear_status_register(struct flash_device *fdev)
{
	bpi_flash_write_cur(fdev, BPI_MICRON_CLEAR_STATUS_REG);
}

void bpi_flash_micron_init(struct flash_device *fdev)
{
	bpi_flash_micron_clear_status_register(fdev);
}

int bpi_flash_micron_sector_erase(struct flash_device *fdev, size_t addr)
{
	bpi_flash_set_addr(fdev, addr);
	bpi_flash_write_cur(fdev, BPI_MICRON_BLOCK_LOCK_SETUP);
	bpi_flash_write_cur(fdev, BPI_MICRON_BLOCK_UNLOCK);

	if (bpi_flash_micron_read_status_register(fdev) & 0x30) {
		fprintf(stderr, "Failed to unlock block\n");
		bpi_flash_write_cur(fdev, CFI_READ_ARRAY);
		return -1;
	}

	bpi_flash_write_cur(fdev, BPI_MICRON_BLOCK_ERASE_SETUP);
	bpi_flash_write_cur(fdev, BPI_MICRON_BLOCK_ERASE_CONFIRM);

	while (!(bpi_flash_micron_read_status_register(fdev) & 0x80)) {};

	if (bpi_flash_micron_read_status_register(fdev) & 0x30) {
		fprintf(stderr, "Failed to erase block\n");
		bpi_flash_write_cur(fdev, CFI_READ_ARRAY);
		return -1;
	}

	bpi_flash_write_cur(fdev, CFI_READ_ARRAY);

	return 0;
}

int bpi_flash_micron_buffered_program(struct flash_device *fdev, size_t addr, size_t len, const void *src)
{
	const uint8_t *s = (const uint8_t *)src;

	bpi_flash_set_addr(fdev, addr);
	bpi_flash_write_cur(fdev, BPI_MICRON_BUFFERED_PROGRAM_SETUP);
	bpi_flash_write_cur(fdev, len-1);

	for (size_t i = 0; i < len; i++) {
		bpi_flash_write_word(fdev, addr+i, s[0] | (s[1] << 8));
		s += 2;
	}

	bpi_flash_set_addr(fdev, addr);
	bpi_flash_write_cur(fdev, BPI_MICRON_BUFFERED_PROGRAM_CONFIRM);

	while (!(bpi_flash_micron_read_status_register(fdev) & 0x80)) {};

	if (bpi_flash_micron_read_status_register(fdev) & 0x30) {
		fprintf(stderr, "Failed to write block\n");
		bpi_flash_write_cur(fdev, CFI_READ_ARRAY);
		return -1;
	}

	bpi_flash_write_cur(fdev, CFI_READ_ARRAY);

	return 0;
}

const struct flash_ops bpi_flash_micron_ops = {
	.init = bpi_flash_micron_init,
	.sector_erase = bpi_flash_micron_sector_erase,
	.buffered_program = bpi_flash_micron_buffered_program
};


void bpi_flash_release(struct flash_device *fdev)
{
	bpi_flash_deselect(fdev);
}

int bpi_flash_init(struct flash_device *fdev)
{
	int ret = 0;

	if (!fdev)
		return -1;

	// CFI query
	bpi_flash_write_word(fdev, CFI_QUERY_ADDR, CFI_QUERY_DATA);

	if (bpi_flash_read_word(fdev, CFI_ID_0) != 'Q') {
		// may be Intel flash in sync read mode; attempt switch to async
		bpi_flash_write_word(fdev, 0xf94f, BPI_INTEL_READ_CONFIG_REG_SETUP);
		bpi_flash_write_word(fdev, 0xf94f, BPI_INTEL_SET_READ_CONFIG_REG);
		bpi_flash_write_word(fdev, CFI_QUERY_ADDR, CFI_QUERY_DATA);
	}

	if (bpi_flash_read_word(fdev, CFI_ID_0) != 'Q' && ((bpi_flash_read_cur(fdev) ^ bpi_flash_read_cur(fdev)) & 0x44)) {
		// may be AMD flash in write buffer abort; perform write buffer abort reset
		bpi_flash_amd_write_buffer_abort_reset(fdev);
		bpi_flash_write_word(fdev, CFI_QUERY_ADDR, CFI_QUERY_DATA);
	}

	if (bpi_flash_read_word(fdev, CFI_ID_0) != 'Q' ||
			bpi_flash_read_word(fdev, CFI_ID_1) != 'R' ||
			bpi_flash_read_word(fdev, CFI_ID_2) != 'Y') {
		fprintf(stderr, "Failed to read flash ID\n");
		ret = -1;
		goto err;
	}

	fdev->protocol = bpi_flash_read_word(fdev, CFI_PRI_CMD_SET_0) | (bpi_flash_read_word(fdev, CFI_PRI_CMD_SET_1) << 8);

	printf("Command set: %d\n", fdev->protocol);

	switch (fdev->protocol) {
	case 0x0001:
		// Intel command set (P30)
		fdev->ops = &bpi_flash_intel_ops;
		break;
	case 0x0002:
		// AMD command set (S29, MT28)
		fdev->ops = &bpi_flash_amd_ops;
		break;
	case 0x0200:
		// Micron
		fdev->ops = &bpi_flash_micron_ops;
		break;
	default:
		fprintf(stderr, "Unknown command set: %d\n", fdev->protocol);
		ret = -1;
		goto err;
	}

	uint8_t flash_size = bpi_flash_read_word(fdev, CFI_DEVICE_SIZE);
	fdev->size = ((size_t)1) << flash_size;

	printf("Flash size: %d MB\n", 1 << (flash_size-20));

	uint16_t write_buffer_size = bpi_flash_read_word(fdev, CFI_WRITE_BUFFER_SIZE_0) | (bpi_flash_read_word(fdev, CFI_WRITE_BUFFER_SIZE_1) << 8);
	fdev->write_buffer_size = ((size_t)1) << write_buffer_size;

	printf("Write buffer size: %ld B\n", fdev->write_buffer_size);

	fdev->erase_region_count = bpi_flash_read_word(fdev, CFI_ERASE_REGION_COUNT);

	printf("Erase regions: %d\n", fdev->erase_region_count);

	if (fdev->erase_region_count > 0) {
		fdev->erase_region[0].block_count = (bpi_flash_read_word(fdev, CFI_ERASE_REGION_1_INFO_0) | (bpi_flash_read_word(fdev, CFI_ERASE_REGION_1_INFO_1) << 8)) + 1;
		fdev->erase_region[0].block_size = (bpi_flash_read_word(fdev, CFI_ERASE_REGION_1_INFO_2) | (bpi_flash_read_word(fdev, CFI_ERASE_REGION_1_INFO_3) << 8)) * 256;
		fdev->erase_region[0].region_start = 0;
		fdev->erase_region[0].region_end = fdev->erase_region[0].region_start + fdev->erase_region[0].block_count * fdev->erase_region[0].block_size;

		fdev->erase_block_size = fdev->erase_region[0].block_size;

		printf("Erase region 0 block count: %ld\n", fdev->erase_region[0].block_count);
		printf("Erase region 0 block size: %ld B\n", fdev->erase_region[0].block_size);
		printf("Erase region 0 start: 0x%08lx\n", fdev->erase_region[0].region_start);
		printf("Erase region 0 end: 0x%08lx\n", fdev->erase_region[0].region_end);
	} else {
		fprintf(stderr, "No erase regions found!\n");
		ret = -1;
		goto err;
	}

	if (fdev->erase_region_count > 1) {
		fdev->erase_region[1].block_count = (bpi_flash_read_word(fdev, CFI_ERASE_REGION_2_INFO_0) | (bpi_flash_read_word(fdev, CFI_ERASE_REGION_2_INFO_1) << 8)) + 1;
		fdev->erase_region[1].block_size = (bpi_flash_read_word(fdev, CFI_ERASE_REGION_2_INFO_2) | (bpi_flash_read_word(fdev, CFI_ERASE_REGION_2_INFO_3) << 8)) * 256;
		fdev->erase_region[1].region_start = fdev->erase_region[0].region_end;
		fdev->erase_region[1].region_end = fdev->erase_region[1].region_start + fdev->erase_region[1].block_count * fdev->erase_region[1].block_size;

		if (fdev->erase_region[1].block_size > fdev->erase_block_size) {
			fdev->erase_block_size = fdev->erase_region[1].block_size;
		}

		printf("Erase region 1 block count: %ld\n", fdev->erase_region[1].block_count);
		printf("Erase region 1 block size: %ld B\n", fdev->erase_region[1].block_size);
		printf("Erase region 1 start: 0x%08lx\n", fdev->erase_region[1].region_start);
		printf("Erase region 1 end: 0x%08lx\n", fdev->erase_region[1].region_end);
	}

	printf("Erase block size: %ld B\n", fdev->erase_block_size);

	fdev->ops->init(fdev);

err:
	bpi_flash_release(fdev);
	return ret;
}

int bpi_flash_read(struct flash_device *fdev, size_t addr, size_t len, void *dest)
{
	char *d = dest;

	bpi_flash_write_word(fdev, 0, CFI_READ_ARRAY);

	if (addr & 1) {
		*d = bpi_flash_read_word(fdev, addr >> 1) >> 8;
		addr++;
		len--;
		d++;
	}

	while (len > 1) {
		*((uint16_t *)d) = bpi_flash_read_word(fdev, addr >> 1);
		addr += 2;
		len -= 2;
		d += 2;
	}

	if (len) {
		*d = bpi_flash_read_word(fdev, addr >> 1);
		addr++;
		len--;
		d++;
	}

	bpi_flash_deselect(fdev);

	return 0;
}

int bpi_flash_write(struct flash_device *fdev, size_t addr, size_t len, const void *src)
{
	const char *s = src;

	while (len > 0) {
		size_t seg = len;

		// align to buffer size
		if (seg > fdev->write_buffer_size - (addr & (fdev->write_buffer_size-1))) {
			seg = fdev->write_buffer_size - (addr & (fdev->write_buffer_size-1));
		}

		if (fdev->ops->buffered_program(fdev, addr >> 1, seg >> 1, s)) {
			fprintf(stderr, "Buffered write failed\n");
			bpi_flash_deselect(fdev);
			return -1;
		}

		addr += seg;
		len -= seg;
		s += seg;
	}

	bpi_flash_deselect(fdev);

	return 0;
}

int bpi_flash_erase(struct flash_device *fdev, size_t addr, size_t len)
{
	size_t erase_block_size = fdev->erase_block_size;

	while (len > 0) {
		// determine sector size
		erase_block_size = 0;

		for (int k = 0; k < fdev->erase_region_count; k++) {
			if (addr >= fdev->erase_region[k].region_start && addr < fdev->erase_region[k].region_end) {
				erase_block_size = fdev->erase_region[k].block_size;
				break;
			}
		}

		if (!erase_block_size) {
			fprintf(stderr, "Address does not match an erase region\n");
			bpi_flash_deselect(fdev);
			return -1;
		}

		// check size and alignment
		if (addr & (erase_block_size-1) || len < erase_block_size) {
			fprintf(stderr, "Invalid erase request\n");
			bpi_flash_deselect(fdev);
			return -1;
		}

		// block erase
		if (fdev->ops->sector_erase(fdev, addr >> 1)) {
			fprintf(stderr, "Failed to erase sector\n");
			bpi_flash_deselect(fdev);
			return -1;
		}

		if (len <= erase_block_size)
			break;

		addr += erase_block_size;
		len -= erase_block_size;
	}

	bpi_flash_deselect(fdev);

	return 0;
}

const struct flash_driver bpi_flash_driver = {
	.init = bpi_flash_init,
	.release = bpi_flash_release,
	.read = bpi_flash_read,
	.write = bpi_flash_write,
	.erase = bpi_flash_erase
};
