#include <STM32I2C_Blocking.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
 *
 * This driver currently doesn't handle payloads > 255. Honestly, it could, but you could also just
 * use the DMA driver impementation too.
 *
 * 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.
 *
 *
 * 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(STM32I2CMaster::pullups::internal);
 * @endcode
 */

enum class NACK_CHECK : uint8_t
{
	ADDR,
	DATA
};

constexpr unsigned LOOP_TIMEOUT_LIMIT = 1000; // units: iterations

/// Increment counter and operate as boolean to determine if the timeout
/// limit has been exceeded.
/// Returns true if counter exceeds the timeout limit
#define LOOP_TIMEOUT_CHECK(counter) (counter++ < LOOP_TIMEOUT_LIMIT)
#define TIMEOUT_EXCEEDED(counter) (counter > LOOP_TIMEOUT_LIMIT)
#define UPDATE_STATUS_IF_TIMEOUT_EXCEEDED(counter, status_var) \
	if(TIMEOUT_EXCEEDED(counter)) \
	{ \
		status_var = I2C::status::timeout; \
	}

/// Convenience macro to check if status is ok before writing bytes
// You could get the number of bytes written in an error scenario by NOT ignoring the
// second return parameter. You'd need a new macro var to assign it to.
#define WRITE_BYTES_IF_STATUS_OK(status_var, i2c_instance, operation) \
			if(status_var == I2C::status::ok) \
			{ \
				std::tie(status_var, std::ignore) = \
					i2c_write_bytes(i2c_instance, operation.tx_buffer, operation.tx_size); \
			}

/// Convenience macro to check if status is ok before reading bytes
// You could get the number of bytes read in an error scenario by NOT ignoring the
// second return parameter. You'd need a new macro var to assign it to.
#define READ_BYTES_IF_STATUS_OK(status_var, i2c_instance, operation) \
			if(status_var == I2C::status::ok) \
			{ \
				std::tie(status_var, std::ignore) = \
					i2c_read_bytes(i2c_instance, operation.rx_buffer, operation.rx_size); \
			}

static I2C::status i2c_wait_for_stop(I2C_TypeDef* const device)
{
	I2C::status wait_status = I2C::status::ok;
	unsigned loop_counter = 0;

	while(!LL_I2C_IsActiveFlag_STOP(device) && LOOP_TIMEOUT_CHECK(loop_counter))
	{
	}

	UPDATE_STATUS_IF_TIMEOUT_EXCEEDED(loop_counter, wait_status);

	LL_I2C_ClearFlag_STOP(device);

	return wait_status;
}

static I2C::status i2c_wait_for_TXE(I2C_TypeDef* const device)
{
	I2C::status wait_status = I2C::status::ok;
	unsigned loop_counter = 0;

	while(!LL_I2C_IsActiveFlag_TXE(device) && LOOP_TIMEOUT_CHECK(loop_counter))
	{
	}

	UPDATE_STATUS_IF_TIMEOUT_EXCEEDED(loop_counter, wait_status);

	return wait_status;
}

static I2C::status i2c_wait_for_TXIS(I2C_TypeDef* const device)
{
	I2C::status wait_status = I2C::status::ok;
	unsigned loop_counter = 0;

	while(!LL_I2C_IsActiveFlag_TXIS(device) && LOOP_TIMEOUT_CHECK(loop_counter))
	{
	}

	UPDATE_STATUS_IF_TIMEOUT_EXCEEDED(loop_counter, wait_status);

	return wait_status;
}

static I2C::status i2c_wait_for_RXNE(I2C_TypeDef* const device)
{
	I2C::status wait_status = I2C::status::ok;
	unsigned loop_counter = 0;

	while(!LL_I2C_IsActiveFlag_RXNE(device) && LOOP_TIMEOUT_CHECK(loop_counter))
	{
	}

	UPDATE_STATUS_IF_TIMEOUT_EXCEEDED(loop_counter, wait_status);

	return wait_status;
}


