#include <STM32I2C_DMA.hpp>
#include <system_time.hpp>
#include <array>
#include <cassert>
#include "bsp.hpp"
#include "_bsp_private.hpp"
#include "STM32ClockControl.hpp"
#include <stm32h7xx_ll_dma.h>
#include "__STM32_I2C_Common.hpp"
#include <event_queue_interface.hpp>

/* Useful Developer Notes
 *
 ** Working in start/stop mode:
 *
 *   In order to start a transfer in Master mode, I2C Control Register 2 must be written with the
 *   Start condition request, the slave address, the transfer direction, the number of bytes to be
 *   transferred, and the End of Transfer mode. End of Transfer mode is configured by the AUTOEND
 *   bit. If it is set, the Stop condition is automatically sent after the programmed number of
 *   bytes is transferred.
 *
 *   If the AUTOEND bit is not set, the end of transfer is managed by software. After the programmed
 *   number of bytes is transferred, the Transfer Complete (TC) flag is set and an interrupt is
 *   generated, if enabled. Then a Repeated Start or a Stop condition can be requested by
 *   software.
 *
 *   The data transfer can be managed by interrupts or by the DMA.
 *
 *   For payloads > 255: the Reload bit must be set in I2C_CR2
 *   No reload == NBYTES data transferred, followed by STOP or RESTART
 *   Reload == NBYTES is reloaded after NBYTES of data is transfered; data transfer will resume
 *     Interrupt is generated if enabled, and TCR flag is set.
 *
 ** How the DMA UART Driver Works
 *
 * This driver needs to be initialized with references to two DMA channels: one for RX, one for TX
 * The DMA channel declaration controls the DMA device selection and the channel selection.
 *
 * The device enumeration in the header is used as an index into a set of tables which translate
 * to the actual registers + flags used for each individul device channel.
 *
 * When the device is started (with .start()), the DMA devices are configured
 * according to the I2C driver expectations. The I2C device itself is also configured,
 * and the desired interrupts are enabled. The GPIO lines will be configured to work in the proper
 * state.
 *
 * receive() sets up a DMA transfer on the RX side. It has an optional global timeout, ensuring
 * that the tranmission will not be stuck for an infinite amount of time. When the receive DMA
 * interrupt fires, we disable the DMA device, stop the I2C TX transfer, and stop the timeout
 * timer. If the user supplied a callback, it will be invoked. A status will be passed to the user
 * callback indicating whether the transfer succeeded or failed (currently, only via timeout).
 *
 * transmit() sets up a DMA transfer on the TX side. It has no timeout - transmit DMA is always
 * assumed to complete. When the transmit DMA interrupt fires, we disable the DMA device and stop
 * the I2C TX transfer.
 *
 *
 * Arbitration Loss Error
 * If you're trying to communicate with a device, but you have an arbitration loss error, it likely
 * means that there are not I2C pullups. You either need to have externals on the board, or set
 * internal pullups.
 * @code
 * i2c1.setPullups(STM32I2C_DMA::pullups::internal);
 * @endcode
 */

constexpr std::array<uint32_t, I2C_COUNT> dma_tx_routing = {LL_DMAMUX1_REQ_I2C1_TX, LL_DMAMUX1_REQ_I2C2_TX,
													LL_DMAMUX1_REQ_I2C3_TX, LL_DMAMUX2_REQ_I2C4_TX};

constexpr std::array<uint32_t, I2C_COUNT> dma_rx_routing = {LL_DMAMUX1_REQ_I2C1_RX, LL_DMAMUX1_REQ_I2C2_RX,
													LL_DMAMUX1_REQ_I2C3_RX, LL_DMAMUX2_REQ_I2C4_RX};

static inline void enableDMATx(STM32DMA& ch_, I2C_TypeDef* inst, const uint8_t* buffer, size_t size)
{
	ch_.setAddresses(
		const_cast<uint8_t*>(
			buffer), // This is frowned upon, but the underlying STM32 code doesn't take const.
		reinterpret_cast<uint32_t*>(LL_I2C_DMA_GetRegAddr(inst, LL_I2C_DMA_REG_DATA_TRANSMIT)),
		size);
	ch_.enable();
}

static inline void enableDMARx(STM32DMA& ch_, I2C_TypeDef* inst, uint8_t* const buffer, size_t size)
{
	ch_.setAddresses(
		reinterpret_cast<uint32_t*>(LL_I2C_DMA_GetRegAddr(inst, LL_I2C_DMA_REG_DATA_RECEIVE)),
		buffer, size);
	ch_.enable();
}

