#include "ch32fun.h"
#include <stdint.h>

#ifndef FUN_HIGH
#define FUN_HIGH 1
#define FUN_LOW  0
#endif


/* =====================================================================
 PIN DEFINITIONS (CH32V003J4M6 SOP-8)
 
     +-----U----+      
  A1 | o        | SWIO  
 GND |          | C4    
  A2 | ch32V003 | C2    
 VDD |          | C1    
     +----------+      
     
  SSD1306 display
  
  PS2 mouse
     
===================================================================== */

#define PS2_CLK_PIN     PA1
#define PS2_DAT_PIN     PA2

#define I2C_SDA_PIN     PC1
#define I2C_SCL_PIN     PC2

// 128 x 32 oled display
// Most SSD1306 modules are 0x78. Some are 0x7A.
#define SSD1306_ADDR    0x78

#define PS2_TIMEOUT_US  20000

// ---------------------------------------------------------------
// SSD1306 framebuffer: 128x32 bits = 512 bytes
// ---------------------------------------------------------------

static uint8_t fb[512];

// ---------------------------------------------------------------
// PS/2 low-level helpers
// ---------------------------------------------------------------

// These assume the GPIOs are configured as open-drain outputs.
// Writing high should release the line; external pull-ups pull it high.
//
// If your CH32Fun/GPIO mode does not behave as true open-drain,
// replace the *_hi() functions with input-floating mode and the
// *_lo() functions with output-low mode.

static inline void ps2_clk_lo(void) {  funDigitalWrite(PS2_CLK_PIN, FUN_LOW); }
static inline void ps2_clk_hi(void) {  funDigitalWrite(PS2_CLK_PIN, FUN_HIGH); }
static inline void ps2_dat_lo(void) {  funDigitalWrite(PS2_DAT_PIN, FUN_LOW);  }
static inline void ps2_dat_hi(void) {  funDigitalWrite(PS2_DAT_PIN, FUN_HIGH); }

static inline uint8_t ps2_clk_rd(void) {  return !!funDigitalRead(PS2_CLK_PIN); }
static inline uint8_t ps2_dat_rd(void) {  return !!funDigitalRead(PS2_DAT_PIN); }

static int ps2_wait_clk(uint8_t level, uint32_t timeout_us) {
    while (ps2_clk_rd() != level) {
        if (timeout_us == 0) return -1;
        timeout_us--;
        Delay_Us(1);
    }
    return 0;
}

static int ps2_wait_dat(uint8_t level, uint32_t timeout_us) {
    while (ps2_dat_rd() != level)  {
        if (timeout_us == 0) return -1;
        timeout_us--;
        Delay_Us(1);
    }
    return 0;
}

// ---------------------------------------------------------------
// PS/2 host-to-device send
//
// 11 bits:
//   start bit = 0
//   8 data bits, LSB first
//   odd parity
//   stop bit = 1
// ---------------------------------------------------------------

static int ps2_send_byte(uint8_t byte) {
    uint8_t ones = 0;

    for (uint8_t i = 0; i < 8; i++)    {
        if (byte & (1u << i)) ones++;
    }

    // Odd parity: total number of 1s in data + parity must be odd.
    uint8_t parity = (ones & 1) ? 0 : 1;

    // Inhibit clock.
    ps2_clk_lo();
    Delay_Us(200);

    // Start bit.
    ps2_dat_lo();

    // Release clock; device takes over clock generation.
    ps2_clk_hi();

    // Wait for device to begin clocking.
    if (ps2_wait_clk(1, PS2_TIMEOUT_US)) return -1;
    if (ps2_wait_clk(0, PS2_TIMEOUT_US)) return -2;

    // 8 data bits.
    for (uint8_t i = 0; i < 8; i++) {
        if (byte & (1u << i)) ps2_dat_hi();
        else                   ps2_dat_lo();

        if (ps2_wait_clk(1, PS2_TIMEOUT_US)) return -3;
        if (ps2_wait_clk(0, PS2_TIMEOUT_US)) return -4;
    }

    // Parity bit.
    if (parity) ps2_dat_hi();
    else        ps2_dat_lo();

    if (ps2_wait_clk(1, PS2_TIMEOUT_US)) return -5;
    if (ps2_wait_clk(0, PS2_TIMEOUT_US)) return -6;

    // Stop bit.
    ps2_dat_hi();

    if (ps2_wait_clk(1, PS2_TIMEOUT_US)) return -7;
    if (ps2_wait_clk(0, PS2_TIMEOUT_US)) return -8;

    // ACK from device: device pulls Data low, then releases it.
    if (ps2_wait_dat(0, PS2_TIMEOUT_US)) return -9;
    if (ps2_wait_dat(1, PS2_TIMEOUT_US)) return -10;

    return 0;
}

