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

// ==========================================
// PIN CONFIGURATION (CH32V003J4M6 - SOP8)
// ==========================================
#define FLOW_PIN      6 // PD6 = bit 6 of GPIOD
#define ALARM_PIN     2 // PA2 = bit 2 of GPIOA
#define BUTTON_PIN    4 // PC4 = bit 4 of GPIOC

// Constants
#define FLOW_CONVERSION_FACTOR 5.5
#define DEBOUNCE_DELAY_MS      200

// Fixed ADC limits for the resistor ladder
#define ADC_UP_BTN_MAX      15   // UP connects direct to GND (~0V -> ADC ~0)
#define ADC_UP_BTN_MIN      0
#define ADC_DOWN_BTN_MAX    200  // DOWN connects via 1k divider (~0.45V -> ADC ~93)
#define ADC_DOWN_BTN_MIN    50

// Globals
volatile uint32_t pulse_count = 0;
uint32_t flow_rate_ml_min = 0;   
uint32_t threshold_ml_min = 2000; 
uint8_t  alarm_active = 0;       

// ==========================================
// COMPACT 5x8 FONT (ASCII 32 to 126)
// ==========================================
const uint8_t font5x8[] = {
    0x00,0x00,0x00,0x00,0x00, 0x00,0x00,0x4f,0x00,0x00, 0x00,0x07,0x00,0x07,0x00, 0x14,0x7f,0x14,0x7f,0x14,
    0x24,0x2a,0x7f,0x2a,0x12, 0x23,0x13,0x08,0x64,0x62, 0x36,0x49,0x55,0x22,0x50, 0x00,0x05,0x03,0x00,0x00,
    0x00,0x1c,0x22,0x41,0x00, 0x00,0x41,0x22,0x1c,0x00, 0x08,0x2a,0x1c,0x2a,0x08, 0x08,0x08,0x3e,0x08,0x08,
    0x00,0x50,0x30,0x00,0x00, 0x08,0x08,0x08,0x08,0x08, 0x00,0x60,0x60,0x00,0x00, 0x20,0x10,0x08,0x04,0x02,
    0x3e,0x51,0x49,0x45,0x3e, 0x00,0x42,0x7f,0x40,0x00, 0x42,0x61,0x51,0x49,0x46, 0x21,0x41,0x45,0x4b,0x31,
    0x18,0x14,0x12,0x7f,0x10, 0x27,0x45,0x45,0x45,0x39, 0x3c,0x4a,0x49,0x49,0x30, 0x01,0x71,0x09,0x05,0x03,
    0x36,0x49,0x49,0x49,0x36, 0x06,0x49,0x49,0x29,0x1e, 0x00,0x36,0x36,0x00,0x00, 0x00,0x56,0x36,0x00,0x00,
    0x00,0x08,0x14,0x22,0x41, 0x14,0x14,0x14,0x14,0x14, 0x41,0x22,0x14,0x08,0x00, 0x02,0x01,0x51,0x09,0x06,
    0x32,0x49,0x79,0x41,0x3e, 0x7e,0x11,0x11,0x11,0x7e, 0x7f,0x49,0x49,0x49,0x36, 0x3e,0x41,0x41,0x41,0x22,
    0x7f,0x41,0x41,0x22,0x1c, 0x7f,0x49,0x49,0x49,0x41, 0x7f,0x09,0x09,0x01,0x01, 0x3e,0x41,0x41,0x51,0x32,
    0x7f,0x08,0x08,0x08,0x7f, 0x00,0x41,0x7f,0x41,0x00, 0x20,0x40,0x41,0x3f,0x01, 0x7f,0x08,0x14,0x22,0x41,
    0x7f,0x40,0x40,0x40,0x40, 0x7f,0x02,0x04,0x02,0x7f, 0x7f,0x04,0x08,0x10,0x7f, 0x3e,0x41,0x41,0x41,0x3e,
    0x7f,0x09,0x09,0x09,0x06, 0x3e,0x41,0x51,0x21,0x5e, 0x7f,0x09,0x19,0x29,0x46, 0x46,0x49,0x49,0x49,0x31,
    0x01,0x01,0x7f,0x01,0x01, 0x3f,0x40,0x40,0x40,0x3f, 0x1f,0x20,0x40,0x20,0x1f, 0x3f,0x40,0x38,0x40,0x3f,
    0x63,0x14,0x08,0x14,0x63, 0x03,0x04,0x78,0x04,0x03, 0x61,0x51,0x49,0x45,0x43, 0x00,0x7f,0x41,0x41,0x00,
    0x02,0x04,0x08,0x10,0x20, 0x00,0x41,0x41,0x7f,0x00, 0x04,0x02,0x01,0x02,0x04, 0x40,0x40,0x40,0x40,0x40,
    0x00,0x01,0x02,0x04,0x00, 0x20,0x54,0x54,0x54,0x78, 0x7f,0x48,0x44,0x44,0x38, 0x38,0x44,0x44,0x44,0x20,
    0x38,0x44,0x44,0x48,0x7f, 0x38,0x54,0x54,0x54,0x18, 0x08,0x7e,0x09,0x01,0x02, 0x08,0x14,0x54,0x54,0x3c,
    0x7f,0x08,0x04,0x04,0x78, 0x00,0x44,0x7d,0x40,0x00, 0x20,0x40,0x44,0x3d,0x00, 0x7f,0x10,0x28,0x44,0x00,
    0x00,0x41,0x7f,0x40,0x00, 0x7c,0x04,0x18,0x04,0x78, 0x7c,0x08,0x04,0x04,0x78, 0x38,0x44,0x44,0x44,0x38,
    0x7c,0x14,0x14,0x14,0x08, 0x08,0x14,0x14,0x18,0x7c, 0x7c,0x08,0x04,0x04,0x08, 0x48,0x54,0x54,0x54,0x20,
    0x04,0x3f,0x44,0x40,0x20, 0x3c,0x40,0x40,0x20,0x7c, 0x1c,0x20,0x40,0x20,0x1c, 0x3c,0x40,0x30,0x40,0x3c,
    0x44,0x28,0x10,0x28,0x44, 0x0c,0x50,0x50,0x50,0x3c, 0x44,0x64,0x54,0x4c,0x44, 0x00,0x08,0x36,0x41,0x00,
    0x00,0x00,0x7f,0x00,0x00, 0x00,0x41,0x36,0x08,0x00, 0x08,0x08,0x2a,0x1c,0x08, 0x08,0x1c,0x2a,0x08,0x08
};

