#include <stdio.h>
#include <math.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "driver/ledc.h"
#include "driver/gpio.h"

// === Servo pins ===
#define SERVO1_GPIO 2  // Base / X
#define SERVO2_GPIO 4  // Shoulder / Y
#define SERVO3_GPIO 5  // Elbow / Z
#define SERVO4_GPIO 6  // Pen Up/Down
#define SERVO5_GPIO 7  // Extra motor 1
#define SERVO6_GPIO 8  // Extra motor 2

// === Buttons ===
#define MODE_BUTTON 19
#define ACTION_BUTTON 18
#define RESET_BUTTON 17

// === Servo config ===
#define SERVO_MIN_PULSEWIDTH_US 500
#define SERVO_MAX_PULSEWIDTH_US 2500
#define SERVO_MAX_DEGREE 180
#define SERVO_FREQUENCY 50

// === State ===
volatile int mode = 0; // start in NOT SELECTED state

uint32_t angle_to_duty(float angle) {
    float us = SERVO_MIN_PULSEWIDTH_US +
               (SERVO_MAX_PULSEWIDTH_US - SERVO_MIN_PULSEWIDTH_US) * (angle / SERVO_MAX_DEGREE);
    uint32_t duty = (us * 8192) / 20000;
    return duty;
}

void move_servo(int channel, float angle) {
    ledc_set_duty(LEDC_LOW_SPEED_MODE, channel, angle_to_duty(angle));
    ledc_update_duty(LEDC_LOW_SPEED_MODE, channel);
}

void pen_up() {
    move_servo(LEDC_CHANNEL_3, 0);
}
void pen_down() {
    move_servo(LEDC_CHANNEL_3, 90);
}

void go_home() {
    move_servo(LEDC_CHANNEL_0, 90);
    move_servo(LEDC_CHANNEL_1, 90);
    move_servo(LEDC_CHANNEL_2, 90);
    pen_up();
    move_servo(LEDC_CHANNEL_4, 90);
    move_servo(LEDC_CHANNEL_5, 90);
}

void draw_circle() {
    printf("Circle drawing started\n");
    pen_down();
    for (int i = 0; i < 36; i++) {
        move_servo(LEDC_CHANNEL_0, 90 + 10 * cos(i * M_PI / 18));
        move_servo(LEDC_CHANNEL_1, 90 + 10 * sin(i * M_PI / 18));
        vTaskDelay(100 / portTICK_PERIOD_MS);
    }
    pen_up();
    printf("Circle drawing done\n");
}

void draw_square() {
    printf("Square drawing started\n");
    pen_down();
    for (int i = 0; i < 4; i++) {
        move_servo(LEDC_CHANNEL_0, 90 + 20 * ((i == 1 || i == 2) ? 1 : -1));
        move_servo(LEDC_CHANNEL_1, 90 + 20 * ((i >= 2) ? 1 : -1));
        vTaskDelay(500 / portTICK_PERIOD_MS);
    }
    pen_up();
    printf("Square drawing done\n");
}

void draw_rectangle() {
    printf("Rectangle drawing started\n");
    pen_down();
    for (int i = 0; i < 4; i++) {
        move_servo(LEDC_CHANNEL_0, 90 + ((i % 2) ? 30 : 10) * ((i == 1 || i == 2) ? 1 : -1));
        move_servo(LEDC_CHANNEL_1, 90 + ((i >= 2) ? 20 : -20));
        vTaskDelay(500 / portTICK_PERIOD_MS);
    }
    pen_up();
    printf("Rectangle drawing done\n");
}

void print_cube() {
    printf("Cube printing started\n");
    pen_down();
    for (int z = 0; z < 5; z++) {
        draw_square();
        move_servo(LEDC_CHANNEL_2, 90 + z * 5); // lift Z for layers
        vTaskDelay(200 / portTICK_PERIOD_MS);
    }
    pen_up();
    printf("Cube printing done\n");
}

void pick_and_place() {
    printf("Pick and place started\n");
    move_servo(LEDC_CHANNEL_4, 45); // extra motor example
    move_servo(LEDC_CHANNEL_5, 45);
    vTaskDelay(500 / portTICK_PERIOD_MS);
    move_servo(LEDC_CHANNEL_4, 135);
    move_servo(LEDC_CHANNEL_5, 135);
    printf("Pick and place done\n");
}

void mode_button_task(void *pvParameter) {
    int last_state = 1;
    while (1) {
        int state = gpio_get_level(MODE_BUTTON);
        if (state == 0 && last_state == 1) {
            mode++;
            if (mode > 3) mode = 1;
            printf("Selected Mode: %d\n", mode);
            if (mode == 1) printf("Mode 1: Action Button Press Once for circle, Twice for square, Thrice for rectangle\n");
            if (mode == 2) printf("Mode 2: Press Action Button for cube printing\n");
            if (mode == 3) printf("Mode 3: Press Action Button for pick and place\n");
            vTaskDelay(300 / portTICK_PERIOD_MS);
        }
        last_state = state;
        vTaskDelay(50 / portTICK_PERIOD_MS);
    }
}