// ---------------------------------------------------------------
// PS/2 device-to-host receive
//
// Waits for a full 11-bit frame from the mouse.
// ---------------------------------------------------------------

static int ps2_read_byte(uint8_t *out, uint32_t timeout_us) {
    // Wait for idle clock high, then the falling edge that starts the
    // frame. The start bit is valid on this edge (we don't bother
    // checking its value, we just use the edge to sync up).
    if (ps2_wait_clk(1, timeout_us)) return -1;
    if (ps2_wait_clk(0, timeout_us)) return -2;

    // Move past the start bit's low pulse so we're aligned for the
    // next falling edge (which will carry data bit 0).
    if (ps2_wait_clk(1, timeout_us)) return -3;

    uint8_t byte = 0;
    uint8_t ones = 0;

    // 8 data bits, LSB first. Data is only guaranteed valid on the
    // falling edge of clock, so sample immediately after clock goes
    // low, then wait for the rising edge before looking for the next
    // falling edge.
    for (uint8_t i = 0; i < 8; i++)   {
        if (ps2_wait_clk(0, timeout_us)) return -4;

        uint8_t bit = ps2_dat_rd();
        if (bit) {
            byte |= (1u << i);
            ones++;
        }

        if (ps2_wait_clk(1, timeout_us)) return -5;
    }

    // Parity bit.
    if (ps2_wait_clk(0, timeout_us)) return -6;
    uint8_t parity = ps2_dat_rd();
    if (ps2_wait_clk(1, timeout_us)) return -7;

    // Stop bit.
    if (ps2_wait_clk(0, timeout_us)) return -8;
    uint8_t stop = ps2_dat_rd();

    if (!stop) return -9;

    // Odd parity check: data ones + parity bit should be odd.
    if (((ones + parity) & 1) == 0) return -10;

    *out = byte;
    return 0;
}

static void ps2_flush(void){
    uint8_t dummy;
    while (ps2_read_byte(&dummy, 20000) == 0)    {
        // Discard pending bytes.
    }
}

// ---------------------------------------------------------------
// Basic PS/2 mouse initialization
//
// 0xFF = Reset
// 0xF6 = Set Defaults
// 0xF4 = Enable Data Reporting
//
// Many mice will work without all of this, but reset/enable is
// a reasonably common sequence.
// ---------------------------------------------------------------

static void ps2_mouse_init(void) {
    uint8_t r;

    Delay_Ms(100);

    ps2_flush();

    // Reset. Ignore / tolerate errors during bring-up.
    ps2_send_byte(0xFF);

    // Usually ACK 0xFA.
    ps2_read_byte(&r, 300000);

    // Usually BAT completion 0xAA.
    ps2_read_byte(&r, 100000);

    // Set Defaults.
    ps2_send_byte(0xF6);
    ps2_read_byte(&r, 100000);

    // Enable streaming/reporting.
    ps2_send_byte(0xF4);
    ps2_read_byte(&r, 100000);
}

// ---------------------------------------------------------------
// Bit-banged I2C master
// ---------------------------------------------------------------