// ==========================================
// I2C & SSD1306 DRIVER
// ==========================================
#define OLED_ADDR 0x3C

// Resilient I2C Write with Timeouts (prevents MCU freezing if OLED is disconnected)
void I2C_Write(uint8_t addr, uint8_t reg, uint8_t data) {
    uint16_t timeout;

    timeout = 10000; while((I2C1->STAR2 & I2C_STAR2_BUSY) && --timeout);
    if (!timeout) return;

    I2C1->CTLR1 |= I2C_CTLR1_START;
    timeout = 10000; while(!(I2C1->STAR1 & I2C_STAR1_SB) && --timeout);
    
    I2C1->DATAR = addr << 1;
    timeout = 10000; while(!(I2C1->STAR1 & I2C_STAR1_ADDR) && --timeout);
    (void)I2C1->STAR2; // Read STAR2 to clear ADDR flag
    
    I2C1->DATAR = reg;
    timeout = 10000; while(!(I2C1->STAR1 & I2C_STAR1_TXE) && --timeout);
    
    I2C1->DATAR = data;
    timeout = 10000; while(!(I2C1->STAR1 & I2C_STAR1_BTF) && --timeout);
    
    I2C1->CTLR1 |= I2C_CTLR1_STOP;
}

void OLED_Cmd(uint8_t cmd) { I2C_Write(OLED_ADDR, 0x00, cmd); }
void OLED_Data(uint8_t data) { I2C_Write(OLED_ADDR, 0x40, data); }

void OLED_Init(void) {
    const uint8_t init_sequence[] = {
        0xAE, 0xD5, 0x80, 0xA8, 0x3F, 0xD3, 0x00, 0x40,
        0x8D, 0x14, 0x20, 0x02, 0xA1, 0xC8, 0xDA, 0x12,
        0x81, 0xCF, 0xD9, 0xF1, 0xDB, 0x40, 0xA4, 0xA6, 0xAF
    };
    for(int i=0; i<sizeof(init_sequence); i++) OLED_Cmd(init_sequence[i]);
}

