/* flashimg/cpu_main/flashio.c - DMA-safe flash storage
 *
 * Copyright (C) 2025  Luke T. Shumaker <lukeshu@lukeshu.com>
 * SPDX-License-Identifier: AGPL-3.0-or-later
 */

#include <string.h>

#include <libcr/coroutine.h>

#define LOG_NAME FLASH
#include <libmisc/log.h>

#define IMPLEMENTATION_FOR_FLASHIO_H YES
#include "flashio.h"

LO_IMPLEMENTATION_C(io_preader_to, struct flashio, flashio);
LO_IMPLEMENTATION_C(io_pwriter,    struct flashio, flashio);
LO_IMPLEMENTATION_C(io_flusher,    struct flashio, flashio);

#define DATA_START ((const char *)(XIP_NOALLOC_BASE))
#define DATA_SIZE  PICO_FLASH_SIZE_BYTES
static_assert(DATA_SIZE % FLASH_SECTOR_SIZE == 0);

/* There are some memcpy()s (and memcmp()s?) in here that maybe should
 * be replaced with SSI DMA.  */

static void flashio_buffer_flush(struct flashio *self) {
	if (!self->dat_dirty)
		return;
	self->dat_dirty = false;
	if (memcmp(self->dat, DATA_START+self->dat_pos, FLASH_SECTOR_SIZE) == 0)
		return;
	log_infoln("write flash sector @ ", (base16_u32_, self->dat_pos), "...");
	bool saved = cr_save_and_disable_interrupts();
	flash_range_erase(self->dat_pos, FLASH_SECTOR_SIZE);
	flash_range_program(self->dat_pos, self->dat, FLASH_SECTOR_SIZE);
	cr_restore_interrupts(saved);
	log_debugln("... written");
}

static void flashio_buffer_load(struct flashio *self, size_t sector_base, bool noinit) {
	if (self->dat_ok && self->dat_pos == sector_base)
		return;
	flashio_buffer_flush(self);
	self->dat_ok = true;
	self->dat_pos = sector_base;
	if (!noinit)
		memcpy(self->dat, DATA_START+sector_base, FLASH_SECTOR_SIZE);
}

size_t_and_error flashio_pread_to(struct flashio *self, lo_interface io_writer dst, uoff_t _src_offset, size_t count) {
	assert(self);
	assert(!LO_IS_NULL(dst));

	if (_src_offset > DATA_SIZE)
		return ERROR_AND(size_t, 0, error_new(E_POSIX_EINVAL, "offset is past the chip size"));
	size_t src_offset = (size_t) _src_offset;

	if (src_offset == DATA_SIZE)
		return ERROR_AND(size_t, 0, error_new(E_EOF));

	/* Assume that somewhere down the line the pointer we pass to
	 * io_write() will be passed to DMA.  We don't want the DMA
	 * engine to hit (slow) XIP (for instance, this can cause
	 * reads/writes to the SSP to get out of sync with eachother).
	 * So, copy the data to a buffer in (fast) RAM first.  It's
	 * lame that the DMA engine can only have a DREQ on one side
	 * of the channel.
	 */
	size_t sector_base = LM_ROUND_DOWN(src_offset, FLASH_SECTOR_SIZE);
	if (src_offset + count > sector_base + FLASH_SECTOR_SIZE)
		count = (sector_base + FLASH_SECTOR_SIZE) - src_offset;
	assert(src_offset + count <= DATA_SIZE);

	cr_mutex_lock(&self->mu);
	flashio_buffer_load(self, sector_base, false);
	size_t_and_error ret = io_write(dst, &self->dat[src_offset-sector_base], count);
	cr_mutex_unlock(&self->mu);

	return ret;
}

size_t_and_error flashio_pwritev(struct flashio *self, const struct wr_iovec *iov, int iovcnt, uoff_t _offset) {
	assert(self);
	assert(iov);
	assert(iovcnt > 0);

	size_t total_count = 0;
	for (int i = 0; i < iovcnt; i++)
		total_count += iov[i].iov_len;
	assert(total_count > 0);

	if (_offset >= DATA_SIZE)
		return ERROR_AND(size_t, 0, error_new(E_POSIX_ENOSPC, "cannot write past the chip size"));
	size_t offset = (size_t) _offset;

	cr_mutex_lock(&self->mu);
	size_t total_done = 0;
	for (int i = 0; i < iovcnt; i++) {
		for (size_t iov_done = 0; iov_done < iov[i].iov_len;) {
			if (offset >= DATA_SIZE)
				cr_mutex_unlock(&self->mu);
				return ERROR_AND(size_t, total_done, error_new(E_POSIX_ENOSPC, "cannot write past the chip size"));
			size_t sector_base = LM_ROUND_DOWN(offset, FLASH_SECTOR_SIZE);
			size_t count = iov[i].iov_len - iov_done;
			if (offset + count > sector_base + FLASH_SECTOR_SIZE)
				count = (sector_base + FLASH_SECTOR_SIZE) - offset;
			assert(offset + count <= DATA_SIZE);

			flashio_buffer_load(self, sector_base, count != FLASH_SECTOR_SIZE);
			memcpy(&self->dat[offset-sector_base], iov[i].iov_write_from+iov_done, count);
			total_done += count;
			iov_done += count;
			offset += count;
		}
	}
	cr_mutex_unlock(&self->mu);
	return ERROR_AND(size_t, total_done, ERROR_NULL);
}

error flashio_flush(struct flashio *self) {
	assert(self);

	cr_mutex_lock(&self->mu);
	flashio_buffer_flush(self);
	cr_mutex_unlock(&self->mu);

	return ERROR_NULL;
}