#include "step_motor_init.h"

#include <stdbool.h>

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_check.h"
#include "esp_timer.h"

static mcpwm_timer_handle_t s_timer;
static mcpwm_oper_handle_t  s_oper[2];
static mcpwm_cmpr_handle_t  s_cmpr[4]; /* [0]=A+ [1]=A- [2]=B+ [3]=B- */
static mcpwm_gen_handle_t   s_gen[4];
static esp_timer_handle_t   s_openloop_timer;

static uint32_t s_pwm_period_ticks = STEPPER_PWM_RESOL_HZ / STEPPER_PWM_FREQ_HZ;
static bool s_running = false;
static bool s_output_enabled = false;
static bool s_openloop_running = false;
static uint8_t s_openloop_step_idx = 0;
static stepper_motor_direction_t s_direction = STEPPER_DIR_FORWARD;
static const char *TAG = "stepper_motor";

static const uint16_t s_openloop_phase_table[4][4] = {
    { STEPPER_STARTUP_DUTY_RAW, 0, STEPPER_STARTUP_DUTY_RAW, 0 },
    { 0, STEPPER_STARTUP_DUTY_RAW, STEPPER_STARTUP_DUTY_RAW, 0 },
    { 0, STEPPER_STARTUP_DUTY_RAW, 0, STEPPER_STARTUP_DUTY_RAW },
    { STEPPER_STARTUP_DUTY_RAW, 0, 0, STEPPER_STARTUP_DUTY_RAW },
};

static uint32_t duty_raw_to_ticks(uint16_t duty_raw)
{
    if (duty_raw > STEPPER_DUTY_MAX) {
        duty_raw = STEPPER_DUTY_MAX;
    }
    return ((uint64_t)duty_raw * s_pwm_period_ticks) / STEPPER_DUTY_MAX;
    // return ((uint64_t)duty_raw * s_pwm_period_ticks);
    // return ((uint64_t)duty_raw );
}

static void apply_openloop_phase(uint8_t step_idx)
{
    const uint16_t *phase = s_openloop_phase_table[step_idx & 0x03];
    stepper_motor_set_duty_raw(phase[0], phase[1], phase[2], phase[3]);
}

static void stepper_motor_step_once(stepper_motor_direction_t direction)
{
    if (direction == STEPPER_DIR_FORWARD) {
        s_openloop_step_idx = (uint8_t)((s_openloop_step_idx + 1) & 0x03);
    } else {
        s_openloop_step_idx = (uint8_t)((s_openloop_step_idx + 3) & 0x03);
    }
    apply_openloop_phase(s_openloop_step_idx);
}

static void openloop_timer_cb(void *arg)
{
    (void)arg;

    stepper_motor_step_once(s_direction);
}

void stepper_motor_output_enable(void)
{
    ESP_ERROR_CHECK(mcpwm_generator_set_force_level(s_gen[0], -1, true));
    ESP_ERROR_CHECK(mcpwm_generator_set_force_level(s_gen[1], -1, true));
    ESP_ERROR_CHECK(mcpwm_generator_set_force_level(s_gen[2], -1, true));
    ESP_ERROR_CHECK(mcpwm_generator_set_force_level(s_gen[3], -1, true));
    s_output_enabled = true;
}

void stepper_motor_output_disable(void)
{
    ESP_ERROR_CHECK(mcpwm_generator_set_force_level(s_gen[0], 0, true));
    ESP_ERROR_CHECK(mcpwm_generator_set_force_level(s_gen[1], 0, true));
    ESP_ERROR_CHECK(mcpwm_generator_set_force_level(s_gen[2], 0, true));
    ESP_ERROR_CHECK(mcpwm_generator_set_force_level(s_gen[3], 0, true));
    s_output_enabled = false;
}

bool stepper_motor_output_is_enabled(void)
{
    return s_output_enabled;
}

stepper_motor_direction_t stepper_motor_get_direction(void)
{
    return s_direction;
}

void stepper_motor_set_duty_raw(uint16_t a_plus, uint16_t a_minus, uint16_t b_plus, uint16_t b_minus)
{
    if (!s_running) {
        return;
    }

    ESP_ERROR_CHECK(mcpwm_comparator_set_compare_value(s_cmpr[0], duty_raw_to_ticks(a_plus)));
    ESP_ERROR_CHECK(mcpwm_comparator_set_compare_value(s_cmpr[1], duty_raw_to_ticks(a_minus)));
    ESP_ERROR_CHECK(mcpwm_comparator_set_compare_value(s_cmpr[2], duty_raw_to_ticks(b_plus)));
    ESP_ERROR_CHECK(mcpwm_comparator_set_compare_value(s_cmpr[3], duty_raw_to_ticks(b_minus)));
}

