#include <stdio.h>
#include <stdint.h>
#include <stdbool.h>
#include <stdlib.h>
#include "pico/stdlib.h"
#include "hardware/i2c.h"
#include "hardware/flash.h"
#include "ssd1306.h"

#define FLASH_TARGET_OFFSET (PICO_FLASH_SIZE_BYTES - 4096)
#define DIAMETER_STORAGE_ADDR (XIP_BASE + FLASH_TARGET_OFFSET)

const uint ENABLE_PIN = 10;
const uint STEP_PIN = 16;
const uint DIR_PIN = 17;
const uint MS1 = 13;
const uint MS2 = 14;
const uint MS3 = 15;

const uint LEFT = 18;
const uint DOWN = 19;
const uint RIGHT = 20;
const uint UP = 21;

const uint SDA_PIN = 4;
const uint SCL_PIN = 5;
const uint I2C_ADDR = 0x3C;
ssd1306_t disp;

float DIAMETER = 20.0;

void save_diameter_to_flash(float diameter) {
    uint32_t stored_value = *((uint32_t*)&diameter);
    uint32_t page_buffer[FLASH_PAGE_SIZE / sizeof(uint32_t)];
    page_buffer[0] = stored_value;
    
    flash_range_erase(FLASH_TARGET_OFFSET, FLASH_SECTOR_SIZE);
    flash_range_program(FLASH_TARGET_OFFSET, (const uint8_t*)page_buffer, FLASH_PAGE_SIZE);
}

float load_diameter_from_flash() {
    uint32_t* flash_ptr = (uint32_t*)(DIAMETER_STORAGE_ADDR);
    uint32_t stored_value = *flash_ptr;
    
    if (stored_value == 0xFFFFFFFF) {
        return 20.0;
    }
    
    return *((float*)&stored_value);
}

const float LEAD = 8.0;
const float STEPS_PER_REV = 200 * 8;
const float STEPS_PER_MM = STEPS_PER_REV / LEAD;

const uint max_visible = 3;
int selection = 0;
int total_items = 5;
int display_top = 0;
bool update = true;

//functions
void init();
void update_menu();
float infuse(float current_val);
float set_diameter(float current_val);
float manual_move(float current_val);
float snake_game(float current_val);

typedef struct {
	char *label;
	int id;
	float (*action)(float);
} MenuItem;
	
MenuItem menu[] = {
	{"Infuse", 0, infuse},
	{"Stop", 1, NULL},
	{"Set Diameter", 2, set_diameter},
	{"Move", 3, manual_move},
	{"Snake", 4, snake_game},
};
	
int main() {
	
	init();
	
	while(true) {
		
		if (update) {
			update_menu();
			update = false;
		}
		if (gpio_get(UP) && selection > 0)
		{
			selection--;
			if (selection < display_top) { //scrolling up
				display_top = selection;
			}
			update = true;
			sleep_ms(200);
		}
		if (gpio_get(DOWN) && selection < total_items - 1)
		{
			selection++;
			if (selection >= display_top + max_visible) //scrolling down
			{
				display_top = selection - max_visible + 1;
			}
			update = true;
			sleep_ms(200);
		}
		if (gpio_get(RIGHT))
		{
			switch (selection) {
				case 0:
					menu[selection].action(DIAMETER);
					break;
				
				case 1:
					gpio_put(ENABLE_PIN, 1); //disable stepper
					break;
					
				case 2:
					DIAMETER = menu[selection].action(DIAMETER);
					break;
					
				case 3:
					menu[selection].action(DIAMETER);
					break;
				
				case 4:
					menu[selection].action(DIAMETER);
					break;
				
				default:
					break;
			}	
			update = true;
			sleep_ms(200);
		}
		
		sleep_ms(10);
		
	}

	return 0;
}