static inline void i2c_delay(void) {   
    // Delay_Us(0); 
   asm volatile ("nop"::);
   asm volatile ("nop"::);
    asm volatile ("nop"::);
   asm volatile ("nop"::);
   
}

static void i2c_init(void) {
    // Use open-drain mode. External or module pull-ups required.
    funPinMode(I2C_SCL_PIN, GPIO_CFGLR_OUT_10Mhz_OD);
    funPinMode(I2C_SDA_PIN, GPIO_CFGLR_OUT_10Mhz_OD);

    funDigitalWrite(I2C_SCL_PIN, FUN_HIGH);
    funDigitalWrite(I2C_SDA_PIN, FUN_HIGH);
}



static inline void i2c_scl_hi(void) {
    funDigitalWrite(I2C_SCL_PIN, FUN_HIGH);
    i2c_delay();
}

static inline void i2c_scl_lo(void) {
    funDigitalWrite(I2C_SCL_PIN, FUN_LOW);
    i2c_delay();
}

static inline void i2c_sda_hi(void) {
    funDigitalWrite(I2C_SDA_PIN, FUN_HIGH);
}

static inline void i2c_sda_lo(void){
    funDigitalWrite(I2C_SDA_PIN, FUN_LOW);
}



static void i2c_start(void) {
    i2c_sda_hi();
    i2c_scl_hi();
    i2c_sda_lo();
    i2c_scl_lo();
}

static void i2c_stop(void) {
    i2c_sda_lo();
    i2c_scl_hi();
    i2c_sda_hi();
}

static void i2c_write(uint8_t b) {
    for (int8_t i = 7; i >= 0; i--)   {
        if (b & (1u << i)) i2c_sda_hi();
        else               i2c_sda_lo();

        i2c_scl_hi();
        i2c_scl_lo();
    }

    // ACK clock pulse. We do not read the ACK here.
    i2c_sda_hi();
    i2c_scl_hi();
    i2c_scl_lo();
}

// ---------------------------------------------------------------
// SSD1306 commands
// ---------------------------------------------------------------

static void ssd1306_cmd(uint8_t cmd) {
    i2c_start();
    i2c_write(SSD1306_ADDR);
    i2c_write(0x00); // Co = 0, D/C# = 0: command
    i2c_write(cmd);
    i2c_stop();
}

static void ssd1306_init(void) {
    Delay_Ms(100);

    ssd1306_cmd(0xAE); ssd1306_cmd(0xD5); ssd1306_cmd(0x80); 
    ssd1306_cmd(0xA8); ssd1306_cmd(0x1F);
    ssd1306_cmd(0xD3); ssd1306_cmd(0x00); ssd1306_cmd(0x40); ssd1306_cmd(0x8D); ssd1306_cmd(0x14);
    ssd1306_cmd(0x20); ssd1306_cmd(0x00); ssd1306_cmd(0xA1); ssd1306_cmd(0xC8); 
    ssd1306_cmd(0xDA); ssd1306_cmd(0x02); 
    ssd1306_cmd(0x81); ssd1306_cmd(0xCF); ssd1306_cmd(0xD9); ssd1306_cmd(0xF1);
    ssd1306_cmd(0xDB); ssd1306_cmd(0x40); ssd1306_cmd(0xA4); ssd1306_cmd(0xA6); ssd1306_cmd(0xAF);
}

static void ssd1306_set_addr(void) {
    // Column 0..127
    ssd1306_cmd(0x21);
    ssd1306_cmd(0x00);
    ssd1306_cmd(0x7F);

    // Page 0..7
    ssd1306_cmd(0x22);
    ssd1306_cmd(0x00);
    ssd1306_cmd(0x07);
}