void STM32I2C_DMA::start() noexcept
{
	assert(!started_);

	auto i2c_inst = i2c_instance[device_];
	assert(i2c_inst); // if failed, device is invalid

	configure_i2c_gpio(device_, pull_setting_);

	STM32ClockControl::i2cEnable(device_);

	// Disable before modifying configuration registers
	LL_I2C_Disable(i2c_inst);

	static LL_I2C_InitTypeDef initializer = {
		.PeripheralMode = LL_I2C_MODE_I2C,
		.Timing = i2c_timings[static_cast<uint8_t>(baudrate_)],
		.AnalogFilter = LL_I2C_ANALOGFILTER_ENABLE,
		.DigitalFilter = 0x00,
		.OwnAddress1 = 0x00,
		.TypeAcknowledge = LL_I2C_ACK,
		.OwnAddrSize = LL_I2C_OWNADDRESS1_7BIT
	};

	configureDMA();

	LL_I2C_EnableDMAReq_RX(i2c_inst);
	LL_I2C_EnableDMAReq_TX(i2c_inst);

	auto r = LL_I2C_Init(i2c_inst, &initializer);
	assert(r == 0);

	enableI2CInterrupts(device_);

	started_ = true;
}

void STM32I2C_DMA::stop() noexcept
{
	assert(started_);

	auto i2c_inst = i2c_instance[device_];
	assert(i2c_inst); // if failed, device is invalid

	disableI2CInterrupts(device_);

	LL_I2C_DeInit(i2c_inst);

	LL_I2C_DisableDMAReq_RX(i2c_inst);
	LL_I2C_DisableDMAReq_TX(i2c_inst);

	STM32ClockControl::i2cDisable(device_);

	started_ = false;
}

void STM32I2C_DMA::configureDMA() noexcept
{
	tx_channel_.setConfiguration(LL_DMA_DIRECTION_MEMORY_TO_PERIPH | LL_DMA_PRIORITY_LOW |
									 LL_DMA_MODE_NORMAL | LL_DMA_PERIPH_NOINCREMENT |
									 LL_DMA_MEMORY_INCREMENT | LL_DMA_PDATAALIGN_BYTE |
									 LL_DMA_MDATAALIGN_BYTE,
								 dma_tx_routing[device_]);
	rx_channel_.setConfiguration(LL_DMA_DIRECTION_PERIPH_TO_MEMORY | LL_DMA_PRIORITY_LOW |
									 LL_DMA_MODE_NORMAL | LL_DMA_PERIPH_NOINCREMENT |
									 LL_DMA_MEMORY_INCREMENT | LL_DMA_PDATAALIGN_BYTE |
									 LL_DMA_MDATAALIGN_BYTE,
								 dma_rx_routing[device_]);

	// TODO: should these be dispatched?
	tx_channel_.registerCallback([this](STM32DMA::status dma_status) {
		if(dma_status != STM32DMA::status::ok)
		{
			if(hasI2CCallback(device_))
			{
				CentralEventQueue::callback(std::bind(getI2CCallback(device_), I2C::status::error));
			}
		}
	});
	rx_channel_.registerCallback([this](STM32DMA::status dma_status) noexcept {
		// Now, this doesn't fire in the event of OK because we're already
		// going to fire an I2C transfer complete interrupt.
		if(dma_status != STM32DMA::status::ok)
		{
			if(hasI2CCallback(device_))
			{
				CentralEventQueue::callback(std::bind(getI2CCallback(device_), I2C::status::error));
			}
		}
	});

	tx_channel_.start();
	rx_channel_.start();
}


void STM32I2C_DMA::setPullups(STM32I2C::pullups pull) noexcept
{
	assert(!started_);
	pull_setting_ = pull;
}

void STM32I2C_DMA::baudrate(STM32I2C::baud rate) noexcept
{
	assert(!started_);
	baudrate_ = rate;
}