void init() {
	stdio_init_all();

	//init pins
	gpio_init(ENABLE_PIN);
	gpio_init(STEP_PIN);
	gpio_init(DIR_PIN);
	gpio_init(MS1);
	gpio_init(MS2);
	gpio_init(MS3);
	gpio_init(LEFT);
	gpio_init(DOWN);
	gpio_init(RIGHT);
	gpio_init(UP);
	
	//stepper driver stuff
	gpio_set_dir(ENABLE_PIN, GPIO_OUT);
	gpio_set_dir(STEP_PIN, GPIO_OUT);
	gpio_set_dir(DIR_PIN, GPIO_OUT);
	gpio_set_pulls(ENABLE_PIN, false, true); // pin, pullup, pulldown
	gpio_set_pulls(STEP_PIN, false, true);
	gpio_set_pulls(DIR_PIN, false, true);
	
	gpio_put(ENABLE_PIN, 1); //disable stepper driver
	
	gpio_set_dir(MS1, GPIO_OUT);
	gpio_set_dir(MS2, GPIO_OUT);
	gpio_set_dir(MS3, GPIO_OUT);
	gpio_set_pulls(MS1, false, true);
	gpio_set_pulls(MS2, false, true);
	gpio_set_pulls(MS3, false, true);
	
	//set microsteps ... 1/8
	gpio_put(MS1, 1);
	gpio_put(MS2, 1);
	gpio_put(MS3, 0);
	
	//switches
	gpio_set_dir(LEFT, GPIO_IN);
	gpio_set_dir(DOWN, GPIO_IN);
	gpio_set_dir(RIGHT, GPIO_IN);
	gpio_set_dir(UP, GPIO_IN);
	gpio_set_pulls(LEFT, false, true);
	gpio_set_pulls(DOWN, false, true);
	gpio_set_pulls(RIGHT, false, true);
	gpio_set_pulls(UP, false, true);
	
	//init i2c
	i2c_init(i2c0, 400000);
	gpio_set_function(SDA_PIN, GPIO_FUNC_I2C);
	gpio_set_function(SCL_PIN, GPIO_FUNC_I2C);
	gpio_set_pulls(SDA_PIN, true, false);
	gpio_set_pulls(SCL_PIN, true, false);
	
	//init display
    disp.external_vcc=false;
    ssd1306_init(&disp, 128, 64, I2C_ADDR, i2c0);
    ssd1306_clear(&disp);
    
    DIAMETER = load_diameter_from_flash();
}

float set_diameter(float current_val) {
	float temp_val = current_val;
	bool editing = true;
	sleep_ms(200);
	
	while (editing) {
		ssd1306_clear(&disp);
		ssd1306_draw_string(&disp, 0, 0, 1, "Syringe Diameter:");
		
		char buffer[10];
		sprintf(buffer, "%.1f mm", temp_val);
		ssd1306_draw_string(&disp, 0, 20, 2, buffer);
		
		ssd1306_show(&disp);
		
		if (gpio_get(UP)) 
		{
			temp_val += 0.1f;
			sleep_ms(60);
		}		
		if (gpio_get(DOWN) && temp_val > 0.0f) 
		{
			temp_val -= 0.1f;
			sleep_ms(200);
		}		
		if (gpio_get(RIGHT)) 
		{
			save_diameter_to_flash(temp_val);
			editing = false;
			sleep_ms(200);
		}	
	}
	return temp_val;
}

void update_menu (){
	ssd1306_clear(&disp);
	
	for (int i = 0; i < max_visible; i++) {
        int item_index = display_top + i;
        
        if (item_index >= total_items) break;

        int y_pos = i * 16;

        if (item_index == selection) {
            ssd1306_draw_string(&disp, 0, y_pos, 1, "> ");
            ssd1306_draw_string(&disp, 15, y_pos, 1, menu[item_index].label);
        } else {
            ssd1306_draw_string(&disp, 15, y_pos, 1, menu[item_index].label);
        }
    }	
	ssd1306_show(&disp);
}

