#include <avr/io.h>
#include <avr/interrupt.h>
#include <stdio.h>

#define F_CPU 8000000UL // CPU Clock Frequency in Hz

#define BAUD_RATE 9600
#define BAUD_RATE_DIVISOR ((F_CPU / (16UL * BAUD_RATE)) - 1)

// Define segments for each digit
#define A (1<<PB5)
#define B (1<<PB7)
#define C (1<<PB1)
#define D (1<<PB3)
#define E (1<<PB4)
#define F (1<<PB6)
#define G (1<<PB0)
#define DP (1<<PB2)



// Segment patterns for digits 0 to 9
const uint8_t numbers[] = {
    (A | B | C | D | E | F),   // 0
    (B | C),                   // 1
    (A | B | G | E | D),       // 2
    (A | B | C | D | G),       // 3
    (F | G | B | C),           // 4
    (A | F | G | C | D),       // 5
    (A | F | E | D | C | G),   // 6
    (A | B | C),               // 7
    (A | B | C | D | E | F | G), // 8
    (A | B | C | F | G)        // 9
};

// Array to store digits to be displayed
volatile uint8_t digits[] = {0, 0, 0, 0, 0, 0};

// Current digit index
volatile uint8_t currentDigit = 0;

// Function to display a single digit
void displayDigit(uint8_t digit) {
    PORTB = numbers[digit];
}

// Function to turn on a specific digit
void turnOnDigit(uint8_t digit) {
    switch (digit) {
        case 0:
            PORTA |= (1<<PA1);
            break;
        case 1:
            PORTA |= (1<<PA0);
            break;
        case 2:
            PORTD |= (1<<PD2);
            break;
        case 3:
            PORTD |= (1<<PD3);
            break;
        case 4:
            PORTD |= (1<<PD4);
            break;
        case 5:
            PORTD |= (1<<PD5);
            break;
        default:
            break;
    }
}

// Function to turn off all digits
void turnOffAllDigits() {
    PORTA &= ~((1<<PA1) | (1<<PA0));
    PORTD &= ~((1<<PD2) | (1<<PD3) | (1<<PD4) | (1<<PD5));
}

// Timer/Counter1 Compare Match A interrupt
ISR(TIMER1_COMPA_vect) {
    turnOffAllDigits(); // Turn off all digits
    displayDigit(digits[currentDigit]); // Display the current digit
    turnOnDigit(currentDigit); // Turn on the corresponding digit
    currentDigit = (currentDigit + 1) % 6; // Move to the next digit
}

// USART RX Complete interrupt
ISR(USART_RXC_vect) {
    static uint8_t d = 0;
    uint8_t data = UDR; // Read the received data

    if(data == 'X') {
        d = 0; // Reset currentDigit to start receiving digits
    } else if (data >= 0 && data <= 9 && currentDigit < 6) {
        digits[d] = data;
        d++;
    }
    // Echo back the received data
    while (!(UCSRA & (1 << UDRE))); // Wait for the transmit buffer to be empty
    UDR = data; // Send the received data back
}

int uart_putchar(char c, FILE *stream) {
    if (c == '\n') {
        uart_putchar('\r', stream);
    }
    // Wait for empty transmit buffer
    while (!(UCSRA & (1 << UDRE)));
    // Put data into buffer, sends the data
    UDR = c;
    return 0;
}


int main() {
    // Set PORTB as output for segments
    DDRB = 0xFF;
    // Set PORTA and PORTD as output for digit control
    DDRA |= (1<<PA0) | (1<<PA1);
    DDRD |= (1<<PD2) | (1<<PD3) | (1<<PD4) | (1<<PD5);

    // Initialize Timer/Counter1 in CTC mode
    TCCR1B |= (1 << WGM12);
    // Set prescaler to 1024
    TCCR1B |= (1 << CS12) | (0 << CS11) | (1 << CS10);
    // Set compare value for 1ms interrupt at 8MHz CPU clock
    OCR1A = 7;
    // Enable Timer/Counter1 Compare Match A interrupt
    TIMSK |= (1 << OCIE1A);

    // Initialize USART
    // Set baud rate to 9600
    UBRRH = (BAUD_RATE_DIVISOR >> 8);
    UBRRL = BAUD_RATE_DIVISOR;
    // Enable receiver and RX complete interrupt
    UCSRB |= (1<<TXEN) | (1 << RXEN) | (1 << RXCIE);
    // Set frame format: 8 data bits, 1 stop bit
    UCSRC = (1 << UCSZ1) | (1 << UCSZ0);
    // Enable global interrupts
    sei();

    
    static FILE mystdout;
    fdev_setup_stream(&mystdout, uart_putchar, NULL, _FDEV_SETUP_WRITE);
    stdout = &mystdout;
    printf("Started");
    while (1) {
    }

    return 0;
}