void stepper_motor_openloop_start_dir(stepper_motor_direction_t direction)
{
    if (s_openloop_timer == NULL) {
        esp_timer_create_args_t timer_args = {
            .callback = openloop_timer_cb,
            .arg = NULL,
            .dispatch_method = ESP_TIMER_TASK,
            .name = "step_ol",
            .skip_unhandled_events = true,
        };
        ESP_ERROR_CHECK(esp_timer_create(&timer_args, &s_openloop_timer));
    }

    if (!s_running) {
        return;
    }

    s_direction = (direction == STEPPER_DIR_REVERSE) ? STEPPER_DIR_REVERSE : STEPPER_DIR_FORWARD;

    if (s_openloop_running) {
        ESP_ERROR_CHECK(esp_timer_stop(s_openloop_timer));
        s_openloop_running = false;
    }

    if (!s_output_enabled) {
        stepper_motor_output_enable();
    }

    s_openloop_step_idx = 0;
    apply_openloop_phase(s_openloop_step_idx);

    const uint64_t period_us = 1000000ULL / (uint64_t)(STEPPER_OPENLOOP_STEP_HZ * 4U);
    ESP_ERROR_CHECK(esp_timer_start_periodic(s_openloop_timer, period_us));
    s_openloop_running = true;
    ESP_LOGI(TAG,
             "openloop start: dir=%s duty=%u step_hz=%u",
             s_direction == STEPPER_DIR_FORWARD ? "forward" : "reverse",
             STEPPER_STARTUP_DUTY_RAW,
             STEPPER_OPENLOOP_STEP_HZ);
}

void stepper_motor_openloop_start(void)
{
    stepper_motor_openloop_start_dir(s_direction);
}

void stepper_motor_forward(void)
{
    stepper_motor_openloop_start_dir(STEPPER_DIR_FORWARD);
}

void stepper_motor_reverse(void)
{
    stepper_motor_openloop_start_dir(STEPPER_DIR_REVERSE);
}

void stepper_motor_run_steps(stepper_motor_direction_t direction, uint32_t steps)
{
    if (!s_running || steps == 0) {
        return;
    }

    if (s_openloop_timer != NULL && s_openloop_running) {
        ESP_ERROR_CHECK(esp_timer_stop(s_openloop_timer));
        s_openloop_running = false;
    }

    s_direction = (direction == STEPPER_DIR_REVERSE) ? STEPPER_DIR_REVERSE : STEPPER_DIR_FORWARD;

    if (!s_output_enabled) {
        stepper_motor_output_enable();
    }

    for (uint32_t i = 0; i < steps; ++i) {
        stepper_motor_step_once(s_direction);
        vTaskDelay(pdMS_TO_TICKS(1000 / (STEPPER_OPENLOOP_STEP_HZ * 4U)));
    }

    stepper_motor_stop();
}

void stepper_motor_openloop_stop(void)
{
    if (s_openloop_timer != NULL && s_openloop_running) {
        ESP_ERROR_CHECK(esp_timer_stop(s_openloop_timer));
        s_openloop_running = false;
    }

    stepper_motor_set_duty_raw(0, 0, 0, 0);
    stepper_motor_output_disable();
    ESP_LOGI(TAG, "openloop stop");
}

void stepper_motor_stop(void)
{
    stepper_motor_openloop_stop();
}