float infuse(float current_val) {
	float radius = current_val / 2.0f;
	float area = 3.14159265359f * radius * radius;
	float ml_per_mm = area / 1000.0;
	float steps_per_ml = STEPS_PER_MM / ml_per_mm;
	float flowrate = 0.1;
	float target_volume = 0.1;
	
	
	float flowrate_tempval = 0.1;
	float target_volume_tempval = 0.1;
	bool editing = true;
	bool editing_flowrate = true;
	bool editing_volume = true;
	
	sleep_ms(300);
	while(editing) {
		while(editing_flowrate) {
			ssd1306_clear(&disp);
			ssd1306_draw_string(&disp, 0, 0, 1, "Infuse");
			ssd1306_draw_string(&disp, 0, 20, 1, "Flowrate:");
	
			ssd1306_draw_string(&disp, 0, 30, 1, "> ");
			char buffer_flowrate[32];
			sprintf(buffer_flowrate, "%.1f ml/h", flowrate_tempval);
			ssd1306_draw_string(&disp, 15, 30, 1, buffer_flowrate);
	
	
			ssd1306_draw_string(&disp, 0, 40, 1, "Target volume");
	
			char buffer_volume[32];
			sprintf(buffer_volume, "%.1f ml", target_volume_tempval);
			ssd1306_draw_string(&disp, 0, 50, 1, buffer_volume);
			ssd1306_show(&disp); //update display
	
			if (gpio_get(UP)) 
			{
				flowrate_tempval += 0.1f;
				sleep_ms(60);
			}		
			if (gpio_get(DOWN) && flowrate_tempval > 0.1f) 
			{
				flowrate_tempval -= 0.1f;
				sleep_ms(200);
			}		
			if (gpio_get(RIGHT)) 
			{
				flowrate = flowrate_tempval;
				editing_flowrate = false;
				editing_volume = true;
				sleep_ms(200);
			}
	
		}

		while(editing_volume) {
			ssd1306_clear(&disp);
			ssd1306_draw_string(&disp, 0, 0, 1, "Infuse");
			ssd1306_draw_string(&disp, 0, 20, 1, "Flowrate:");
	
			char buffer_flowrate[10];
			sprintf(buffer_flowrate, "%.1f ml/h", flowrate_tempval);
			ssd1306_draw_string(&disp, 0, 30, 1, buffer_flowrate);
	
	
			ssd1306_draw_string(&disp, 0, 40, 1, "Target volume");
	
			ssd1306_draw_string(&disp, 0, 50, 1, "> ");
			char buffer_volume[10];
			sprintf(buffer_volume, "%.1f ml", target_volume_tempval);
			ssd1306_draw_string(&disp, 15, 50, 1, buffer_volume);
			ssd1306_show(&disp); //update display
		
			if (gpio_get(UP)) 
			{
				target_volume_tempval += 0.1f;
				sleep_ms(60);
			}		
			if (gpio_get(DOWN) && target_volume_tempval > 0.1f) 
			{
				target_volume_tempval -= 0.1f;
				sleep_ms(200);
			}
			if (gpio_get(LEFT))
			{
				editing_flowrate = true;
				editing_volume = false;
				sleep_ms(200);
			}
			if (gpio_get(RIGHT)) 
			{
				target_volume = target_volume_tempval;
				editing_flowrate = false;
				editing_volume = false;
				editing = false;
				sleep_ms(200);
			}
		}
	}
	
	uint32_t total_steps = (uint32_t)(target_volume * steps_per_ml);
	float ml_per_sec = flowrate / 3600.0;
	float steps_per_sec = ml_per_sec * steps_per_ml;
	uint32_t step_interval_us = (uint32_t)(1000000.0 / steps_per_sec); // interval of the step in microseconds
	
	ssd1306_clear(&disp);
	ssd1306_draw_string(&disp, 0, 0, 2, "Pumping...");
	ssd1306_draw_string(&disp, 0, 25, 1, "Press LEFT to STOP");
	ssd1306_show(&disp);
	
	uint32_t current_steps = 0;
	uint32_t last_step_time = time_us_32();
	gpio_put(DIR_PIN, 1); //forward
	gpio_put(ENABLE_PIN, 0); //enable stepper
	
	while (current_steps < total_steps) {
		uint32_t current_time = time_us_32();
		if (current_time - last_step_time >= step_interval_us) {
			gpio_put(STEP_PIN, 1);
			sleep_us(2);
			gpio_put(STEP_PIN, 0);
			
			current_steps++;
			last_step_time = current_time;
		}
		
		if (gpio_get(LEFT)) {
			gpio_put(ENABLE_PIN, 1); //disable stepper
			break;
		}
	}
	
	gpio_put(ENABLE_PIN, 1); //disable stepper
}

float manual_move(float current_val) {	
	bool active = true;
	gpio_put(ENABLE_PIN, 0); //enable stepper
	sleep_ms(300);
	
	ssd1306_clear(&disp);
	ssd1306_draw_string(&disp, 0, 0, 1, "Manual Move");
	ssd1306_show(&disp);
	
	while (active) {
		if (gpio_get(UP))
		{
			gpio_put(DIR_PIN, 0);
			gpio_put(STEP_PIN, 1);
			sleep_ms(1);
			gpio_put(STEP_PIN, 0);
			sleep_ms(1);
		}
		if (gpio_get(DOWN))
		{
			gpio_put(DIR_PIN, 1);
			gpio_put(STEP_PIN, 1);
			sleep_ms(1);
			gpio_put(STEP_PIN, 0);
			sleep_ms(1);
		}
		if (gpio_get(LEFT) || gpio_get(RIGHT))
		{
			gpio_put(ENABLE_PIN, 1); //disable stepper
		active = false;
		}
	}
	
	gpio_put(ENABLE_PIN, 1); //disable stepper
	return 0;
}

#define SNAKE_W 64
#define SNAKE_H 32
#define CELL_SIZE 2

typedef struct {
	uint8_t x;
	uint8_t y;
} Point;

Point snake[32];
int snake_len = 3;
Point food;
uint8_t dir_x = 1;
uint8_t dir_y = 0;
uint8_t next_dir_x = 1;
uint8_t next_dir_y = 0;
int score = 0;

