// SPDX-License-Identifier: BSD-2-Clause-Views
/*
 * Copyright (c) 2020-2023 The Regents of the University of California
 */

#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_write32(fdev->addr_reg, addr);
}

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

    reg_write32(fdev->ctrl_reg, FLASH_REGION_OE | FLASH_WE_N);
    reg_read32(fdev->data_reg); // dummy read
    val = reg_read32(fdev->data_reg);
    reg_write32(fdev->ctrl_reg, 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)
{
    reg_write32(fdev->data_reg, data);
    reg_write32(fdev->ctrl_reg, FLASH_REGION_OE | FLASH_DQ_OE | FLASH_OE_N);
    reg_read32(fdev->data_reg); // dummy read
    reg_write32(fdev->ctrl_reg, 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_write32(fdev->ctrl_reg, 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
};