#include <string.h>
#include <avr/io.h>
#include <avr/pgmspace.h>
#include <util/delay.h>
#include "pt6524.h"
#include "spi.h"

// PT6524 Command/Address
// Datasheet specifies 41H, sent with DORD=1 (LSB-first)
// Hardware automatically reverses 0x41 to 0x82 on the wire
#define PT6524_ADDRESS  0x41

// Control bits
#define CTRL_BU   (1 << 0)  // Normal/Power Saving Mode (0=Normal, 1=Power Saving)
#define CTRL_SC   (1 << 1)  // Segment ON/OFF (0=ON, 1=OFF)
#define CTRL_DR   (1 << 2)  // Bias Drive (0=1/3 bias, 1=1/2 bias)
#define CTRL_P3   (1 << 3)  // Port select bit 3
#define CTRL_P2   (1 << 4)  // Port select bit 2
#define CTRL_P1   (1 << 5)  // Port select bit 1
#define CTRL_P0   (1 << 6)  // Port select bit 0
#define CTRL_CU   (1 << 7)  // Current drain control (0=Normal, 1=Low)

// Display framebuffer - stores segment states
// Each bit represents one segment/COM intersection
// Organized as: D1, D2, D3, D4, D5, D6, ... D160
// Where D1=SG1:COM1, D2=SG1:COM2, D3=SG1:COM3, D4=SG1:COM4,
//       D5=SG2:COM1, D6=SG2:COM2, etc.
static uint8_t framebuffer[(PT6524_TOTAL_SEGMENTS * PT6524_TOTAL_COMS + 7) / 8];

// Initialize PT6524 LCD driver
void pt6524_init(void) {
	// Set DISP_ON low BEFORE configuring as output
	DISP_ON_PORT &= ~DISP_ON;
	DISP_ON_DDR |= DISP_ON;

    // Initialize SPI
    spi_init();

    // Keep display off during initialization
    DISP_ON_LOW();

    pt6524_clear_all();

    // Enable display
    DISP_ON_HIGH();
}

// Set a specific segment on or off
// address: Linear address (0 to PT6524_MAX_ADDRESS-1)
//          Address 0 = SG1:COM1, 1 = SG1:COM2, 2 = SG1:COM3, 3 = SG1:COM4,
//          Address 4 = SG2:COM1, 5 = SG2:COM2, etc.
// state: 0 = off, non-zero = on
void pt6524_set_segment(uint8_t address, uint8_t state) {
    // Validate address against actual framebuffer size
    if (address >= PT6524_MAX_ADDRESS)
        return;

    // Calculate bit position in framebuffer
    // Linear address directly maps to bit index in framebuffer
    uint8_t byte_pos = address / 8;
    uint8_t bit_pos = address % 8;

    // Set or clear the bit
    if (state) {
        framebuffer[byte_pos] |= (1 << bit_pos);
    } else {
        framebuffer[byte_pos] &= ~(1 << bit_pos);
    }
}

// Clear all segments
void pt6524_clear_all(void) {
    memset(framebuffer, 0, sizeof(framebuffer));
}

// Get direct access to framebuffer for efficient bulk updates
uint8_t* pt6524_get_framebuffer(void) {
    return framebuffer;
}

// Write a single byte to framebuffer
void pt6524_write_byte(uint8_t byte_index, uint8_t value) {
    if (byte_index < sizeof(framebuffer))
        framebuffer[byte_index] = value;
}

// Write a block of bytes to framebuffer
void pt6524_write_block(uint8_t byte_index, const uint8_t* data, uint8_t length) {
    if (byte_index + length <= sizeof(framebuffer))
        memcpy(&framebuffer[byte_index], data, length);
}

#define DATA_BITS_PER_BLOCK 52

