#pragma once

#include <utility>

#include <cstdint>

#include <avr/pgmspace.h>

#include "../clock.hpp"
#include "../io/io.hpp"
#include "../util/util.hpp"

namespace eink {

template <std::uint16_t Width, std::uint16_t Height, typename Spi, io::P RstPin, io::P DcPin, io::P BusyPin>
class Eink {
	using word_t = typename Spi::word_t;

	enum class Cmd : std::uint8_t {
		DRIVER_OUTPUT_CONTROL = 0x01,
		DEEP_SLEEP_MODE = 0x10,
		DATA_ENTRY_MODE = 0x11,
		SW_RESET = 0x12,
		READ_TEMPERATURE_SENSOR = 0x18,
		UPDATE_DISPLAY = 0x20,
		DISPLAY_UPDATE_CONTROL_2 = 0x22,
		WRITE_RAM_BLACK = 0x24,
		WRITE_RAM_RED = 0x26,
		BORDER_WAVEFORM_CONTROL = 0x3C,
		SET_RAM_X_ADDR_POSITIONS = 0x44,
		SET_RAM_Y_ADDR_POSITIONS = 0x45,
		SET_RAM_X_ADDR = 0x4E,
		SET_RAM_Y_ADDR = 0x4F,
	};

	static io::Pin<RstPin> m_rst;
	static io::Pin<DcPin> m_dc;
	static io::Pin<BusyPin> m_bsy;
	static constexpr auto THREE_WIRE_SPI = (DcPin == io::P::NONE ? true : false);
	static_assert((THREE_WIRE_SPI && sizeof(word_t) > 1) || !THREE_WIRE_SPI,
	              "Three wire SPI requires SPI word size of at least 9 bits");

	static constexpr auto BLOCK_SIZE = 5;

	enum class Color : std::uint8_t {
		BLACK = 0b00,
		WHITE = 0b01,
		RED = 0b10,
		ERROR = 0b11,
	};

	class ImageBlock {
	  public:
		inline Color &operator[](const std::size_t idx)
		{
			return data[idx];
		}

		inline const Color &operator[](const std::size_t idx) const
		{
			return data[idx];
		}

	  private:
		Color data[BLOCK_SIZE];
	};

  public:
	static void init()
	{
		m_rst.dir(io::Dir::OUT);
		m_dc.dir(io::Dir::OUT);
		m_bsy.dir(io::Dir::IN);
		Spi::init();

		reset();

		waitUntilIdle();
		sendCommand(Cmd::SW_RESET);
		waitUntilIdle();

		sendCommand(Cmd::DRIVER_OUTPUT_CONTROL);
		sendData(0xC7);
		sendData(0x00);
		sendData(0x01);

		sendCommand(Cmd::DATA_ENTRY_MODE);
		sendData(0x02);

		sendCommand(Cmd::SET_RAM_X_ADDR_POSITIONS);
		sendData(Width / 8 - 1);
		sendData(0x00);

		sendCommand(Cmd::SET_RAM_Y_ADDR_POSITIONS);
		sendData(0x00);
		sendData(0x00);
		sendData(Height - 1);
		sendData(0x00);

		sendCommand(Cmd::BORDER_WAVEFORM_CONTROL);
		sendData(0x05);

		sendCommand(Cmd::READ_TEMPERATURE_SENSOR);
		sendData(0x80);

		sendCommand(Cmd::SET_RAM_X_ADDR);
		sendData(Width / 8 - 1);

		sendCommand(Cmd::SET_RAM_Y_ADDR);
		sendData(0x00);
		sendData(0x00);

		waitUntilIdle();
	}

	static void sendCommand(const Cmd command)
	{
		m_dc = false;
		spiTransfer(static_cast<word_t>(command));
	}

	static void sendData(word_t data)
	{
		if constexpr (THREE_WIRE_SPI) {
			data |= 1 << 8;
		}

		m_dc = true;
		spiTransfer(data);
	}

	static void waitUntilIdle()
	{
		while (m_bsy) {
			_delay_ms(100);
		}
	}

	static void reset()
	{
		m_rst = true;
		_delay_ms(200);
		m_rst = false;
		_delay_ms(10);
		m_rst = true;
		_delay_ms(200);
	}

	template <typename Image>
	static void draw(const Image &image)
	{
		constexpr auto lookup = [](std::uint8_t bits) {
			auto block = ImageBlock{};
			util::for_constexpr(
			    [&](const auto idx) {
				    block[idx.value] = static_cast<Color>(bits % 3);
				    bits /= 3;
			    },
			    std::make_index_sequence<BLOCK_SIZE>{});
			return block;
		};

		constexpr auto sendImageChannel = [lookup](const auto command, const auto image) {
			sendCommand(command);
			auto buffer = std::uint8_t{0};
			auto bufferPos = std::uint8_t{0};
			for (auto i = std::uint16_t{0}; i < Width * Height / BLOCK_SIZE; i++) {
				const auto block = lookup(pgm_read_byte(&image[i]));
				for (auto p = std::uint8_t{0}; p < BLOCK_SIZE; ++p) {
					const auto pixel = std::uint8_t{(command == Cmd::WRITE_RAM_BLACK) ? (block[p] != Color::BLACK)
					                                                                  : (block[p] == Color::RED)};
					buffer |= pixel << (7 - bufferPos++);
					if (bufferPos == 8) {
						sendData(buffer);
						buffer = 0;
						bufferPos = 0;
					}
				}
			}
		};

		sendImageChannel(Cmd::WRITE_RAM_BLACK, image.data());
		sendImageChannel(Cmd::WRITE_RAM_RED, image.data());

		sendCommand(Cmd::DISPLAY_UPDATE_CONTROL_2);
		sendData(0xF7);
		sendCommand(Cmd::UPDATE_DISPLAY);
		waitUntilIdle();
	}

	static void clear()
	{
		sendCommand(Cmd::WRITE_RAM_BLACK);
		for (auto i = std::uint16_t{0}; i < Width * Height / 8; i++) {
			sendData(0xff);
		}
		sendCommand(Cmd::WRITE_RAM_RED);
		for (auto i = std::uint16_t{0}; i < Width * Height / 8; i++) {
			sendData(0x00);
		}

		sendCommand(Cmd::DISPLAY_UPDATE_CONTROL_2);
		sendData(0xF7);
		sendCommand(Cmd::UPDATE_DISPLAY);
		waitUntilIdle();
	}

	static void sleep()
	{
		sendCommand(Cmd::DEEP_SLEEP_MODE);
		sendData(0x01);
		_delay_ms(100);
	}

	static void spiTransfer(const word_t data)
	{
		Spi::select(true);
		Spi::transfer(data);
		Spi::select(false);
	}
};

} // namespace eink