static void ssd1306_send_fb(void) {
    ssd1306_cmd(0x21); ssd1306_cmd(0); ssd1306_cmd(127);
    ssd1306_cmd(0x22); ssd1306_cmd(0); ssd1306_cmd(7);

    i2c_start(); i2c_write(SSD1306_ADDR);  i2c_write(0x40); // Co = 0, D/C# = 1: data

    for (uint16_t i = 0; i < 512; i++)     {
        i2c_write(fb[i]);
    }

    i2c_stop();
}

// ---------------------------------------------------------------
// Simple framebuffer drawing
// ---------------------------------------------------------------

static void fb_clear(void) {
    for (uint16_t i = 0; i < 512; i++)   {
        fb[i] = 0;
    }
}

static void fb_pixel(int x, int y, uint8_t on) {
    if (x < 0 || x >= 128 || y < 0 || y >= 32) return;

    uint16_t idx = (uint16_t)((y >> 3) * 128 + x);
    uint8_t mask = (uint8_t)(1u << (y & 7));

    if (on) fb[idx] |= mask;
    else    fb[idx] &= (uint8_t)~mask;
}

static void draw_box(int x, int y, int w, int h, uint8_t filled) {
    for (int yy = y; yy < y + h; yy++)     {
        for (int xx = x; xx < x + w; xx++)       {
            if (filled ||
                yy == y ||
                yy == y + h - 1 ||
                xx == x ||
                xx == x + w - 1)    {
                fb_pixel(xx, yy, 1);
            }
        }
    }
}



static void screen_update(const uint8_t *pkt, int x, int y) {
    fb_clear();

    // Button indicators:
    // Left, Middle, Right.
    draw_box(2,   2, 18, 8, pkt[0] & 0x01);
    draw_box(55,  2, 18, 8, pkt[0] & 0x04);
    draw_box(108, 2, 18, 8, pkt[0] & 0x02);

    // Cursor.
    draw_box(x - 2, y - 2, 5, 5, 1);

    ssd1306_set_addr();
    ssd1306_send_fb();
}

// ---------------------------------------------------------------
// Main
// ---------------------------------------------------------------

int main(void) {
    // If your ch32fun version uses SystemInit() instead, adjust here.
    SystemInit();

    funGpioInitAll();

    // PS/2 pins: open-drain, released high initially.
    funPinMode(PS2_CLK_PIN, GPIO_CFGLR_OUT_10Mhz_OD);
    funPinMode(PS2_DAT_PIN, GPIO_CFGLR_OUT_10Mhz_OD);

    ps2_clk_hi();
    ps2_dat_hi();

    i2c_init();
    ssd1306_init();

    fb_clear();
    ssd1306_set_addr();
    ssd1306_send_fb();

    ps2_mouse_init();

    int x = 64;
    int y = 16;

    uint8_t pkt[3] = {0, 0, 0};

    screen_update(pkt, x, y);


    while (1)    {
        uint8_t b;

        // Wait for first byte of a mouse packet.
        // For standard PS/2 mouse packets, byte 0 has bit 3 set.
        if (ps2_read_byte(&b, 300000)) continue;

        if ((b & 0x08) == 0)   {
            // Not aligned to packet start. Discard and resync.
            continue;
        }

        pkt[0] = b;

        if (ps2_read_byte(&pkt[1], 30000)) continue;
        if (ps2_read_byte(&pkt[2], 30000)) continue;

        // Decode X/Y movement.
        int dx = pkt[1];
        int dy = pkt[2];

        if (pkt[0] & 0x10) dx -= 256; // X sign bit
        if (pkt[0] & 0x20) dy -= 256; // Y sign bit
        
        // pkt[0] bits 0,1,2 are the mouse button status.

        // PS/2 positive Y is usually upward; OLED Y increases downward.
        x += dx;
        y -= dy;

        // Keep cursor on screen.
        if (x < 2)   x = 2;
        if (x > 125) x = 125;

        if (y < 2)  y = 2;
        if (y > 29)  y = 29;

        screen_update(pkt, x, y);
    }

    return 0;
}