/*
 * Copyright (c) 2021 Arm Limited and Contributors. All rights reserved.
 *
 * SPDX-License-Identifier: Apache-2.0
 * 
 * ST7789 TFT LCD Display Driver
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "hardware/gpio.h"
#include "hardware/spi.h"
#include "pico/binary_info.h"
#include "pico/stdlib.h"

#include "st7789.h"

// ST7789 Commands
#define ST7789_NOP        0x00
#define ST7789_SWRESET    0x01
#define ST7789_RDDID      0x04
#define ST7789_RDDST      0x09

#define ST7789_SLPIN      0x10
#define ST7789_SLPOUT     0x11
#define ST7789_PTLON      0x12
#define ST7789_NORON      0x13

#define ST7789_INVOFF     0x20
#define ST7789_INVON      0x21
#define ST7789_DISPOFF    0x28
#define ST7789_DISPON     0x29
#define ST7789_CASET      0x2A
#define ST7789_RASET      0x2B
#define ST7789_RAMWR      0x2C
#define ST7789_RAMRD      0x2E

#define ST7789_PTLAR      0x30
#define ST7789_VSCRDEF    0x33
#define ST7789_MADCTL     0x36
#define ST7789_VSCRSADD   0x37
#define ST7789_COLMOD     0x3A

#define ST7789_MADCTL_MY  0x80
#define ST7789_MADCTL_MX  0x40
#define ST7789_MADCTL_MV  0x20
#define ST7789_MADCTL_ML  0x10
#define ST7789_MADCTL_RGB 0x00

#define ST7789_RDID1      0xDA
#define ST7789_RDID2      0xDB
#define ST7789_RDID3      0xDC
#define ST7789_RDID4      0xDD

static const struct st7789_config *config;
static uint16_t width;
static uint16_t height;
static uint16_t x_offset;
static uint16_t y_offset;

static inline void cs_select() {
    if (config->gpio_cs >= 0) {
        asm volatile("nop \n nop \n nop");
        gpio_put(config->gpio_cs, 0);
        asm volatile("nop \n nop \n nop");
    }
}

static inline void cs_deselect() {
    if (config->gpio_cs >= 0) {
        asm volatile("nop \n nop \n nop");
        gpio_put(config->gpio_cs, 1);
        asm volatile("nop \n nop \n nop");
    }
}

static inline void dc_command() {
    asm volatile("nop \n nop \n nop");
    gpio_put(config->gpio_dc, 0);
    asm volatile("nop \n nop \n nop");
}

static inline void dc_data() {
    asm volatile("nop \n nop \n nop");
    gpio_put(config->gpio_dc, 1);
    asm volatile("nop \n nop \n nop");
}

static inline void reset_pulse() {
    gpio_put(config->gpio_rst, 1);
    sleep_ms(5);
    gpio_put(config->gpio_rst, 0);
    sleep_ms(20);
    gpio_put(config->gpio_rst, 1);
    sleep_ms(150);
}

static void write_command(uint8_t cmd) {
    dc_command();
    cs_select();
    spi_write_blocking(config->spi, &cmd, 1);
    cs_deselect();
}

static void write_data(const uint8_t *data, size_t len) {
    dc_data();
    cs_select();
    spi_write_blocking(config->spi, data, len);
    cs_deselect();
}

static void write_command_with_data(uint8_t cmd, const uint8_t *data, size_t len) {
    write_command(cmd);
    write_data(data, len);
}

static void set_window(uint16_t x0, uint16_t y0, uint16_t x1, uint16_t y1) {
    uint8_t data[4];

    // Add offsets for display positioning
    x0 += x_offset;
    x1 += x_offset;
    y0 += y_offset;
    y1 += y_offset;

    // Column address set
    data[0] = (x0 >> 8) & 0xFF;
    data[1] = x0 & 0xFF;
    data[2] = (x1 >> 8) & 0xFF;
    data[3] = x1 & 0xFF;
    write_command_with_data(ST7789_CASET, data, 4);

    // Row address set
    data[0] = (y0 >> 8) & 0xFF;
    data[1] = y0 & 0xFF;
    data[2] = (y1 >> 8) & 0xFF;
    data[3] = y1 & 0xFF;
    write_command_with_data(ST7789_RASET, data, 4);

    // Write to RAM
    write_command(ST7789_RAMWR);
}

void st7789_init(const struct st7789_config *c, uint16_t w, uint16_t h) {
    config = c;
    width = w;
    height = h;
    
    // Set offsets for 240x240 display on ST7789
    // The Adafruit 1.54" TFT needs these offsets
    x_offset = 0;
    y_offset = 80;  // 240x240 display with 80 pixel row offset

    // Initialize SPI
    spi_init(config->spi, 62500 * 1000); // 62.5 MHz
    gpio_set_function(config->gpio_din, GPIO_FUNC_SPI);
    gpio_set_function(config->gpio_clk, GPIO_FUNC_SPI);

    // Initialize CS pin
    if (config->gpio_cs >= 0) {
        gpio_init(config->gpio_cs);
        gpio_set_dir(config->gpio_cs, GPIO_OUT);
        gpio_put(config->gpio_cs, 1);
    }

    // Initialize DC pin
    gpio_init(config->gpio_dc);
    gpio_set_dir(config->gpio_dc, GPIO_OUT);

    // Initialize RST pin
    gpio_init(config->gpio_rst);
    gpio_set_dir(config->gpio_rst, GPIO_OUT);

    // Initialize backlight pin
    gpio_init(config->gpio_bl);
    gpio_set_dir(config->gpio_bl, GPIO_OUT);
    gpio_put(config->gpio_bl, 1); // Turn on backlight

    // Reset display
    reset_pulse();

    // Software reset
    write_command(ST7789_SWRESET);
    sleep_ms(150);

    // Out of sleep mode
    write_command(ST7789_SLPOUT);
    sleep_ms(10);

    // Set color mode to 16-bit (RGB565)
    uint8_t colmod_data = 0x05;
    write_command_with_data(ST7789_COLMOD, &colmod_data, 1);
    sleep_ms(10);

    // Memory data access control
    uint8_t madctl_data = ST7789_MADCTL_MX | ST7789_MADCTL_MY | ST7789_MADCTL_RGB;
    write_command_with_data(ST7789_MADCTL, &madctl_data, 1);

    // Normal display mode on
    write_command(ST7789_NORON);
    sleep_ms(10);

    // Display on
    write_command(ST7789_DISPON);
    sleep_ms(10);
}

void st7789_fill(uint16_t color) {
    set_window(0, 0, width - 1, height - 1);

    dc_data();
    cs_select();

    // Send color data in chunks for better performance
    uint8_t data[2] = {(color >> 8) & 0xFF, color & 0xFF};
    uint32_t pixel_count = width * height;
    
    // Send pixels in batches
    uint8_t buffer[256];  // 128 pixels at a time
    for (int i = 0; i < 128; i++) {
        buffer[i * 2] = data[0];
        buffer[i * 2 + 1] = data[1];
    }
    
    uint32_t full_chunks = pixel_count / 128;
    uint32_t remaining = pixel_count % 128;
    
    for (uint32_t i = 0; i < full_chunks; i++) {
        spi_write_blocking(config->spi, buffer, 256);
    }
    
    if (remaining > 0) {
        spi_write_blocking(config->spi, buffer, remaining * 2);
    }

    cs_deselect();
}

void st7789_put(uint16_t color) {
    uint8_t data[2] = {(color >> 8) & 0xFF, color & 0xFF};
    
    dc_data();
    cs_select();
    spi_write_blocking(config->spi, data, 2);
    cs_deselect();
}

void st7789_set_cursor(uint16_t x, uint16_t y) {
    set_window(x, y, width - 1, height - 1);
}

void st7789_write(const uint16_t *data, size_t len) {
    dc_data();
    cs_select();
    
    for (size_t i = 0; i < len; i++) {
        uint8_t bytes[2] = {(data[i] >> 8) & 0xFF, data[i] & 0xFF};
        spi_write_blocking(config->spi, bytes, 2);
    }
    
    cs_deselect();
}

void st7789_vertical_scroll(uint16_t row) {
    uint8_t data[2] = {(row >> 8) & 0xFF, row & 0xFF};
    write_command_with_data(ST7789_VSCRSADD, data, 2);
}

void st7789_draw_pixel(uint16_t x, uint16_t y, uint16_t color) {
    if (x >= width || y >= height) return;
    set_window(x, y, x, y);
    uint8_t data[2] = {(color >> 8) & 0xFF, color & 0xFF};
    dc_data();
    cs_select();
    spi_write_blocking(config->spi, data, 2);
    cs_deselect();
}

void st7789_draw_rect(uint16_t x, uint16_t y, uint16_t w, uint16_t h, uint16_t color) {
    // Top and bottom horizontal lines
    st7789_fill_rect(x, y, w, 1, color);
    st7789_fill_rect(x, y + h - 1, w, 1, color);
    // Left and right vertical lines
    st7789_fill_rect(x, y, 1, h, color);
    st7789_fill_rect(x + w - 1, y, 1, h, color);
}

void st7789_fill_rect(uint16_t x, uint16_t y, uint16_t w, uint16_t h, uint16_t color) {
    if (x >= width || y >= height) return;
    if (x + w > width) w = width - x;
    if (y + h > height) h = height - y;
    
    set_window(x, y, x + w - 1, y + h - 1);
    
    dc_data();
    cs_select();
    
    uint8_t data[2] = {(color >> 8) & 0xFF, color & 0xFF};
    uint32_t pixel_count = w * h;
    
    // Send pixels in batches
    uint8_t buffer[256];
    for (int i = 0; i < 128; i++) {
        buffer[i * 2] = data[0];
        buffer[i * 2 + 1] = data[1];
    }
    
    uint32_t full_chunks = pixel_count / 128;
    uint32_t remaining = pixel_count % 128;
    
    for (uint32_t i = 0; i < full_chunks; i++) {
        spi_write_blocking(config->spi, buffer, 256);
    }
    
    if (remaining > 0) {
        spi_write_blocking(config->spi, buffer, remaining * 2);
    }
    
    cs_deselect();
}

void st7789_draw_circle(uint16_t x0, uint16_t y0, uint16_t r, uint16_t color) {
    int16_t f = 1 - r;
    int16_t ddF_x = 1;
    int16_t ddF_y = -2 * r;
    int16_t x = 0;
    int16_t y = r;
    
    st7789_draw_pixel(x0, y0 + r, color);
    st7789_draw_pixel(x0, y0 - r, color);
    st7789_draw_pixel(x0 + r, y0, color);
    st7789_draw_pixel(x0 - r, y0, color);
    
    while (x < y) {
        if (f >= 0) {
            y--;
            ddF_y += 2;
            f += ddF_y;
        }
        x++;
        ddF_x += 2;
        f += ddF_x;
        
        st7789_draw_pixel(x0 + x, y0 + y, color);
        st7789_draw_pixel(x0 - x, y0 + y, color);
        st7789_draw_pixel(x0 + x, y0 - y, color);
        st7789_draw_pixel(x0 - x, y0 - y, color);
        st7789_draw_pixel(x0 + y, y0 + x, color);
        st7789_draw_pixel(x0 - y, y0 + x, color);
        st7789_draw_pixel(x0 + y, y0 - x, color);
        st7789_draw_pixel(x0 - y, y0 - x, color);
    }
}

void st7789_fill_circle(uint16_t x0, uint16_t y0, uint16_t r, uint16_t color) {
    int16_t f = 1 - r;
    int16_t ddF_x = 1;
    int16_t ddF_y = -2 * r;
    int16_t x = 0;
    int16_t y = r;
    
    st7789_fill_rect(x0 - r, y0, 2 * r + 1, 1, color);
    
    while (x < y) {
        if (f >= 0) {
            y--;
            ddF_y += 2;
            f += ddF_y;
        }
        x++;
        ddF_x += 2;
        f += ddF_x;
        
        st7789_fill_rect(x0 - x, y0 + y, 2 * x + 1, 1, color);
        st7789_fill_rect(x0 - x, y0 - y, 2 * x + 1, 1, color);
        st7789_fill_rect(x0 - y, y0 + x, 2 * y + 1, 1, color);
        st7789_fill_rect(x0 - y, y0 - x, 2 * y + 1, 1, color);
    }
}

void st7789_draw_line(uint16_t x0, uint16_t y0, uint16_t x1, uint16_t y1, uint16_t color) {
    int16_t dx = abs(x1 - x0);
    int16_t dy = abs(y1 - y0);
    int16_t sx = (x0 < x1) ? 1 : -1;
    int16_t sy = (y0 < y1) ? 1 : -1;
    int16_t err = dx - dy;
    
    while (1) {
        st7789_draw_pixel(x0, y0, color);
        
        if (x0 == x1 && y0 == y1) break;
        
        int16_t e2 = 2 * err;
        if (e2 > -dy) {
            err -= dy;
            x0 += sx;
        }
        if (e2 < dx) {
            err += dx;
            y0 += sy;
        }
    }
}