void stepper_motor_init(void)
{
    gpio_config_t gpio_cfg = {
        .pin_bit_mask = (1ULL << STEPPER_PIN_A1) | (1ULL << STEPPER_PIN_A2) |
                        (1ULL << STEPPER_PIN_B1) | (1ULL << STEPPER_PIN_B2),
        .mode  = GPIO_MODE_OUTPUT,
        .pull_up_en = GPIO_PULLUP_DISABLE,
        .pull_down_en = GPIO_PULLDOWN_DISABLE,
        .intr_type = GPIO_INTR_DISABLE,
    };
    ESP_ERROR_CHECK(gpio_config(&gpio_cfg));
    gpio_set_level(STEPPER_PIN_A1, 0);
    gpio_set_level(STEPPER_PIN_A2, 0);
    gpio_set_level(STEPPER_PIN_B1, 0);
    gpio_set_level(STEPPER_PIN_B2, 0);
    ESP_LOGI(TAG, "Step1: GPIO done");

    mcpwm_timer_config_t timer_cfg = {
        .group_id      = STEPPER_MCPWM_GROUP,
        .clk_src       = MCPWM_TIMER_CLK_SRC_DEFAULT,
        .resolution_hz = STEPPER_PWM_RESOL_HZ,
        .period_ticks  = s_pwm_period_ticks,
        .count_mode    = MCPWM_TIMER_COUNT_MODE_UP,
        .flags         = { .update_period_on_empty = true },
    };
    ESP_ERROR_CHECK(mcpwm_new_timer(&timer_cfg, &s_timer));
    ESP_LOGI(TAG, "Step2: timer created | pwm=%lu Hz | period=%lu",
             (unsigned long)STEPPER_PWM_FREQ_HZ,
             (unsigned long)s_pwm_period_ticks);

    mcpwm_operator_config_t oper_cfg = { .group_id = STEPPER_MCPWM_GROUP };
    ESP_ERROR_CHECK(mcpwm_new_operator(&oper_cfg, &s_oper[0]));
    ESP_ERROR_CHECK(mcpwm_operator_connect_timer(s_oper[0], s_timer));
    ESP_ERROR_CHECK(mcpwm_new_operator(&oper_cfg, &s_oper[1]));
    ESP_ERROR_CHECK(mcpwm_operator_connect_timer(s_oper[1], s_timer));
    ESP_LOGI(TAG, "Step3: operators connected");

    mcpwm_comparator_config_t cmpr_cfg = { .flags.update_cmp_on_tez = true };
    ESP_ERROR_CHECK(mcpwm_new_comparator(s_oper[0], &cmpr_cfg, &s_cmpr[0]));
    ESP_ERROR_CHECK(mcpwm_new_comparator(s_oper[0], &cmpr_cfg, &s_cmpr[1]));
    ESP_ERROR_CHECK(mcpwm_new_comparator(s_oper[1], &cmpr_cfg, &s_cmpr[2]));
    ESP_ERROR_CHECK(mcpwm_new_comparator(s_oper[1], &cmpr_cfg, &s_cmpr[3]));
    ESP_LOGI(TAG, "Step4: comparators created");

    mcpwm_generator_config_t gen_cfg = { .gen_gpio_num = STEPPER_PIN_A1 };
    ESP_ERROR_CHECK(mcpwm_new_generator(s_oper[0], &gen_cfg, &s_gen[0]));
    gen_cfg.gen_gpio_num = STEPPER_PIN_A2;
    ESP_ERROR_CHECK(mcpwm_new_generator(s_oper[0], &gen_cfg, &s_gen[1]));
    gen_cfg.gen_gpio_num = STEPPER_PIN_B1;
    ESP_ERROR_CHECK(mcpwm_new_generator(s_oper[1], &gen_cfg, &s_gen[2]));
    gen_cfg.gen_gpio_num = STEPPER_PIN_B2;
    ESP_ERROR_CHECK(mcpwm_new_generator(s_oper[1], &gen_cfg, &s_gen[3]));
    ESP_LOGI(TAG, "Step5: generators created");

    for (int i = 0; i < 4; ++i) {
        ESP_ERROR_CHECK(mcpwm_generator_set_actions_on_timer_event(s_gen[i],
            MCPWM_GEN_TIMER_EVENT_ACTION(MCPWM_TIMER_DIRECTION_UP,
                MCPWM_TIMER_EVENT_EMPTY, MCPWM_GEN_ACTION_LOW),//MCPWM_GEN_ACTION_LOW
            MCPWM_GEN_TIMER_EVENT_ACTION_END()));
        ESP_ERROR_CHECK(mcpwm_generator_set_actions_on_compare_event(s_gen[i],
            MCPWM_GEN_COMPARE_EVENT_ACTION(MCPWM_TIMER_DIRECTION_UP,
                s_cmpr[i], MCPWM_GEN_ACTION_HIGH),//
            MCPWM_GEN_COMPARE_EVENT_ACTION_END()));
    }
    ESP_LOGI(TAG, "Step6: pwm actions configured");

    stepper_motor_output_disable();
    ESP_ERROR_CHECK(mcpwm_timer_enable(s_timer));
    ESP_ERROR_CHECK(mcpwm_timer_start_stop(s_timer, MCPWM_TIMER_START_NO_STOP));

    s_running = true;
    stepper_motor_set_duty_raw(0, 0, 0, 0);

    ESP_LOGI(TAG, "MCPWM init done, motor standby on boot");
}