#include "lcd2004.hpp"

#include "pins.h"
#include "blk_box_drivers/i2c.h"
#include "driver/i2c_master.h"
#include "driver/gpio.h"
#include "freertos/FreeRTOS.h"

#include "esp_log.h"
#include "esp_rom_sys.h"

#include <stdint.h>

const static char* TAG = "LCD2004";

// masks
const uint8_t ENABLE_MASK = 0x04;
const uint8_t BACKLIGHT_MASK = 0x08;
const uint8_t DATA_MASK = 0x01;

const uint8_t CMD_CLEAR_DISPLAY = 0x01;

const uint8_t CMD_RETURN_HOME = 0x02;

const uint8_t CMD_ENTRY_MODE_SET = 0x04;
const uint8_t ENTRY_MODE_INC_DEC = 0x02;
const uint8_t ENTRY_MODE_SHIFT = 0x01;

const uint8_t CMD_DISPLAY_CONTROL = 0x08;
const uint8_t DISPLAY_CONTROL_SHOW = 0x04;
const uint8_t DISPLAY_CONTROL_CURSOR = 0x02;
const uint8_t DISPLAY_CONTROL_BLINK = 0x01;

const uint8_t CMD_CURSOR_OR_DISPLAY_SHIFT = 0x10;
const uint8_t CURSOR_OR_DISPLAY_SHIFT_SC = 0x08;
const uint8_t CURSOR_OR_DISPLAY_SHIFT_RL = 0x04;

const uint8_t CMD_FUNCTION_SET = 0x20;
const uint8_t FUNCTION_SET_DL = 0x10;
const uint8_t FUNCTION_SET_N = 0x08;
const uint8_t FUNCTION_SET_F = 0x04;

const uint8_t CMD_SET_CGRAM_ADDRESS = 0x40;
const uint8_t CMD_SET_DDRAM_ADDRESS = 0x80;

LCD2004I2C::LCD2004I2C() {
    backlight = 0;
    display_control = 0;
    entry_mode = 0;
}

esp_err_t LCD2004I2C::init(uint8_t addr, bool initial_backlight) {
    if (initialized) return ESP_ERR_INVALID_STATE;
    i2c_device_config_t dev_config = {
        .dev_addr_length = I2C_ADDR_BIT_LEN_7,
        .device_address = addr,
        .scl_speed_hz = LCD2004_CLK_SPEED_HZ,
        .scl_wait_us = 0, // default
        .flags = {
            .disable_ack_check = 0,
        }
    };

    esp_err_t err = i2c_master_bus_add_device(i2c_main_bus_handle, &dev_config, &lcd_device_handle);
    if (err != ESP_OK) return err;
    ESP_LOGD(TAG, "LCD2004 device added to bus");
    backlight = initial_backlight ? BACKLIGHT_MASK : 0;

    // Initialize the LCD
    // set to 8-bit mode to bring it to a known state
    for (int i = 0; i < 3; i++) {
        err = send4_cmd(CMD_FUNCTION_SET | FUNCTION_SET_DL);
        if (err != ESP_OK) return err;
        vTaskDelay(pdMS_TO_TICKS(5));
    }
    
    // set to 4-bit mode, while in 8 bit mode
    err = send4_cmd(CMD_FUNCTION_SET);
    if (err != ESP_OK) return err;

    // use all 4 rows
    err = send_cmd(CMD_FUNCTION_SET | FUNCTION_SET_N);
    if (err != ESP_OK) return err;

    // display, set cursor visible
    err = send_cmd(CMD_DISPLAY_CONTROL | DISPLAY_CONTROL_SHOW);
    if (err != ESP_OK) return err;

    // clear display
    err = send_cmd(CMD_CLEAR_DISPLAY);
    if (err != ESP_OK) return err;
    esp_rom_delay_us(1600);

    // set cursor to increment to the right
    err = send_cmd(CMD_ENTRY_MODE_SET | ENTRY_MODE_INC_DEC);
    if (err != ESP_OK) return err;

    display_control = DISPLAY_CONTROL_SHOW;
    entry_mode = ENTRY_MODE_INC_DEC;
    initialized = true;

    ESP_LOGI(TAG, "LCD2004 initialized!");
    return ESP_OK;
}

// i2c wrappers

