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

#define REGISTER_SERVOA_POSITIONH 0
#define REGISTER_SERVOA_POSITIONL 1
#define REGISTER_SERVOA_KPA 2
#define REGISTER_SERVOA_KPB 3
#define REGISTER_SERVOA_KPC 4
#define REGISTER_SERVOA_KPD 5
#define REGISTER_SERVOA_KIA 6
#define REGISTER_SERVOA_KIB 7
#define REGISTER_SERVOA_KIC 8
#define REGISTER_SERVOA_KID 9
#define REGISTER_SERVOA_KDA 10
#define REGISTER_SERVOA_KDB 11
#define REGISTER_SERVOA_KDC 12
#define REGISTER_SERVOA_KDD 13
#define REGISTER_SERVOB_POSITIONH 14
#define REGISTER_SERVOB_POSITIONL 15
#define REGISTER_SERVOB_KPA 16
#define REGISTER_SERVOB_KPB 17
#define REGISTER_SERVOB_KPC 18
#define REGISTER_SERVOB_KPD 19
#define REGISTER_SERVOB_KIA 20
#define REGISTER_SERVOB_KIB 21
#define REGISTER_SERVOB_KIC 22
#define REGISTER_SERVOB_KID 23
#define REGISTER_SERVOB_KDA 24
#define REGISTER_SERVOB_KDB 25
#define REGISTER_SERVOB_KDC 26
#define REGISTER_SERVOB_KDD 27

// 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

#define Slave_Address 0x69

// I2C Slave Register Map
#define REGISTER_COUNT 27
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)
{
    // Select ADC channel with MUX bits, clear left-adjust (ADMUX[5] = 0)
    ADMUX = (1 << REFS0) | (channel & 0x07);

    // 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 || motor->target == 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 (motor->target > 0)
    {
        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);
}

uint16_t i = 64535;

uint8_t reg_pointer = 0;
bool expecting_address = true;

int main(void)
{
    DDRA = (1 << 7); // LED
    PORTA = (1 << 7);
    I2C_Slave_Init(Slave_Address);

    *(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 = 64535;
        }
        update_motor(&motor_a);
        update_motor(&motor_b);
        if (TWCR & (1 << TWINT))
        {
            /* Wait to be addressed */
            int8_t status = TWSR & 0xF8;                 /* Read TWI status register with masking lower three bits */
            if (status == 0x60 || status == 0x68) /* Check weather own SLA+W received & ack returned (TWEA = 1) */
            {
                do
                {
                    int8_t byte = I2C_Slave_Receive();
                    if (byte == -1)
                        break;

                    if (expecting_address)
                    {
                        reg_pointer = byte;
                        expecting_address = false;
                    }
                    else
                    {
                        ServoMotor *currentMotor;
                        uint8_t offset = 0;

                        // Decide which motor
                        if (reg_pointer >= REGISTER_SERVOB_POSITIONH)
                        {
                            offset = REGISTER_SERVOB_POSITIONH;
                            currentMotor = &motor_b;
                        }
                        else
                        {
                            offset = 0;
                            currentMotor = &motor_a;
                        }

                        uint8_t local_reg = reg_pointer - offset;

                        switch (local_reg)
                        {
                        case REGISTER_SERVOA_POSITIONH:
                            currentMotor->target &= 0x00FF;
                            currentMotor->target |= ((uint16_t)byte << 8);
                            break;
                        case REGISTER_SERVOA_POSITIONL:
                            currentMotor->target &= 0xFF00;
                            currentMotor->target |= byte;
                            break;

                        case REGISTER_SERVOA_KPA:
                        case REGISTER_SERVOA_KPB:
                        case REGISTER_SERVOA_KPC:
                        case REGISTER_SERVOA_KPD:
                            *((uint8_t *)&currentMotor->kp + (local_reg - REGISTER_SERVOA_KPA)) = byte;
                            break;

                        case REGISTER_SERVOA_KIA:
                        case REGISTER_SERVOA_KIB:
                        case REGISTER_SERVOA_KIC:
                        case REGISTER_SERVOA_KID:
                            *((uint8_t *)&currentMotor->ki + (local_reg - REGISTER_SERVOA_KIA)) = byte;
                            break;

                        case REGISTER_SERVOA_KDA:
                        case REGISTER_SERVOA_KDB:
                        case REGISTER_SERVOA_KDC:
                        case REGISTER_SERVOA_KDD:
                            *((uint8_t *)&currentMotor->kd + (local_reg - REGISTER_SERVOA_KDA)) = byte;
                            break;

                        default:
                            // Optional: save to general-purpose register map
                            registers[reg_pointer] = byte;
                            break;
                        }

                        reg_pointer++; // Optional auto-increment
                    }
                } while (1);
                expecting_address = true; // Reset for next transaction
            }
            if (status == 0xA8 || status == 0xB0) /* Check weather own SLA+R received & ack returned (TWEA = 1) */
            {
                // READ
                char ret;
                ServoMotor *currentMotor;
                uint8_t offset = 0;

                // Choose motor and adjust reg_pointer for local motor indexing
                if (reg_pointer >= REGISTER_SERVOB_POSITIONH)
                {
                    offset = REGISTER_SERVOB_POSITIONH;
                    currentMotor = &motor_b;
                }
                else
                {
                    offset = 0;
                    currentMotor = &motor_a;
                }

                uint8_t local_reg = reg_pointer - offset;

                switch (local_reg)
                {
                case REGISTER_SERVOA_POSITIONH:
                    ret = (currentMotor->current >> 8) & 0xFF;
                    break;
                case REGISTER_SERVOA_POSITIONL:
                    ret = currentMotor->current & 0xFF;
                    break;

                case REGISTER_SERVOA_KPA:
                case REGISTER_SERVOA_KPB:
                case REGISTER_SERVOA_KPC:
                case REGISTER_SERVOA_KPD:
                    ret = *((uint8_t *)&currentMotor->kp + (local_reg - REGISTER_SERVOA_KPA));
                    break;

                case REGISTER_SERVOA_KIA:
                case REGISTER_SERVOA_KIB:
                case REGISTER_SERVOA_KIC:
                case REGISTER_SERVOA_KID:
                    ret = *((uint8_t *)&currentMotor->ki + (local_reg - REGISTER_SERVOA_KIA));
                    break;

                case REGISTER_SERVOA_KDA:
                case REGISTER_SERVOA_KDB:
                case REGISTER_SERVOA_KDC:
                case REGISTER_SERVOA_KDD:
                    ret = *((uint8_t *)&currentMotor->kd + (local_reg - REGISTER_SERVOA_KDA));
                    break;

                default:
                    ret = registers[reg_pointer];
                    break;
                }
                I2C_Slave_Transmit(ret);
                reg_pointer++;
                break;
            }
            if (status == 0x70 || status == 0x78) /* Check weather general call received & ack returned (TWEA = 1) */
                continue;                         /* If yes then return 2 to indicate ack returned */
            else
                continue; /* Else continue */
        }
    }
}