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

// Register Map Definitions
#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 Pins
#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

#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 28
volatile uint8_t registers[REGISTER_COUNT];

// I2C State
volatile uint8_t reg_pointer = 0;
volatile bool expecting_address = true;

// Servo Motor Structure
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;
    float kp, ki, kd;
    int16_t target;
    int16_t current;
    float integral;
    int16_t last_error;
    int8_t pot_dir;  // Direction multiplier (1 or -1)
    int8_t motor_dir; // Direction multiplier (1 or -1)
    volatile void *ocr_a;
    volatile void *ocr_b;
    bool ocr_a_16bit;
    bool ocr_b_16bit;
} ServoMotor;

// Motor Configuration
ServoMotor motor_a = {
    .pot_channel = MOTOR_A_POT,
    .pin_a_port = &PORTD,
    .pin_a_bit = MOTOR_A_PIN_A,
    .pin_a_ddr = &DDRD,
    .pin_b_port = &PORTD,
    .pin_b_bit = MOTOR_A_PIN_B,
    .pin_b_ddr = &DDRD,
    .kp = 1.0f,
    .ki = 0.0f,
    .kd = 0.0f,
    .target = 0,
    .current = 0,
    .integral = 0,
    .last_error = 0,
    .pot_dir = 1,
    .motor_dir = 1,
    .ocr_a = &MOTOR_A_PIN_A_OCR,
    .ocr_b = &MOTOR_A_PIN_B_OCR,
    .ocr_a_16bit = true,
    .ocr_b_16bit = false
};

ServoMotor motor_b = {
    .pot_channel = MOTOR_B_POT,
    .pin_a_port = &PORTB,
    .pin_a_bit = MOTOR_B_PIN_A,
    .pin_a_ddr = &DDRB,
    .pin_b_port = &PORTD,
    .pin_b_bit = MOTOR_B_PIN_B,
    .pin_b_ddr = &DDRD,
    .kp = 1.0f,
    .ki = 0.0f,
    .kd = 0.0f,
    .target = 0,
    .current = 0,
    .integral = 0,
    .last_error = 0,
    .pot_dir = 1,
    .motor_dir = 1,
    .ocr_a = &MOTOR_B_PIN_A_OCR,
    .ocr_b = &MOTOR_B_PIN_B_OCR,
    .ocr_a_16bit = false,
    .ocr_b_16bit = true
};

// Function to set OCR registers (8-bit or 16-bit)
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;
    }
}

// Setup Timer1 for PWM (used by both motors)
void setup_pwm_motor_a(void) {
    // Setup Timer1 (shared)
    DDRD |= (1 << PD5); // OC1A output
    TCCR1A |= (1 << COM1A1); // Non-inverting PWM
    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; // Top value for PWM (8-bit resolution)

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

// Setup Timer0 for PWM
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); // Prescaler 8

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

// Initialize ADC
void adc_init(void) {
    // AREF = AVcc, ADC Left Adjust Result = 0
    ADMUX = (1 << REFS0);
    
    // Enable ADC, prescaler = 128 (16MHz/128 = 125kHz)
    ADCSRA = (1 << ADEN) | (1 << ADPS2) | (1 << ADPS1) | (1 << ADPS0);
}

// Read ADC value
uint16_t read_adc(uint8_t channel) {
    // Select ADC channel with safety mask
    ADMUX = (ADMUX & 0xF8) | (channel & 0x07);
    
    // Start single conversion
    ADCSRA |= (1 << ADSC);
    
    // Wait for conversion to complete
    while (ADCSRA & (1 << ADSC));
    
    // Return 10-bit ADC result
    return ADC;
}

// Control motor with PWM and direction
void control_motor(ServoMotor *motor, uint8_t pwm, int8_t direction) {
    // Apply motor direction correction
    direction *= motor->motor_dir;
    
    // Stop motor if PWM is 0 or target is 0
    if (pwm == 0 || motor->target == 0) {
        // Coast mode: 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;
    }
    
    // Apply direction based on target sign
    if (direction > 0) {
        // Forward: PWM on A, B LOW
        *(motor->pin_b_port) &= ~(1 << motor->pin_b_bit);
        set_ocr(motor->ocr_a, motor->ocr_a_16bit, pwm);
        set_ocr(motor->ocr_b, motor->ocr_b_16bit, 0);
    } else {
        // Reverse: PWM on B, A LOW
        *(motor->pin_a_port) &= ~(1 << motor->pin_a_bit);
        set_ocr(motor->ocr_a, motor->ocr_a_16bit, 0);
        set_ocr(motor->ocr_b, motor->ocr_b_16bit, pwm);
    }
}