I2C::status STM32I2C_DMA::transfer(I2C::op_t& op,
											 const I2C::cb_t& cb) noexcept
{
	auto i2c_inst = i2c_instance[device_];
	assert(i2c_inst); // Instance is not valid if failed
	uint32_t generate_mode = LL_I2C_GENERATE_STOP;
	uint32_t end_mode = LL_I2C_MODE_AUTOEND;
	uint32_t transfer_size = 0;
	bool handle_transfer = true;

	if(busy_)
	{
		return I2C::status::busy;
	}

	busy_ = true;

	// Note that the STM32 calls don't shift the address for us, so we need
	// to do this manually.
	uint32_t address = static_cast<uint32_t>(op.address << 1);

	// Reset per-transfer settings
	//transfer_completed_ = false;
	setI2CCallback(device_, nullptr);

	active_op_ = op;

	/** A Note on Large Transfers (> 255 bytes)
	 *
	 * For large transfers, we need to use the "reload" capability to ensure that the Transfer
	 * Complete Reload (TCR) flag/interrupt will fire. After the programmed number of bytes has been
	 * transferred, you'll get an interrupt. The additional number of bytes to be transferred is
	 * programmed, the TCR bit is set once more, and then the data transfer will resume.
	 *
	 * Note that when RELOAD is set, AUTOEND has no effect. We need to STOP the transfer manually.
	 *
	 * To handle this logic, use the check_and_adjust_transfer_size() function, which will set
	 * variables that are accessible in an interrupt context and can be used for handling these
	 * large transfers without necessary intervention on your part.
	 *
	 * Also note that since the DMA block can handle sizes > 255, we will pass the full
	 * buffer size to the enableDMATx/Rx functions.
	 */
	switch(op.op)
	{
		case I2C::operation::continueWriteStop:
		case I2C::operation::write: {
			// Is this case really separated, because with continue-write I don't want to generate
			// a start? Or is it fine? I think I might need to separate them for a different
			// behavior - e.g. GENERATE_NOSTARTSTOP + AutoEnd?
			std::tie(transfer_size, end_mode) =
				check_and_adjust_transfer_size(device_, op.tx_size, LL_I2C_MODE_AUTOEND);
			generate_mode = LL_I2C_GENERATE_START_WRITE;
			enableDMATx(tx_channel_, i2c_inst, op.tx_buffer, op.tx_size);
			break;
		}
		case I2C::operation::writeNoStop:
		case I2C::operation::continueWriteNoStop: {
			std::tie(transfer_size, end_mode) =
				check_and_adjust_transfer_size(device_, op.tx_size, LL_I2C_MODE_SOFTEND);
			// For continue, does this need to be separated as a case?
			generate_mode = LL_I2C_GENERATE_START_WRITE;
			enableDMATx(tx_channel_, i2c_inst, op.tx_buffer, op.tx_size);
			break;
		}
		case I2C::operation::read: {
			std::tie(transfer_size, end_mode) =
				check_and_adjust_transfer_size(device_, op.rx_size, LL_I2C_MODE_AUTOEND);
			generate_mode = LL_I2C_GENERATE_START_READ;
			enableDMARx(rx_channel_, i2c_inst, op.rx_buffer, op.rx_size);
			break;
		}
		case I2C::operation::writeRead: {
			std::tie(transfer_size, end_mode) = check_and_adjust_transfer_size(device_, op.tx_size, LL_I2C_MODE_SOFTEND);
			generate_mode = LL_I2C_GENERATE_START_WRITE;

			// Override auto-end setting to specify that we need to kick off a read
			// if we are in RELOAD mode.
			setEndTransferAction(device_, end_transfer_option::START_RX);

			// Register intermediate callback that enables receive
			setI2CCallback(device_, [this](auto i2c_status)
			{
				(void) i2c_status;
				// Disable TX DMA, put us in RX mode.
				tx_channel_.disable();
				auto [next_transfer_size, next_end_mode] = check_and_adjust_transfer_size(device_, active_op_.rx_size, LL_I2C_MODE_AUTOEND);
				auto inst = i2c_instance[device_];
				enableDMARx(rx_channel_, inst, active_op_.rx_buffer, active_op_.rx_size);
				LL_I2C_HandleTransfer(inst, active_op_.address << 1, LL_I2C_ADDRSLAVE_7BIT, next_transfer_size, next_end_mode, LL_I2C_GENERATE_START_READ);
				setI2CCallback(device_, [this](auto s)
				{
					assert(s == I2C::status::ok);
					//transfer_completed_ = true;
					tx_channel_.disable();
					rx_channel_.disable();
				});
			});
			enableDMATx(tx_channel_, i2c_inst, op.tx_buffer, op.tx_size);
			break;
		}
		case I2C::operation::ping: {
			transfer_size = 1;
			end_mode = LL_I2C_MODE_AUTOEND;
			generate_mode = LL_I2C_GENERATE_START_READ;
			break;
		}
		case I2C::operation::stop: {
			LL_I2C_GenerateStopCondition(i2c_inst);
			handle_transfer = false;
			break;
		}
		case I2C::operation::restart: {
			LL_I2C_GenerateStartCondition(i2c_inst);
			handle_transfer = false;
			break;
		}
	}

	if(!hasI2CCallback(device_))
	{
		setI2CCallback(device_, [this, cb](auto s)
		{
			tx_channel_.disable();
			rx_channel_.disable();
			busy_ = false;
			if(cb)
			{
				CentralEventQueue::callback(std::bind(cb, s));
			}
		});
	}

	if(handle_transfer)
	{
		// Start the transfer
		LL_I2C_HandleTransfer(i2c_inst, address, LL_I2C_ADDRSLAVE_7BIT, transfer_size, end_mode,
							  generate_mode);
	}

	return I2C::status::ok;
}