/// Sends `cmd` to the device as a command.
esp_err_t LCD2004I2C::send_cmd(uint8_t cmd) {
    uint8_t control_bits = backlight;

    uint8_t high = control_bits | (cmd & 0xF0);
    uint8_t high_enable = high | ENABLE_MASK;
    uint8_t low = control_bits | ((cmd << 4) & 0xF0);
    uint8_t low_enable = low | ENABLE_MASK;

    uint8_t i2c_data[] = {
        high_enable,
        high,
        low_enable,
        low
    };

    return i2c_master_transmit(lcd_device_handle, i2c_data, sizeof(i2c_data), I2C_MASTER_TIMEOUT_MS);
}

esp_err_t LCD2004I2C::send_data(uint8_t data) {
    uint8_t control_bits = DATA_MASK | backlight;

    uint8_t high = control_bits | (data & 0xF0);
    uint8_t high_enable = high | ENABLE_MASK;
    uint8_t low = control_bits | ((data << 4) & 0xF0);
    uint8_t low_enable = low | ENABLE_MASK;

    // data pins are on pins D7 to D4
    uint8_t i2c_data[] = {
        high_enable,
        high,
        low_enable,
        low
    };

    return i2c_master_transmit(lcd_device_handle, i2c_data, sizeof(i2c_data), I2C_MASTER_TIMEOUT_MS);
}

/// Sends the ***HIGH*** 4 bits of `cmd` as a command.
esp_err_t LCD2004I2C::send4_cmd(uint8_t cmd) {
    uint8_t control_bits = backlight;

    uint8_t high = control_bits | (cmd & 0xF0);
    uint8_t high_enable = high | ENABLE_MASK;

    // data pins are on pins D7 to D4
    uint8_t i2c_data[] = {
        high_enable,
        high
    };

    return i2c_master_transmit(lcd_device_handle, i2c_data, sizeof(i2c_data), I2C_MASTER_TIMEOUT_MS);
}

// LCD2004 class implementations

bool LCD2004I2C::get_backlight() {
    return backlight != 0;
}

esp_err_t LCD2004I2C::set_backlight(bool backlight) {
    if (get_backlight() == backlight) {
        return ESP_OK;
    }

    uint8_t data = backlight ? BACKLIGHT_MASK : 0;
    esp_err_t err = i2c_master_transmit(lcd_device_handle, &data, 1, I2C_MASTER_TIMEOUT_MS);
    if (err == ESP_OK) {
        this->backlight = data;
    }
    return err;
}

esp_err_t LCD2004I2C::clear() {
    esp_err_t err = send_cmd(CMD_CLEAR_DISPLAY);
    if (err != ESP_OK) return err;
    esp_rom_delay_us(1600);
    entry_mode |= ENTRY_MODE_INC_DEC;
    return ESP_OK;
}

esp_err_t LCD2004I2C::move_cursor(uint8_t row, uint8_t col) {
    const uint8_t ROW_OFFSETS[4] = {0x00, 0x40, 0x14, 0x54};
    if (row >= LCD2004_ROWS || col >= LCD2004_COLS) return ESP_ERR_INVALID_ARG;
    uint8_t addr = ROW_OFFSETS[row] + col;
    return send_cmd(CMD_SET_DDRAM_ADDRESS | addr);
}

bool LCD2004I2C::get_show_hide() {
    return (display_control & DISPLAY_CONTROL_SHOW) != 0;
}

esp_err_t LCD2004I2C::show_hide(bool show) {
    if (get_show_hide() == show) {
        return ESP_OK;
    }

    if (show) {
        display_control |= DISPLAY_CONTROL_SHOW;
    } else {
        display_control &= ~DISPLAY_CONTROL_SHOW;
    }
    return send_cmd(CMD_DISPLAY_CONTROL | display_control);
}

bool LCD2004I2C::get_show_cursor() {
    return (display_control & DISPLAY_CONTROL_CURSOR) != 0;
}

esp_err_t LCD2004I2C::show_cursor(bool show_cursor) {
    if (get_show_cursor() == show_cursor) {
        return ESP_OK;
    }

    if (show_cursor) {
        display_control |= DISPLAY_CONTROL_CURSOR;
    } else {
        display_control &= ~DISPLAY_CONTROL_CURSOR;
    }
    return send_cmd(CMD_DISPLAY_CONTROL | display_control);
}

bool LCD2004I2C::get_blink_cursor() {
    return (display_control & DISPLAY_CONTROL_BLINK) != 0;
}