static I2C::status i2c_update_status_if_NACK(I2C_TypeDef* const device, I2C::status initial, NACK_CHECK nack)
{
	I2C::status nack_status = initial;

	if(LL_I2C_IsActiveFlag_NACK(device))
	{
		LL_I2C_ClearFlag_NACK(device);
		nack_status = (nack == NACK_CHECK::ADDR) ? I2C::status::addrNACK : I2C::status::dataNACK;
	}

	return nack_status;
}

// Function to write a byte to the I2C bus
static I2C::status i2c_write_byte(I2C_TypeDef* const device, uint8_t byte)
{
	I2C::status write_status = I2C::status::ok;

	// Write byte to data register
	LL_I2C_TransmitData8(device, byte);

	// Wait for byte transmission to complete
	write_status = i2c_wait_for_TXE(device);

	return i2c_update_status_if_NACK(device, write_status, NACK_CHECK::DATA);
}

// Function to read a byte from the I2C bus
static std::pair<I2C::status, uint8_t> i2c_read_byte(I2C_TypeDef* const device)
{
	I2C::status read_status = I2C::status::ok;
	uint8_t data_byte = 0x0;

	// Wait for byte reception to complete
	read_status = i2c_wait_for_RXNE(device);
	if(read_status == I2C::status::ok)
	{
		data_byte =	LL_I2C_ReceiveData8(device);
	}

	return {read_status, data_byte};
}

// Function to write multiple bytes to the I2C bus
// Returns the status and the number of bytes written
static std::pair<I2C::status, unsigned> i2c_write_bytes(I2C_TypeDef* const device,
	const uint8_t* data, unsigned size)
{
	I2C::status write_status = I2C::status::ok;
	unsigned written_count = 0;

  // Write data to the bus
  for (;
  	(written_count < size && (write_status == I2C::status::ok));
  	written_count++)
	{
		write_status = i2c_write_byte(device, data[written_count]);
		LL_I2C_SetTransferSize(device, size - written_count - 1);
	}

	return {write_status, written_count};
}

// Function to read multiple bytes from the I2C bus
// Returns the status and the number of bytes read
static std::pair<I2C::status, unsigned> i2c_read_bytes(I2C_TypeDef* const device,
	uint8_t* data, uint32_t size)
{
	I2C::status read_status = I2C::status::ok;
	unsigned read_count = 0;

	for (;
		(read_count < size && (read_status == I2C::status::ok));
		read_count++)
	{
		std::tie(read_status, data[read_count]) = i2c_read_byte(device);
		LL_I2C_SetTransferSize(device, size - read_count - 1);
	}

	return {read_status, read_count};
}

void STM32I2CMaster::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
	};

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

	started_ = true;
}

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

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

	LL_I2C_DeInit(i2c_inst);

	STM32ClockControl::i2cDisable(device_);

	started_ = false;
}

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

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

