#include <avr/io.h>
#include <avr/interrupt.h>
#include <util/delay.h>
#include <stdlib.h>
#include <stdbool.h>
#include <avr/iom32.h>

// Motor A
#define MOTOR_A_POT      2
#define MOTOR_A_PIN_A    PD5
#define MOTOR_A_PIN_B    PD7
#define MOTOR_A_PIN_A_OCR    OCR1A
#define MOTOR_A_PIN_B_OCR    OCR2

// Motor B
#define MOTOR_B_POT      3
#define MOTOR_B_PIN_A    PB3
#define MOTOR_B_PIN_B    PD4
#define MOTOR_B_PIN_A_OCR    OCR0
#define MOTOR_B_PIN_B_OCR    OCR1B

// I2C Slave Register Map
#define REGISTER_COUNT 16
volatile uint8_t registers[REGISTER_COUNT];

// I2C State
volatile uint8_t reg_address = 0;
volatile bool reg_address_received = false;

typedef struct {
    uint8_t pot_channel;
    volatile uint8_t *pin_a_port;
    volatile uint8_t *pin_a_ddr;
    uint8_t pin_a_bit;
    volatile uint8_t *pin_b_port;
    volatile uint8_t *pin_b_ddr;
    uint8_t pin_b_bit;

    volatile uint8_t *ocr;

    float kp, ki, kd;
    int16_t target;
    int16_t current;
    float integral;
    int16_t last_error;

    bool pot_dir;  // 1 or -1
    bool motor_dir; // 1 or -1
    volatile void *ocr_a;
    volatile void *ocr_b;
    bool ocr_a_16bit;
    bool ocr_b_16bit;
} ServoMotor;

ServoMotor motor_a = {
        .pot_channel = 2,
        .pin_a_port = &PORTD, .pin_a_bit = PD5, .pin_a_ddr = &DDRD,
        .pin_b_port = &PORTD, .pin_b_bit = PD7, .pin_b_ddr = &DDRD,
        .kp = 1.0f, .ki = 0.0f, .kd = 0.0f,
        .pot_dir = 1,
        .motor_dir = 1,
        .ocr_a = &MOTOR_A_PIN_A_OCR,
        .ocr_b = &MOTOR_A_PIN_B_OCR,
};

ServoMotor motor_b = {
        .pot_channel = 3,
        .pin_a_port = &PORTB, .pin_a_bit = PB3, .pin_a_ddr = &DDRB,
        .pin_b_port = &PORTD, .pin_b_bit = PD4, .pin_b_ddr = &DDRD,
        .kp = 1.0f, .ki = 0.0f, .kd = 0.0f,
        .pot_dir = 1,
        .motor_dir = 1,
        .ocr_a = &MOTOR_B_PIN_A_OCR,
        .ocr_b = &MOTOR_B_PIN_B_OCR,
};

void set_ocr(volatile void *reg, bool is_16bit, uint16_t value) {
    if (is_16bit) {
        *((volatile uint16_t *) reg) = value;
    } else {
        *((volatile uint8_t *) reg) = (uint8_t) value;
    }
}

void setup_pwm_motor_a(void) {
    // Setup Timer1 (shared)
    DDRD |= (1 << PD5); // OC1A
    TCCR1A |= (1 << COM1A1);
    TCCR1A |= (1 << COM1B1); // Also needed for Motor B on OC1B (PD4)
    TCCR1B |= (1 << WGM13) | (1 << WGM12) | (1 << CS11); // Fast PWM, prescaler 8
    TCCR1A |= (1 << WGM11); // Complete mode 14
    ICR1 = 255;

    // Setup Timer2 (OC2 for PD7)
    DDRD |= (1 << PD7);
    TCCR2 |= (1 << WGM20) | (1 << WGM21);
    TCCR2 |= (1 << COM21); // Non-inverting
    TCCR2 |= (1 << CS21);
}

void setup_pwm_motor_b(void) {
    // Setup Timer0 (OC0 for PB3)
    DDRB |= (1 << PB3);
    TCCR0 |= (1 << WGM00) | (1 << WGM01); // Fast PWM
    TCCR0 |= (1 << COM01); // Non-inverting
    TCCR0 |= (1 << CS01);

    // OC1B on PD4 (Timer1 already configured in setup_pwm_motor_a)
    DDRD |= (1 << PD4); // Just make sure it's output
}


uint16_t read_adc(uint8_t channel) {
    // Mask channel to stay within 0–7
    channel &= 0x07;

    // Select ADC channel with MUX bits, clear left-adjust (ADMUX[5] = 0)
    ADMUX = (ADMUX & 0xF0) | channel;

    // Start single conversion
    ADCSRA |= (1 << ADSC);

    // Wait for conversion to finish
    while (ADCSRA & (1 << ADSC));

    // Return 10-bit ADC result
    return ADC;
}