// PID control loop for a motor
void update_motor(ServoMotor *motor) {
    // Read current position from potentiometer
    motor->current = motor->pot_dir * read_adc(motor->pot_channel);
    
    // Calculate error
    int16_t error = motor->target - motor->current;
    
    // Update integral term with anti-windup
    motor->integral += error;
    if (motor->integral > 1000) motor->integral = 1000;
    if (motor->integral < -1000) motor->integral = -1000;
    
    // Calculate derivative term
    int16_t derivative = error - motor->last_error;
    motor->last_error = error;
    
    // Calculate PID output
    float output = motor->kp * error + motor->ki * motor->integral + motor->kd * derivative;
    
    // Determine direction and PWM value
    int8_t direction = (output >= 0) ? 1 : -1;
    uint8_t pwm = abs((int16_t)output);
    
    // Cap PWM at 255 (8-bit)
    if (pwm > 255)
        pwm = 255;
    
    // Apply control to motor
    control_motor(motor, pwm, direction);
}

// I2C Interrupt Service Routine
ISR(TWI_vect) {
    uint8_t status = TWSR & 0xF8; // Read TWI status register with masking lower three bits
    
    // Own SLA+W received & ACK returned
    if (status == 0x60 || status == 0x68) {
        TWCR |= (1 << TWINT); // Clear interrupt flag to receive next byte
        return;
    }
    
    // Data received & ACK returned in SLA+W mode
    if (status == 0x80 || status == 0x90) {
        uint8_t received_byte = TWDR;
        
        if (expecting_address) {
            reg_pointer = received_byte;
            expecting_address = false;
        } else {
            ServoMotor *currentMotor;
            uint8_t offset = 0;
            
            // Determine which motor based on register address
            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;
            
            // Process register write based on local register address
            switch (local_reg) {
                case REGISTER_SERVOA_POSITIONH:
                    currentMotor->target &= 0x00FF;
                    currentMotor->target |= ((uint16_t)received_byte << 8);
                    break;
                    
                case REGISTER_SERVOA_POSITIONL:
                    currentMotor->target &= 0xFF00;
                    currentMotor->target |= received_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)) = received_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)) = received_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)) = received_byte;
                    break;
                    
                default:
                    // Store in general register map
                    if (reg_pointer < REGISTER_COUNT) {
                        registers[reg_pointer] = received_byte;
                    }
                    break;
            }
            
            // Auto-increment register pointer
            reg_pointer++;
        }
        
        TWCR |= (1 << TWINT); // Clear interrupt flag
        return;
    }
    
    // STOP or REPEATED START received in slave receiver mode
    if (status == 0xA0) {
        expecting_address = true; // Reset for next transaction
        TWCR |= (1 << TWINT); // Clear interrupt flag
        return;
    }
    
    // Own SLA+R received & ACK returned
    if (status == 0xA8 || status == 0xB0) {
        ServoMotor *currentMotor;
        uint8_t offset = 0;
        uint8_t data_to_send = 0;
        
        // Determine which motor based on register address
        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;
        
        // Process register read based on local register address
        switch (local_reg) {
            case REGISTER_SERVOA_POSITIONH:
                data_to_send = (currentMotor->current >> 8) & 0xFF;
                break;
                
            case REGISTER_SERVOA_POSITIONL:
                data_to_send = currentMotor->current & 0xFF;
                break;
                
            case REGISTER_SERVOA_KPA:
            case REGISTER_SERVOA_KPB:
            case REGISTER_SERVOA_KPC:
            case REGISTER_SERVOA_KPD:
                data_to_send = *((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:
                data_to_send = *((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:
                data_to_send = *((uint8_t *)&currentMotor->kd + (local_reg - REGISTER_SERVOA_KDA));
                break;
                
            default:
                if (reg_pointer < REGISTER_COUNT) {
                    data_to_send = registers[reg_pointer];
                }
                break;
        }
        
        // Send data
        TWDR = data_to_send;
        
        // Auto-increment register pointer
        reg_pointer++;
        
        // Send ACK if not the last byte
        if (reg_pointer < REGISTER_COUNT) {
            TWCR = (1 << TWEN) | (1 << TWINT) | (1 << TWEA) | (1 << TWIE);
        } else {
            // Send NACK for last byte
            TWCR = (1 << TWEN) | (1 << TWINT) | (1 << TWIE);
        }
        return;
    }
    
    // Data byte transmitted & ACK received
    if (status == 0xB8) {
        ServoMotor *currentMotor;
        uint8_t offset = 0;
        uint8_t data_to_send = 0;
        
        // Determine which motor based on register address
        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;
        
        // Process next register read
        switch (local_reg) {
            case REGISTER_SERVOA_POSITIONH:
                data_to_send = (currentMotor->current >> 8) & 0xFF;
                break;
                
            case REGISTER_SERVOA_POSITIONL:
                data_to_send = currentMotor->current & 0xFF;
                break;
                
            case REGISTER_SERVOA_KPA:
            case REGISTER_SERVOA_KPB:
            case REGISTER_SERVOA_KPC:
            case REGISTER_SERVOA_KPD:
                data_to_send = *((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:
                data_to_send = *((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:
                data_to_send = *((uint8_t *)&currentMotor->kd + (local_reg - REGISTER_SERVOA_KDA));
                break;
                
            default:
                if (reg_pointer < REGISTER_COUNT) {
                    data_to_send = registers[reg_pointer];
                }
                break;
        }
        
        // Send data
        TWDR = data_to_send;
        
        // Auto-increment register pointer
        reg_pointer++;
        
        // Send ACK if not the last byte
        if (reg_pointer < REGISTER_COUNT) {
            TWCR = (1 << TWEN) | (1 << TWINT) | (1 << TWEA) | (1 << TWIE);
        } else {
            // Send NACK for last byte
            TWCR = (1 << TWEN) | (1 << TWINT) | (1 << TWIE);
        }
        return;
    }
    
    // Data byte transmitted & NACK received or last byte transmitted & ACK received
    if (status == 0xC0 || status == 0xC8) {
        expecting_address = true;
        TWCR = (1 << TWEN) | (1 << TWINT) | (1 << TWEA) | (1 << TWIE);
        return;
    }
    
    // Default: re-enable TWI interrupt
    TWCR |= (1 << TWINT);
}

// Heartbeat counter for LED blinking
volatile uint16_t heartbeat_counter = 0;

// Main function
int main(void) {
    // Set LED pin as output (PA7)
    DDRA |= (1 << PA7);
    PORTA |= (1 << PA7); // Turn on LED initially
    
    // Initialize I2C as slave
    I2C_Slave_Init(Slave_Address);
    
    // Initialize ADC
    adc_init();
    
    // Configure motor pins as outputs
    *(motor_a.pin_a_ddr) |= (1 << motor_a.pin_a_bit);
    *(motor_a.pin_b_ddr) |= (1 << motor_a.pin_b_bit);
    *(motor_b.pin_a_ddr) |= (1 << motor_b.pin_a_bit);
    *(motor_b.pin_b_ddr) |= (1 << motor_b.pin_b_bit);
    
    // Configure PWM for both motors
    setup_pwm_motor_a();
    setup_pwm_motor_b();
    
    // Initialize all registers to 0
    for (uint8_t i = 0; i < REGISTER_COUNT; i++) {
        registers[i] = 0;
    }
    
    // Enable global interrupts
    sei();
    
    // Main loop
    while (1) {
        // Update motor control
        update_motor(&motor_a);
        update_motor(&motor_b);
        
        // Heartbeat LED - toggle every ~0.5 seconds
        heartbeat_counter++;
        if (heartbeat_counter >= 10000) {
            PORTA ^= (1 << PA7);
            heartbeat_counter = 0;
        }
        
        // Small delay for stability
        _delay_us(50);
    }
    
    return 0; // Never reached
}