I2C::status STM32I2CMaster::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
	I2C::status op_status = I2C::status::ok;

	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);

	/** A Note on Large Transfers (> 255 bytes)
	 *
	 * This driver doesn't support them in the curent implementation. Sorry.
	 * If you're transferring that much data, use the DMA driver, which does support larger payloads.
	 * Or add support here. You can use the DMA driver's RELOAD behavior as a reference. Rather than
	 * reloading in an interrupt, you'll need to have the proper conditions at each write/read point
	 * below.
	 */
	assert(op.tx_size <= 255 && op.rx_size <= 255);

	switch(op.op)
	{
		case I2C::operation::continueWriteStop:
				LL_I2C_HandleTransfer(i2c_inst, address, LL_I2C_ADDRSLAVE_7BIT, op.tx_size, LL_I2C_MODE_AUTOEND,
						  LL_I2C_GENERATE_NOSTARTSTOP);
				// You could get the number of bytes written in an error scenario by NOT ignoring the
				// second return parameter.
				std::tie(op_status, std::ignore) = i2c_write_bytes(i2c_inst, op.tx_buffer, op.tx_size);
				op_status = i2c_wait_for_stop(i2c_inst);
				break;
		case I2C::operation::write:
			LL_I2C_HandleTransfer(i2c_inst, address, LL_I2C_ADDRSLAVE_7BIT, op.tx_size, LL_I2C_MODE_AUTOEND,
						  LL_I2C_GENERATE_START_WRITE);
			op_status = i2c_wait_for_TXIS(i2c_inst);
			op_status = i2c_update_status_if_NACK(i2c_inst, op_status, NACK_CHECK::ADDR);
			WRITE_BYTES_IF_STATUS_OK(op_status, i2c_inst, op);
			op_status = i2c_wait_for_stop(i2c_inst);
			break;
		case I2C::operation::writeNoStop:
			LL_I2C_HandleTransfer(i2c_inst, address, LL_I2C_ADDRSLAVE_7BIT, op.tx_size, LL_I2C_MODE_SOFTEND,
						  LL_I2C_GENERATE_NOSTARTSTOP);
			op_status = i2c_wait_for_TXIS(i2c_inst);
			op_status = i2c_update_status_if_NACK(i2c_inst, op_status, NACK_CHECK::ADDR);
			WRITE_BYTES_IF_STATUS_OK(op_status, i2c_inst, op);
			break;
		case I2C::operation::continueWriteNoStop: {
			LL_I2C_HandleTransfer(i2c_inst, address, LL_I2C_ADDRSLAVE_7BIT, op.tx_size, LL_I2C_MODE_SOFTEND,
						  LL_I2C_GENERATE_NOSTARTSTOP);
			// You could get the number of bytes written in an error scenario by NOT ignoring the
			// second return parameter.
			std::tie(op_status, std::ignore) = i2c_write_bytes(i2c_inst, op.tx_buffer, op.tx_size);
			break;
		}
		case I2C::operation::read: {
			LL_I2C_HandleTransfer(i2c_inst, address, LL_I2C_ADDRSLAVE_7BIT, op.rx_size, LL_I2C_MODE_AUTOEND,
						  LL_I2C_GENERATE_START_READ);
			op_status = i2c_wait_for_RXNE(i2c_inst);
			op_status = i2c_update_status_if_NACK(i2c_inst, op_status, NACK_CHECK::ADDR);
			READ_BYTES_IF_STATUS_OK(op_status, i2c_inst, op);
			break;
		}
		case I2C::operation::writeRead: {
			LL_I2C_HandleTransfer(i2c_inst, address, LL_I2C_ADDRSLAVE_7BIT, op.tx_size, LL_I2C_MODE_SOFTEND,
						  LL_I2C_GENERATE_START_WRITE);
			op_status = i2c_wait_for_TXIS(i2c_inst);
			op_status = i2c_update_status_if_NACK(i2c_inst, op_status, NACK_CHECK::ADDR);
			WRITE_BYTES_IF_STATUS_OK(op_status, i2c_inst, op);

			LL_I2C_HandleTransfer(i2c_inst, address, LL_I2C_ADDRSLAVE_7BIT, op.rx_size, LL_I2C_MODE_AUTOEND,
						  LL_I2C_GENERATE_START_READ);

			op_status = i2c_update_status_if_NACK(i2c_inst, op_status, NACK_CHECK::ADDR);
			READ_BYTES_IF_STATUS_OK(op_status, i2c_inst, op);
			break;
		}
		case I2C::operation::ping: {
			LL_I2C_HandleTransfer(i2c_inst, address, LL_I2C_ADDRSLAVE_7BIT, 1, LL_I2C_MODE_AUTOEND,
						  LL_I2C_GENERATE_START_READ);

			op_status = i2c_wait_for_stop(i2c_inst);
			op_status = i2c_update_status_if_NACK(i2c_inst, op_status, NACK_CHECK::ADDR);
			break;
		}
		case I2C::operation::stop: {
			LL_I2C_GenerateStopCondition(i2c_inst);
			break;
		}
		case I2C::operation::restart: {
			LL_I2C_GenerateStartCondition(i2c_inst);
			break;
		}
	}

	busy_ = false;

	if(cb)
	{
		CentralEventQueue::callback(std::bind(cb, op_status));
	}

	return op_status;
}