void action_button_task(void *pvParameter) {
    int last_state = 1;
    int press_count = 0;
    int timeout = 0;

    while (1) {
        int state = gpio_get_level(ACTION_BUTTON);
        if (state == 0 && last_state == 1) {
            press_count++;
            timeout = 0;
        }
        if (press_count > 0 && state == 1) {
            timeout++;
            if (timeout > 10) {
                if (mode == 1) {
                    if (press_count == 1) draw_circle();
                    else if (press_count == 2) draw_square();
                    else if (press_count == 3) draw_rectangle();
                } else if (mode == 2 && press_count == 1) {
                    print_cube();
                } else if (mode == 3 && press_count == 1) {
                    pick_and_place();
                } else {
                    printf("No valid mode selected!\n");
                }
                press_count = 0;
                timeout = 0;
                go_home();
            }
        }
        last_state = state;
        vTaskDelay(100 / portTICK_PERIOD_MS);
    }
}

void reset_button_task(void *pvParameter) {
    int last_state = 1;
    while (1) {
        int state = gpio_get_level(RESET_BUTTON);
        if (state == 0 && last_state == 1) {
            go_home();
            mode = 0;
            printf("Reset done.\nMode Button: 1st Press for Mode1, 2nd Press for Mode2, 3rd for Mode3:\n");
            vTaskDelay(300 / portTICK_PERIOD_MS);
        }
        last_state = state;
        vTaskDelay(50 / portTICK_PERIOD_MS);
    }
}

void app_main() {
    printf("I am Ready.\nMode Button: 1st Press for Mode1, 2nd Press for Mode2, 3rd for Mode3:\n");

    ledc_timer_config_t ledc_timer = {
        .duty_resolution = LEDC_TIMER_13_BIT,
        .freq_hz = SERVO_FREQUENCY,
        .speed_mode = LEDC_LOW_SPEED_MODE,
        .timer_num = LEDC_TIMER_0
    };
    ledc_timer_config(&ledc_timer);

    ledc_channel_config_t channels[6] = {
        {.channel = LEDC_CHANNEL_0, .gpio_num = SERVO1_GPIO, .speed_mode = LEDC_LOW_SPEED_MODE, .hpoint = 0, .timer_sel = LEDC_TIMER_0},
        {.channel = LEDC_CHANNEL_1, .gpio_num = SERVO2_GPIO, .speed_mode = LEDC_LOW_SPEED_MODE, .hpoint = 0, .timer_sel = LEDC_TIMER_0},
        {.channel = LEDC_CHANNEL_2, .gpio_num = SERVO3_GPIO, .speed_mode = LEDC_LOW_SPEED_MODE, .hpoint = 0, .timer_sel = LEDC_TIMER_0},
        {.channel = LEDC_CHANNEL_3, .gpio_num = SERVO4_GPIO, .speed_mode = LEDC_LOW_SPEED_MODE, .hpoint = 0, .timer_sel = LEDC_TIMER_0},
        {.channel = LEDC_CHANNEL_4, .gpio_num = SERVO5_GPIO, .speed_mode = LEDC_LOW_SPEED_MODE, .hpoint = 0, .timer_sel = LEDC_TIMER_0},
        {.channel = LEDC_CHANNEL_5, .gpio_num = SERVO6_GPIO, .speed_mode = LEDC_LOW_SPEED_MODE, .hpoint = 0, .timer_sel = LEDC_TIMER_0}
    };
    for (int i = 0; i < 6; i++) {
        channels[i].duty = 0;
        ledc_channel_config(&channels[i]);
    }

    gpio_reset_pin(MODE_BUTTON);
    gpio_set_direction(MODE_BUTTON, GPIO_MODE_INPUT);
    gpio_pullup_en(MODE_BUTTON);

    gpio_reset_pin(ACTION_BUTTON);
    gpio_set_direction(ACTION_BUTTON, GPIO_MODE_INPUT);
    gpio_pullup_en(ACTION_BUTTON);

    gpio_reset_pin(RESET_BUTTON);
    gpio_set_direction(RESET_BUTTON, GPIO_MODE_INPUT);
    gpio_pullup_en(RESET_BUTTON);

    go_home();

    xTaskCreate(mode_button_task, "mode_button_task", 2048, NULL, 1, NULL);
    xTaskCreate(action_button_task, "action_button_task", 4096, NULL, 1, NULL);
    xTaskCreate(reset_button_task, "reset_button_task", 2048, NULL, 1, NULL);
}