void OLED_SetPos(uint8_t x, uint8_t y) {
    OLED_Cmd(0xB0 + y);
    OLED_Cmd(((x & 0xF0) >> 4) | 0x10);
    OLED_Cmd(x & 0x0F);
}

void OLED_Clear(void) {
    for(uint8_t y=0; y<8; y++) {
        OLED_SetPos(0, y);
        for(uint8_t x=0; x<128; x++) OLED_Data(0x00);
    }
}

// Print string using 5x8 font (scaled to 6x8 with spacing)
void OLED_Print(uint8_t x, uint8_t y, const char* str) {
    OLED_SetPos(x, y);
    while(*str) {
        if (*str >= 32 && *str <= 126) {
            uint16_t idx = (*str - 32) * 5;
            for(uint8_t i=0; i<5; i++) OLED_Data(font5x8[idx + i]);
            OLED_Data(0x00); // 1px spacing
        }
        str++;
    }
}

// ==========================================
// INTERRUPT HANDLER (Flow Sensor Pulses)
// ==========================================
void EXTI7_0_IRQHandler(void) __attribute__((interrupt));
void EXTI7_0_IRQHandler(void) {
    if (EXTI->INTFR & (1 << FLOW_PIN)) {
        pulse_count++;
        EXTI->INTFR = (1 << FLOW_PIN); 
    }
}

// ==========================================
// HARDWARE INITIALIZATION
// ==========================================
void Init_Hardware(void) {
    // 1. Enable Clocks
    RCC->APB2PCENR |= RCC_APB2Periph_GPIOA | RCC_APB2Periph_GPIOC | RCC_APB2Periph_GPIOD | RCC_APB2Periph_ADC1 | RCC_APB2Periph_AFIO;
    RCC->APB1PCENR |= RCC_APB1Periph_I2C1;

    // 2. Configure I2C Pins (PC1 SDA, PC2 SCL) -> AF Open Drain, 10MHz
    GPIOC->CFGLR &= ~(0xFF << 4);
    GPIOC->CFGLR |= (0xDD << 4); // 0xD = Alternate Function Open Drain (11) + 10Mhz (01)

    // 3. Init Hardware I2C (100kHz @ 48MHz SysClk)
    I2C1->CTLR1 |= I2C_CTLR1_SWRST;
    I2C1->CTLR1 &= ~I2C_CTLR1_SWRST;
    I2C1->CTLR2 = 48;          
    I2C1->CKCFGR = 240;        
    I2C1->CTLR1 |= I2C_CTLR1_PE;

    // 4. ADC Setup (Fixed Clocks & SW Trigger Setup)
    RCC->CFGR0 &= ~RCC_ADCPRE;        
    RCC->CFGR0 |= RCC_ADCPRE_DIV8;    // CH32V003 ADC max is 14MHz, 48/8=6MHz is safe

    GPIOC->CFGLR &= ~(0xF << (4 * BUTTON_PIN)); // Analog mode (no pull)
    ADC1->CTLR2 |= ADC_ADON;       
    ADC1->RSQR3 = 2;               // AIN2
    ADC1->SAMPTR2 = 7 << (3 * 2);  

    ADC1->CTLR2 |= ADC_EXTSEL;     // [FIX] Select SWSTART trigger (111)
    ADC1->CTLR2 |= ADC_EXTTRIG;    // [FIX] Enable external trigger for SW to work

    ADC1->CTLR2 |= ADC_RSTCAL;     
    while(ADC1->CTLR2 & ADC_RSTCAL);
    ADC1->CTLR2 |= ADC_CAL;
    while(ADC1->CTLR2 & ADC_CAL);

    // 5. Output / Input Pins
    GPIOA->CFGLR &= ~(0xF << (4 * ALARM_PIN));
    GPIOA->CFGLR |= (GPIO_Speed_10MHz | GPIO_CNF_OUT_PP) << (4 * ALARM_PIN);
    GPIOA->BSHR = (1 << (16 + ALARM_PIN)); 

    GPIOD->CFGLR &= ~(0xF << (4 * FLOW_PIN));
    GPIOD->CFGLR |= (GPIO_CNF_IN_PUPD) << (4 * FLOW_PIN);
    GPIOD->BSHR = (1 << FLOW_PIN); 

    // 6. EXTI Interrupt
    AFIO->EXTICR |= (0b11 << (FLOW_PIN * 2));   
    EXTI->INTENR |= (1 << FLOW_PIN);             
    EXTI->FTENR  |= (1 << FLOW_PIN);             
    NVIC_EnableIRQ(EXTI7_0_IRQn);                
}