void control_motor(ServoMotor *motor, uint8_t pwm, int8_t direction) {
    direction *= motor->motor_dir;

    if (pwm == 0) {
        // Coast: both LOW
        *(motor->pin_a_port) &= ~(1 << motor->pin_a_bit);
        *(motor->pin_b_port) &= ~(1 << motor->pin_b_bit);
        set_ocr(motor->ocr_a, motor->ocr_a_16bit, 0);
        set_ocr(motor->ocr_b, motor->ocr_b_16bit, 0);
        return;
    }

    if (direction > 0) {
        // PWM on A, B LOW
        *(motor->pin_b_port) &= ~(1 << motor->pin_b_bit);   // Direction LOW
        set_ocr(motor->ocr_a, motor->ocr_a_16bit, pwm);
        set_ocr(motor->ocr_b, motor->ocr_b_16bit, 0);
    } else {
        // PWM on B, A LOW
        *(motor->pin_a_port) &= ~(1 << motor->pin_a_bit);   // Direction LOW
        set_ocr(motor->ocr_a, motor->ocr_a_16bit, 0);
        set_ocr(motor->ocr_b, motor->ocr_b_16bit, pwm);
    }
}


void update_motor(ServoMotor *motor) {
    motor->current = motor->pot_dir * read_adc(motor->pot_channel);
    int16_t error = motor->target - motor->current;

    motor->integral += error;
    int16_t derivative = error - motor->last_error;
    motor->last_error = error;

    int16_t output = motor->kp * error + motor->ki * motor->integral + motor->kd * derivative;

    int8_t direction = (output >= 0) ? 1 : -1;
    uint8_t pwm = abs(output);
    if (pwm > 255) pwm = 255;

    control_motor(motor, pwm, direction);
}


uint8_t i = 127;

ISR(TWI_vect) {
    switch (TWSR & 0xF8) {
        case 0x60: // Own SLA+W received, ACK returned
        case 0x68: // Arbitration lost, own SLA+W received, ACK returned
            reg_address_received = false; // Reset for new transfer
            TWCR |= (1 << TWINT) | (1 << TWEA); // ACK next byte
            break;

        case 0x80: // Data received, ACK returned
        case 0x90: // Data received (General Call), ACK returned
            if (!reg_address_received) {
                reg_address = TWDR; // First received byte = register address
                reg_address_received = true;
            } else {
                if (reg_address < REGISTER_COUNT) {
                    registers[reg_address++] = TWDR; // Store received data, then auto-increment address
                }
            }
            TWCR |= (1 << TWINT) | (1 << TWEA); // ACK next byte
            break;

        case 0xA8: // Own SLA+R received, ACK returned
        case 0xB0: // Arbitration lost, own SLA+R received, ACK returned
            if (reg_address < REGISTER_COUNT) {
                TWDR = registers[reg_address++]; // Load data to send
            } else {
                TWDR = 0xFF; // Out of range, send dummy
            }
            TWCR |= (1 << TWINT) | (1 << TWEA); // ACK next byte
            break;

        case 0xB8: // Data transmitted, ACK received
            if (reg_address < REGISTER_COUNT) {
                TWDR = registers[reg_address++];
            } else {
                TWDR = 0xFF;
            }
            TWCR |= (1 << TWINT) | (1 << TWEA); // ACK next byte
            break;

        case 0xC0: // Data transmitted, NACK received (done)
        case 0xC8: // Last byte transmitted, ACK received
        case 0x88: // Data received, NACK returned
            TWCR |= (1 << TWINT) | (1 << TWEA); // Done
            break;

        case 0x00: // Bus error
            TWCR |= (1 << TWSTO) | (1 << TWINT) | (1 << TWEA); // Recover
            break;

        default:
            TWCR |= (1 << TWINT) | (1 << TWEA); // Default ACK
            break;
    }
}


int main(void) {
    ADCSRA |= (1 << ADEN);
    DDRA = (1 << 7); //LED

    TWAR = (0x69 << 1) | (1 << 0);
    TWCR = (1 << 6) | (1 << 2) | (1 << 0);
    *(motor_a.pin_a_ddr) |= (1 << motor_a.pin_a_bit);     // Direction pin output
    *(motor_a.pin_b_ddr) |= (1 << motor_a.pin_b_bit);     // Direction pin output
    *(motor_b.pin_a_ddr) |= (1 << motor_b.pin_a_bit);     // Direction pin output
    *(motor_b.pin_b_ddr) |= (1 << motor_b.pin_b_bit);     // Direction pin output

    setup_pwm_motor_a();
    setup_pwm_motor_b();


    while (1) {
        if (!i++) {
            PORTA ^= (1 << 7);
            i = 127;
        }
        update_motor(&motor_a);
        update_motor(&motor_b);
        _delay_ms(10);
    }
}