void spawn_food() {
	food.x = (rand() % (SNAKE_W - 2)) + 1;
	food.y = (rand() % (SNAKE_H - 2)) + 1;
	for (int i = 0; i < snake_len; i++) {
		if (snake[i].x == food.x && snake[i].y == food.y) {
			spawn_food();
			return;
		}
	}
}

void draw_snake() {
	for (int i = 0; i < snake_len; i++) {
		ssd1306_draw_pixel(&disp, snake[i].x * CELL_SIZE, snake[i].y * CELL_SIZE);
		ssd1306_draw_pixel(&disp, snake[i].x * CELL_SIZE + 1, snake[i].y * CELL_SIZE);
		ssd1306_draw_pixel(&disp, snake[i].x * CELL_SIZE, snake[i].y * CELL_SIZE + 1);
		ssd1306_draw_pixel(&disp, snake[i].x * CELL_SIZE + 1, snake[i].y * CELL_SIZE + 1);
	}
}

void draw_food() {
	ssd1306_draw_pixel(&disp, food.x * CELL_SIZE, food.y * CELL_SIZE);
	ssd1306_draw_pixel(&disp, food.x * CELL_SIZE + 1, food.y * CELL_SIZE);
	ssd1306_draw_pixel(&disp, food.x * CELL_SIZE, food.y * CELL_SIZE + 1);
	ssd1306_draw_pixel(&disp, food.x * CELL_SIZE + 1, food.y * CELL_SIZE + 1);
}

float snake_game(float current_val) {
	srand(1000);
	snake_len = 3;
	snake[0].x = SNAKE_W / 2;
	snake[0].y = SNAKE_H / 2;
	snake[1].x = snake[0].x - 1;
	snake[1].y = snake[0].y;
	snake[2].x = snake[0].x - 2;
	snake[2].y = snake[0].y;
	dir_x = 1;
	dir_y = 0;
	next_dir_x = 1;
	next_dir_y = 0;
	score = 0;
	spawn_food();
	
	bool playing = true;
	int frame = 0;
	
	sleep_ms(300);
	
	while (playing) {
		if (gpio_get(UP) && dir_y != 1) {
			next_dir_x = 0; next_dir_y = -1;
		}
		if (gpio_get(DOWN) && dir_y != -1) {
			next_dir_x = 0; next_dir_y = 1;
		}
		if (gpio_get(LEFT) && dir_x != 1) {
			next_dir_x = -1; next_dir_y = 0;
		}
		if (gpio_get(RIGHT) && dir_x != -1) {
			next_dir_x = 1; next_dir_y = 0;
		}
		
		if (frame % 5 == 0) {
			dir_x = next_dir_x;
			dir_y = next_dir_y;
			
			Point head = {snake[0].x + dir_x, snake[0].y + dir_y};
			
			if (head.x == 0 || head.x >= SNAKE_W - 1 || head.y == 0 || head.y >= SNAKE_H - 1) {
				ssd1306_clear(&disp);
				ssd1306_draw_string(&disp, 20, 20, 1, "Game Over!");
				char buf[10];
				sprintf(buf, "Score: %d", score);
				ssd1306_draw_string(&disp, 20, 35, 1, buf);
				ssd1306_show(&disp);
				sleep_ms(2000);
				playing = false;
				continue;
			}
			
			for (int i = 1; i < snake_len; i++) {
				if (head.x == snake[i].x && head.y == snake[i].y) {
					ssd1306_clear(&disp);
					ssd1306_draw_string(&disp, 20, 20, 1, "Game Over!");
					char buf[10];
					sprintf(buf, "Score: %d", score);
					ssd1306_draw_string(&disp, 20, 35, 1, buf);
					ssd1306_show(&disp);
					sleep_ms(2000);
					playing = false;
					continue;
				}
			}
			
			for (int i = snake_len; i > 0; i--) {
				snake[i] = snake[i-1];
			}
			snake[0] = head;
			
			if (head.x == food.x && head.y == food.y) {
				snake_len++;
				score++;
				spawn_food();
			}
			
			ssd1306_clear(&disp);
			ssd1306_draw_line(&disp, 0, 0, SNAKE_W * CELL_SIZE - 1, 0);
			ssd1306_draw_line(&disp, 0, 0, 0, SNAKE_H * CELL_SIZE - 1);
			ssd1306_draw_line(&disp, SNAKE_W * CELL_SIZE - 1, 0, SNAKE_W * CELL_SIZE - 1, SNAKE_H * CELL_SIZE - 1);
			ssd1306_draw_line(&disp, 0, SNAKE_H * CELL_SIZE - 1, SNAKE_W * CELL_SIZE - 1, SNAKE_H * CELL_SIZE - 1);
			draw_snake();
			draw_food();
			ssd1306_show(&disp);
		}
		
		frame++;
		sleep_ms(25);
	}
	
	return 0;
}