static inline void pt6524_send_block(uint8_t dir_bits, uint16_t start_addr) {
    const uint8_t control = 0x00; // DR=0 for 1/3 bias

    _delay_us(.16);
    spi_write(PT6524_ADDRESS);
    CE_HIGH();

    uint8_t data_payload[7] = {0}; // 52 bits = 6.5 bytes

    // Determine how many bits to send in this block
    uint16_t bits_to_send = PT6524_MAX_ADDRESS - start_addr;
    if (bits_to_send > DATA_BITS_PER_BLOCK) {
        bits_to_send = DATA_BITS_PER_BLOCK;
    }

    // Prepare data payload from framebuffer
    uint16_t start_byte_idx = start_addr / 8;
    uint8_t start_bit_offset = start_addr % 8;
    uint16_t current_byte_idx = start_byte_idx;
    uint16_t bits_prepared = 0;

    for (uint8_t i = 0; i < sizeof(data_payload); i++) {
        if (bits_prepared >= bits_to_send) {
            break;
        }

        uint8_t data_byte = 0;
        if (start_bit_offset == 0) {
            if (current_byte_idx < sizeof(framebuffer)) {
                data_byte = framebuffer[current_byte_idx];
            }
        } else {
            if (current_byte_idx < sizeof(framebuffer)) {
                data_byte = (framebuffer[current_byte_idx] >> start_bit_offset);
            }
            if (current_byte_idx + 1 < sizeof(framebuffer)) {
                data_byte |= (framebuffer[current_byte_idx + 1] << (8 - start_bit_offset));
            }
        }
        
        uint16_t bits_left_to_prepare = bits_to_send - bits_prepared;
        if (bits_left_to_prepare < 8) {
            data_byte &= (1 << bits_left_to_prepare) - 1;
        }

        data_payload[i] = data_byte;
        bits_prepared += 8;
        current_byte_idx++;
    }

    // Write data payload bytes 1-6
    for (uint8_t i = 0; i < 6; i++) {
        spi_write(data_payload[i]);
    }

    uint8_t byte7 = data_payload[6];
    uint8_t byte8 = 0;

    // Prepare bytes 7 and 8 which contain last data nibble and control bits
    switch (dir_bits) {
        case 0: // dir=00
            byte7 |= ((control & CTRL_CU) ? 0x40 : 0) | ((control & CTRL_P0) ? 0x80 : 0);
            byte8 = ((control & CTRL_P1) ? 0x01 : 0) | ((control & CTRL_P2) ? 0x02 : 0) |
                    ((control & CTRL_P3) ? 0x04 : 0) | ((control & CTRL_DR) ? 0x08 : 0) |
                    ((control & CTRL_SC) ? 0x10 : 0) | ((control & CTRL_BU) ? 0x20 : 0);
            break;
        case 1: // dir=01
            byte8 = 0x40;
            break;
        case 2: // dir=10
            byte8 = 0x80;
            break;
        case 3: // dir=11
            byte8 = 0xC0;
            break;
    }
    spi_write(byte7);
    spi_write(byte8);

    CE_LOW();
}

// Update the display by sending framebuffer data via SPI
// This function calls the pt6524_send_block helper for each required transaction.
// With compiler optimizations, this will be expanded into a series of efficient,
// hardcoded SPI writes specific to the chosen PT6524_MAX_ADDRESS.
void pt6524_update_display(void) {
#if PT6524_MAX_ADDRESS > (DATA_BITS_PER_BLOCK * 0)
    pt6524_send_block(0, (DATA_BITS_PER_BLOCK * 0));
#endif
#if PT6524_MAX_ADDRESS > (DATA_BITS_PER_BLOCK * 1)
    pt6524_send_block(1, (DATA_BITS_PER_BLOCK * 1));
#endif
#if PT6524_MAX_ADDRESS > (DATA_BITS_PER_BLOCK * 2)
    pt6524_send_block(2, (DATA_BITS_PER_BLOCK * 2));
#endif
#if PT6524_MAX_ADDRESS > (DATA_BITS_PER_BLOCK * 3)
    pt6524_send_block(3, (DATA_BITS_PER_BLOCK * 3));
#endif
}