esp_err_t LCD2004I2C::blink_cursor(bool blink_cursor) {
    if (get_blink_cursor() == blink_cursor) {
        return ESP_OK;
    }

    if (blink_cursor) {
        display_control |= DISPLAY_CONTROL_BLINK;
    } else {
        display_control &= ~DISPLAY_CONTROL_BLINK;
    }
    return send_cmd(CMD_DISPLAY_CONTROL | display_control);
}

esp_err_t LCD2004I2C::show_blink_cursor(bool show_cursor, bool blink_cursor) {
    if (get_show_cursor() == show_cursor && get_blink_cursor() == blink_cursor) {
        return ESP_OK;
    }

    if (blink_cursor) {
        display_control |= DISPLAY_CONTROL_BLINK;
    } else {
        display_control &= ~DISPLAY_CONTROL_BLINK;
    }
    if (show_cursor) {
        display_control |= DISPLAY_CONTROL_CURSOR;
    } else {
        display_control &= ~DISPLAY_CONTROL_CURSOR;
    }
    return send_cmd(CMD_DISPLAY_CONTROL | display_control);
}

bool LCD2004I2C::get_scroll_direction() {
    return (entry_mode & ENTRY_MODE_INC_DEC) != 0;
}

esp_err_t LCD2004I2C::scroll_direction(bool left_to_right) {
    if (get_scroll_direction() == left_to_right) {
        return ESP_OK;
    }

    if (left_to_right) {
        entry_mode |= ENTRY_MODE_INC_DEC;
    } else {
        entry_mode &= ~ENTRY_MODE_INC_DEC;
    }
    return send_cmd(CMD_ENTRY_MODE_SET | entry_mode);
}

bool LCD2004I2C::get_display_shift() {
    return (entry_mode & ENTRY_MODE_SHIFT) != 0;
}

esp_err_t LCD2004I2C::shift_display(bool shift_display) {
    if (get_display_shift() == shift_display) {
        return ESP_OK;
    }

    if (shift_display) {
        entry_mode |= ENTRY_MODE_SHIFT;
    } else {
        entry_mode &= ~ENTRY_MODE_SHIFT;
    }
    return send_cmd(CMD_ENTRY_MODE_SET | entry_mode);
}

esp_err_t LCD2004I2C::shift_cursor_left() {
    return send_cmd(CMD_CURSOR_OR_DISPLAY_SHIFT);
}

esp_err_t LCD2004I2C::shift_cursor_right() {
    return send_cmd(CMD_CURSOR_OR_DISPLAY_SHIFT | CURSOR_OR_DISPLAY_SHIFT_RL);
}

esp_err_t LCD2004I2C::shift_display_left() {
    return send_cmd(CMD_CURSOR_OR_DISPLAY_SHIFT | CURSOR_OR_DISPLAY_SHIFT_SC);
}

esp_err_t LCD2004I2C::shift_display_right() {
    return send_cmd(CMD_CURSOR_OR_DISPLAY_SHIFT | CURSOR_OR_DISPLAY_SHIFT_SC | CURSOR_OR_DISPLAY_SHIFT_RL);
}

esp_err_t LCD2004I2C::create_custom_char(uint8_t location, const uint8_t char_map[8]) {
    if (char_map == nullptr) return ESP_ERR_INVALID_ARG;
    uint8_t loc = location % 8;
    esp_err_t err = send_cmd(CMD_SET_CGRAM_ADDRESS | (loc << 3));
    if (err != ESP_OK) return err;
    for (int i = 0; i < 8; i++) {
        err = send_data(char_map[i]);
        if (err != ESP_OK) return err;
    }
    return ESP_OK;
}

esp_err_t LCD2004I2C::print_char(char c) {
    if (c == '\x08') {
        // custom char 0 is represented by \x08 since \x00 is a null terminator
        c = 0;
    }
    return send_data(static_cast<uint8_t>(c));
}

esp_err_t LCD2004I2C::print_u8(uint8_t b) {
    if (b == 0x08) {
        // custom char 0 is represented by \x08 since \x00 is a null terminator
        b = 0;
    }
    return send_data(b);
}

esp_err_t LCD2004I2C::write_str(const char* s) {
    // TODO: replace this with a single I2C transmission.
    if (s == nullptr) return ESP_ERR_INVALID_ARG;
    while (*s) {
        esp_err_t err = print_char(*s);
        if (err != ESP_OK) return err;
        s++;
    }
    return ESP_OK;
}