#include "lib/pbdrv/pb-mod.h"
#include "lib/pbdrv/pb-types.h"
#include <Arduino.h>
#include <TM1637Display.h>

#define TOTAL_LEVELS 5
#define ROWS 4
#define COLS 3
#define CLK 2
#define DIO 3
#define SOLVED_PIN 53
#define I2C_MODULE_ADDRESS 0x08 // Address of the puzzle module
#define HANDSHAKE_RECEIVED \
	{ 0x70, 0x75, 0x7a, 0x62, 0x75, 0x73 } // Magic command for the handshake
#define HANDSHAKE_SEND \
	{ 0x67, 0x61, 0x6d, 0x69, 0x6e, 0x67 } // Magic command for the handshake
#define REQUEST_STATE_CMD 0x53 // 'S' to request the game state

const int ROW_PINS[ROWS] = {7, 6, 5, 4};
const int COL_PINS[COLS] = {10, 9, 8};
const char * validButtons[TOTAL_LEVELS] = {"A2", "B1", "D3", "C2", "C3"};
const char bombCode[] = "1234";
const uint8_t SEGMENT_MAP[] = {
	0b00111111, // 0
	0b00000110, // 1
	0b01011011, // 2
	0b01001111, // 3
	0b01100110, // 4
	0b01101101, // 5
	0b01111101, // 6
	0b00000111, // 7
	0b01111111, // 8
	0b01101111, // 9
	0b01110111, // A
	0b01111100, // B
	0b00111001, // C
	0b01011110, // D
	0b01111001, // E
	0b01110001 // F
	// Add other letters if needed
};

// This array of level codes matches the codes you might display per level.
const char * levelCodes[TOTAL_LEVELS] = {"A1", "B2", "D1", "C3", "A2"};

// Puzzle state
pb_global_state_t puzzleState = PB_GS_NOINIT;

TM1637Display display(CLK, DIO);

int currentLevel = 0;

void blink_display(int num) {
	if (num == 1) {
		// Display "1111" with leading zeros shown if necessary
		display.showNumberDecEx(1111, 0b11111111, true);
	} else if (num == 0) {
		// Display "0000" with leading zeros shown if necessary
		display.showNumberDecEx(0, 0b11111111, true);
	}
	delay(500);
	display.clear();
	delay(500);
}

void display_final_code(const char * code) {
	uint8_t segs[4] = {0, 0, 0, 0};
	int numDigits = strlen(code);
	numDigits = numDigits > 4 ? 4 : numDigits;

	for (int i = 0; i < numDigits; i++) {
		segs[i] = display.encodeDigit(code[i] - '0');
	}

	display.setSegments(segs, numDigits, 0);
}

void check_button_press() {
	for (int col = 0; col < COLS; col++) {
		digitalWrite(COL_PINS[col], LOW);
		for (int row = 0; row < ROWS; row++) {
			if (digitalRead(ROW_PINS[row]) == LOW) {
				delay(50);
				if (digitalRead(ROW_PINS[row]) == LOW) {
					char keyPress[3] = {'A' + row, '1' + col, '\0'};
					Serial.print("Keypress detected: ");
					Serial.println(keyPress);
					if (strcmp(keyPress, validButtons[currentLevel]) == 0) {
						currentLevel++;
						if (currentLevel >= TOTAL_LEVELS) {
							pb_hook_mod_state_write(PB_GS_SOLVED);
							Serial.println("Puzzle solved!");
							display.showNumberDec(currentLevel + 1, true);
							digitalWrite(SOLVED_PIN, HIGH);
						}
					} else {
						currentLevel = 0;
					}
					while (digitalRead(ROW_PINS[row]) == LOW) {
					} // Ensure button release
				}
			}
		}
		digitalWrite(COL_PINS[col], HIGH);
	}
}

void initialize_system() {
	for (int i = 0; i < ROWS; i++) {
		pinMode(ROW_PINS[i], INPUT_PULLUP);
	}
	for (int i = 0; i < COLS; i++) {
		pinMode(COL_PINS[i], OUTPUT);
		digitalWrite(COL_PINS[i], HIGH);
	}
	Serial.println("GPIO and display initialized.");
}

void display_code_for_level(int level) {
	char code[3] = {0}; // Temp storage for level code
	strncpy(code, levelCodes[level], 2); // Copy the level-specific code

	uint8_t segs[4] = {0}; // Segments to send to the display

	// Check if the first character is a letter and map it
	if (isalpha(code[0])) {
		if (code[0] >= 'A' && code[0] <= 'F') {
			segs[0] = SEGMENT_MAP[code[0] - 'A'
								  + 10]; // Maps A-F to their segment patterns
		} else {
			// Handle unexpected characters or extend SEGMENT_MAP for more letters
			segs[0] = 0; // Display nothing for undefined letters
		}
	} else {
		// Assume it's a number and map directly
		segs[0] = SEGMENT_MAP[code[0] - '0'];
	}

	// Check if the second character is a digit and map it
	if (isdigit(code[1])) {
		segs[1] = SEGMENT_MAP[code[1] - '0'];
	} else {
		// Handle unexpected characters
		segs[1] = 0; // Display nothing for undefined digits
	}

	// Set only the first two segments, leave others blank
	display.setSegments(segs, 2, 0); // Display on leftmost two digits
}

pb_global_state_t pb_hook_mod_state_read() { return puzzleState; }

void pb_hook_mod_state_write(pb_global_state_t state) { puzzleState = state; }

void pb_hook_ev_main_state_update(pb_global_state_t state) {
	Serial.println("WE IN BOISS");
}

void setup() {
	Serial.begin(115200);
	pinMode(SOLVED_PIN, OUTPUT);
	digitalWrite(SOLVED_PIN, LOW);
	display.setBrightness(0x0f);
	initialize_system();
}

void loop() {
	switch (puzzleState) {
		case PB_GS_PLAYING:
			display_code_for_level(currentLevel);
			check_button_press();
			delay(100);
			break;
		case PB_GS_SOLVED:
			Serial.println("STATE = PB_GS_SOLVED");
			display_final_code(bombCode);
			digitalWrite(SOLVED_PIN, HIGH);
			break;
		case PB_GS_NOINIT:
			Serial.println("STATE = PB_GS_NOINIT");
			blink_display(0);
			break;
		case PB_GS_IDLE:
			Serial.println("STATE = PB_GS_IDLE");
			blink_display(1);
			break;
	}
}