// ==========================================
// HELPER FUNCTIONS
// ==========================================
uint16_t Read_ADC(void) {
    ADC1->CTLR2 |= ADC_SWSTART;
    while(!(ADC1->STATR & ADC_EOC));
    return ADC1->RDATAR;
}

void Update_Display(uint32_t current_flow, uint32_t threshold) {
    char buf[32];
    
    // Page 0 (Top line)
    OLED_Print(16, 0, "CNC WATER FLOW");
    
    // Page 2 (Current Flow)
    snprintf(buf, sizeof(buf), "Flow: %4lu mL/m", current_flow);
    OLED_Print(0, 3, buf);
    
    // Page 4 (Threshold)
    snprintf(buf, sizeof(buf), "Lim : %4lu mL/m", threshold);
    OLED_Print(0, 5, buf);
    
    // Page 6 (Status)
    if (current_flow >= threshold) {
        OLED_Print(0, 7, "Status: SYSTEM OK ");
    } else {
        OLED_Print(0, 7, "Status: ALARM !!! ");
    }
}

// ==========================================
// MAIN LOOP
// ==========================================
int main() {
    SystemInit();
    Init_Hardware();
    
    OLED_Init();
    OLED_Clear();
    Update_Display(flow_rate_ml_min, threshold_ml_min);

    printf("\r\n==================================\r\n");
    printf(" CNC Flow Controller Initialized\r\n");
    printf("==================================\r\n");
    printf("[INFO] Active Threshold: %lu mL/min\r\n", threshold_ml_min);

    uint32_t last_time = SysTick->CNT;
    uint32_t last_button_time = 0;

    while(1) {
        uint32_t current_time = SysTick->CNT;
        uint32_t ticks_per_sec = DELAY_US_TIME * 1000000; 

        // --- 1. EVALUATE FLOW EVERY 1 SECOND ---
        if ((current_time - last_time) >= ticks_per_sec) {
            last_time = current_time;

            NVIC_DisableIRQ(EXTI7_0_IRQn);
            uint32_t hz = pulse_count;
            pulse_count = 0;
            NVIC_EnableIRQ(EXTI7_0_IRQn);

            flow_rate_ml_min = (hz * 10000) / 55;

            // Failsafe Logic
            if (flow_rate_ml_min >= threshold_ml_min) {
                GPIOA->BSHR = (1 << ALARM_PIN); // Alarm OK (MOSFET ON)
                if (alarm_active == 1) {
                    alarm_active = 0;
                    printf("[INFO] Flow restored: %lu mL/min. Alarm CLEARED.\r\n", flow_rate_ml_min);
                }            } else {
                GPIOA->BSHR = (1 << (16 + ALARM_PIN)); // Alarm Triggered (MOSFET OFF)
                if (alarm_active == 0) {
                    alarm_active = 1;
                    printf("[WARN] Low flow detected: %lu mL/min (Thr: %lu). Alarm TRIGGERED!\r\n", flow_rate_ml_min, threshold_ml_min);
                }            }

            Update_Display(flow_rate_ml_min, threshold_ml_min);
        }

        // --- 2. BUTTON HANDLING ---
        uint16_t adc_val = Read_ADC();

        if ((current_time - last_button_time) > (DEBOUNCE_DELAY_MS * DELAY_US_TIME * 1000)) {
            if (adc_val >= ADC_UP_BTN_MIN && adc_val <= ADC_UP_BTN_MAX) {
                threshold_ml_min += 100;
                last_button_time = current_time;
                printf("[SET] Threshold increased to: %lu mL/min\r\n", threshold_ml_min);
                Update_Display(flow_rate_ml_min, threshold_ml_min);
            }
            else if (adc_val >= ADC_DOWN_BTN_MIN && adc_val <= ADC_DOWN_BTN_MAX) {
                if (threshold_ml_min >= 100) threshold_ml_min -= 100;
                last_button_time = current_time;
                printf("[SET] Threshold decreased to: %lu mL/min\r\n", threshold_ml_min);
                Update_Display(flow_rate_ml_min, threshold_ml_min);
            }
        }
    }
}