#include <stdio.h>
#include "pico/stdlib.h"
#include "hardware/gpio.h"
#include "hardware/timer.h"

#define DIGIT1_PIN 13
#define DIGIT2_PIN 14
#define DIGIT3_PIN 16
#define DIGIT4_PIN 17

#define SEG_A 0
#define SEG_B 1
#define SEG_C 2
#define SEG_D 3
#define SEG_E 4
#define SEG_F 5
#define SEG_G 6

#define BUTTON_PIN 15

const uint8_t digit_patterns[] = {
    0b00111111, // 0
    0b00000110, // 1
    0b01011011, // 2
    0b01001111, // 3
    0b01100110, // 4
    0b01101101, // 5
    0b01111101, // 6
    0b00000111, // 7
    0b01111111, // 8
    0b01101111  // 9
};

void init_gpio() {
    gpio_init(DIGIT1_PIN);
    gpio_set_dir(DIGIT1_PIN, GPIO_OUT);
    gpio_init(DIGIT2_PIN);
    gpio_set_dir(DIGIT2_PIN, GPIO_OUT);
    gpio_init(DIGIT3_PIN);
    gpio_set_dir(DIGIT3_PIN, GPIO_OUT);
    gpio_init(DIGIT4_PIN);
    gpio_set_dir(DIGIT4_PIN, GPIO_OUT);

    for (int i = SEG_A; i <= SEG_G; i++) {
        gpio_init(i);
        gpio_set_dir(i, GPIO_OUT);
    }

    gpio_init(BUTTON_PIN);
    gpio_set_dir(BUTTON_PIN, GPIO_IN);
    gpio_pull_up(BUTTON_PIN);
}

void display_digit(uint8_t digit, uint8_t position) {
    gpio_put(DIGIT1_PIN, 0);
    gpio_put(DIGIT2_PIN, 0);
    gpio_put(DIGIT3_PIN, 0);
    gpio_put(DIGIT4_PIN, 0);

    uint8_t pattern = digit_patterns[digit];
    for (int i = SEG_A; i <= SEG_G; i++) {
        gpio_put(i, pattern & 1);
        pattern >>= 1;
    }

    switch (position) {
        case 0: gpio_put(DIGIT1_PIN, 1); break;
        case 1: gpio_put(DIGIT2_PIN, 1); break;
        case 2: gpio_put(DIGIT3_PIN, 1); break;
        case 3: gpio_put(DIGIT4_PIN, 1); break;
    }
}

void update_display(uint16_t count) {
    uint8_t digits[4];
    digits[0] = count / 1000;
    digits[1] = (count / 100) % 10;
    digits[2] = (count / 10) % 10;
    digits[3] = count % 10;

    for (int i = 0; i < 4; i++) {
        display_digit(digits[i], i);
        sleep_ms(5);  // Slightly longer delay to improve visibility
    }
}

void check_button(uint16_t *count, bool *reset_flag) {
    static bool button_pressed = false;
    if (gpio_get(BUTTON_PIN) == 0) {
        if (!button_pressed) {
            button_pressed = true;
            *count = 0;
            *reset_flag = true;
        }
    } else {
        button_pressed = false;
    }
}

int main() {
    stdio_init_all();
    init_gpio();

    uint16_t count = 0;
    bool reset_flag = false;
    absolute_time_t timer = get_absolute_time();

    while (true) {
        update_display(count);
        check_button(&count, &reset_flag);

        if (absolute_time_diff_us(timer, get_absolute_time()) >= 1000000) {
            timer = get_absolute_time();
            if (!reset_flag) {
                count = (count + 1) % 10000;
            }
            reset_flag = false;
        }
    }

    return 0;
}