/*
 * SPDX-FileCopyrightText: 2024-2025 Espressif Systems (Shanghai) CO LTD
 *
 * SPDX-License-Identifier: Apache-2.0
 */

#include "platform_internal.h"

#include <math.h>

#include "driver/gpio.h"
#include "esp_log.h"
#include "esp_timer.h"
#include "esp_rom_sys.h"
#include "freertos/FreeRTOS.h"
#include "freertos/semphr.h"
#include "sdkconfig.h"
#include "soc/gpio_struct.h"
#include "freertos/portmacro.h"

#ifndef CONFIG_TQ_STEPPER_IN1_GPIO
#error "CONFIG_TQ_STEPPER_IN1_GPIO must be defined in Kconfig"
#endif
#ifndef CONFIG_TQ_STEPPER_IN2_GPIO
#error "CONFIG_TQ_STEPPER_IN2_GPIO must be defined in Kconfig"
#endif
#ifndef CONFIG_TQ_STEPPER_IN3_GPIO
#error "CONFIG_TQ_STEPPER_IN3_GPIO must be defined in Kconfig"
#endif
#ifndef CONFIG_TQ_STEPPER_IN4_GPIO
#error "CONFIG_TQ_STEPPER_IN4_GPIO must be defined in Kconfig"
#endif

#define PLATFORM_STEPPER_IN1_PIN         CONFIG_TQ_STEPPER_IN1_GPIO
#define PLATFORM_STEPPER_IN2_PIN         CONFIG_TQ_STEPPER_IN2_GPIO
#define PLATFORM_STEPPER_IN3_PIN         CONFIG_TQ_STEPPER_IN3_GPIO
#define PLATFORM_STEPPER_IN4_PIN         CONFIG_TQ_STEPPER_IN4_GPIO
#define PLATFORM_STEPPER_START_DELAY_US  2200
#define PLATFORM_STEPPER_ACCEL_STEPS     48
#define PLATFORM_STEPPER_DECEL_STEPS     48
#define PLATFORM_STEPPER_STEPS_PER_REV   4128.0f
#define PLATFORM_STEPPER_POWER_OFF_WAIT_MS 1000
#define PLATFORM_STEPPER_MAX_ROTATION_DEG 100.0f
#define PLATFORM_GPIO_LOW_BANK_MAX_PIN   31
#define PLATFORM_GPIO_HIGH_BANK_OFFSET   32

static const char *TAG = "platform_stepper";
static portMUX_TYPE s_stepper_mux = portMUX_INITIALIZER_UNLOCKED;
static SemaphoreHandle_t s_stepper_motion_mutex;
#if ( configSUPPORT_STATIC_ALLOCATION == 1 )
static StaticSemaphore_t s_stepper_motion_mutex_buf;
#endif
static int s_stepper_phase_index;
static bool s_stepper_energized;

static const int s_step_sequence[8][4] = {
    {1, 0, 0, 0},
    {1, 1, 0, 0},
    {0, 1, 0, 0},
    {0, 1, 1, 0},
    {0, 0, 1, 0},
    {0, 0, 1, 1},
    {0, 0, 0, 1},
    {1, 0, 0, 1}
};

static void add_gpio_level_mask(uint32_t *set0,
                                uint32_t *clr0,
                                uint32_t *set1,
                                uint32_t *clr1,
                                int pin,
                                int level)
{
    if (pin <= PLATFORM_GPIO_LOW_BANK_MAX_PIN) {
        uint32_t mask = (1U << pin);
        if (level) {
            *set0 |= mask;
        } else {
            *clr0 |= mask;
        }
        return;
    }

    uint32_t mask = (1U << (pin - PLATFORM_GPIO_HIGH_BANK_OFFSET));
    if (level) {
        *set1 |= mask;
    } else {
        *clr1 |= mask;
    }
}

static void set_motor_pins(int in1, int in2, int in3, int in4)
{
    uint32_t set0 = 0;
    uint32_t clr0 = 0;
    uint32_t set1 = 0;
    uint32_t clr1 = 0;

    add_gpio_level_mask(&set0, &clr0, &set1, &clr1, PLATFORM_STEPPER_IN1_PIN, in1);
    add_gpio_level_mask(&set0, &clr0, &set1, &clr1, PLATFORM_STEPPER_IN2_PIN, in2);
    add_gpio_level_mask(&set0, &clr0, &set1, &clr1, PLATFORM_STEPPER_IN3_PIN, in3);
    add_gpio_level_mask(&set0, &clr0, &set1, &clr1, PLATFORM_STEPPER_IN4_PIN, in4);

    /* Keep phase updates batched while supporting GPIO32-48 on GPIO.out1. */
    if (set0) GPIO.out_w1ts = set0;
    if (clr0) GPIO.out_w1tc = clr0;
    if (set1) GPIO.out1_w1ts.val = set1;
    if (clr1) GPIO.out1_w1tc.val = clr1;
}

static bool ensure_stepper_motion_mutex(void)
{
    if (s_stepper_motion_mutex != NULL) {
        return true;
    }

#if ( configSUPPORT_STATIC_ALLOCATION == 1 )
    s_stepper_motion_mutex = xSemaphoreCreateMutexStatic(&s_stepper_motion_mutex_buf);
#else
    s_stepper_motion_mutex = xSemaphoreCreateMutex();
#endif

    return (s_stepper_motion_mutex != NULL);
}

static void precise_delay_us(int delay_us)
{
    if (delay_us > 0) {
        esp_rom_delay_us((uint32_t)delay_us);
    }
}

static int stepper_normalize_phase_index(int phase_index)
{
    int normalized = phase_index % 8;
    if (normalized < 0) {
        normalized += 8;
    }
    return normalized;
}

static void stepper_apply_phase_locked(int phase_index)
{
    int index = stepper_normalize_phase_index(phase_index);

    s_stepper_phase_index = index;
    s_stepper_energized = true;
    set_motor_pins(s_step_sequence[index][0],
                   s_step_sequence[index][1],
                   s_step_sequence[index][2],
                   s_step_sequence[index][3]);
}

static void stepper_release_locked(void)
{
    s_stepper_energized = false;
    set_motor_pins(0, 0, 0, 0);
}

static void stepper_reenergize_if_needed(int delay_us)
{
    bool needs_settle = false;

    portENTER_CRITICAL(&s_stepper_mux);
    if (!s_stepper_energized) {
        stepper_apply_phase_locked(s_stepper_phase_index);
        needs_settle = true;
    }
    portEXIT_CRITICAL(&s_stepper_mux);

    if (needs_settle) {
        precise_delay_us(delay_us);
    }
}

static void stepper_advance_phase(bool clockwise)
{
    portENTER_CRITICAL(&s_stepper_mux);
    int next_phase = s_stepper_phase_index + (clockwise ? 1 : -1);
    stepper_apply_phase_locked(next_phase);
    portEXIT_CRITICAL(&s_stepper_mux);
}

static bool stepper_timeout_reached(int timeout_ms, int64_t start_ms)
{
    if (timeout_ms <= 0) {
        return false;
    }

    int64_t now_ms = esp_timer_get_time() / 1000;
    return (now_ms - start_ms) >= timeout_ms;
}

static int stepper_initial_delay_us(int target_delay_us)
{
    if (target_delay_us <= 0) {
        target_delay_us = PLATFORM_STEPPER_SPEED_NORMAL;
    }

    return (target_delay_us > PLATFORM_STEPPER_START_DELAY_US) ?
            target_delay_us :
            PLATFORM_STEPPER_START_DELAY_US;
}

static int estimate_rotation_time_ms(int steps, int target_delay_us)
{
    if (steps <= 0) {
        return 0;
    }

    int max_delay_us = stepper_initial_delay_us(target_delay_us);

    int64_t total_us = (int64_t)steps * (int64_t)max_delay_us;
    return (int)((total_us + 999) / 1000);
}

static void calculate_motion_segments(int steps,
                                      int *out_accel_steps,
                                      int *out_constant_steps,
                                      int *out_decel_steps)
{
    int accel_steps = PLATFORM_STEPPER_ACCEL_STEPS;
    int decel_steps = PLATFORM_STEPPER_DECEL_STEPS;

    if (steps <= 0) {
        accel_steps = 0;
        decel_steps = 0;
    } else if (steps < (accel_steps + decel_steps)) {
        accel_steps = steps / 3;
        decel_steps = steps / 3;
    }

    while ((accel_steps + decel_steps) > steps) {
        if (decel_steps >= accel_steps && decel_steps > 0) {
            decel_steps--;
        } else if (accel_steps > 0) {
            accel_steps--;
        } else {
            break;
        }
    }

    if (out_accel_steps != NULL) {
        *out_accel_steps = accel_steps;
    }
    if (out_constant_steps != NULL) {
        *out_constant_steps = steps - accel_steps - decel_steps;
    }
    if (out_decel_steps != NULL) {
        *out_decel_steps = decel_steps;
    }
}

static int calculate_step_delay_us(int step_index,
                                   int accel_steps,
                                   int constant_steps,
                                   int decel_steps,
                                   int target_delay_us)
{
    int initial_delay_us = stepper_initial_delay_us(target_delay_us);
    int current_delay_us = target_delay_us;

    if (step_index < accel_steps) {
        current_delay_us = initial_delay_us -
                (initial_delay_us - target_delay_us) * step_index / accel_steps;
    } else if (step_index >= (accel_steps + constant_steps) && decel_steps > 0) {
        int decel_progress = step_index - (accel_steps + constant_steps);
        current_delay_us = target_delay_us +
                (initial_delay_us - target_delay_us) * decel_progress / decel_steps;
    }

    return current_delay_us;
}

static esp_err_t rotate_with_accel(int steps,
                                   bool clockwise,
                                   int target_delay_us,
                                   int timeout_ms,
                                   int64_t start_ms)
{
    int accel_steps = 0;
    int constant_steps = 0;
    int decel_steps = 0;

    calculate_motion_segments(steps, &accel_steps, &constant_steps, &decel_steps);
    stepper_reenergize_if_needed(stepper_initial_delay_us(target_delay_us));

    for (int i = 0; i < steps; i++) {
        if (stepper_timeout_reached(timeout_ms, start_ms)) {
            return ESP_ERR_TIMEOUT;
        }

        int current_delay_us = calculate_step_delay_us(i,
                                                       accel_steps,
                                                       constant_steps,
                                                       decel_steps,
                                                       target_delay_us);
        stepper_advance_phase(clockwise);
        precise_delay_us(current_delay_us);
    }

    return ESP_OK;
}

esp_err_t platform_stepper_driver_init(void)
{
    if (!ensure_stepper_motion_mutex()) {
        return ESP_ERR_NO_MEM;
    }

    portENTER_CRITICAL(&s_stepper_mux);
    s_stepper_phase_index = 0;
    s_stepper_energized = false;
    portEXIT_CRITICAL(&s_stepper_mux);

    gpio_config_t io_conf = {
        .pin_bit_mask = (1ULL << PLATFORM_STEPPER_IN1_PIN) |
                        (1ULL << PLATFORM_STEPPER_IN2_PIN) |
                        (1ULL << PLATFORM_STEPPER_IN3_PIN) |
                        (1ULL << PLATFORM_STEPPER_IN4_PIN),
        .mode = GPIO_MODE_OUTPUT,
        .pull_up_en = GPIO_PULLUP_DISABLE,
        .pull_down_en = GPIO_PULLDOWN_DISABLE,
        .intr_type = GPIO_INTR_DISABLE
    };

    esp_err_t err = gpio_config(&io_conf);
    if (err == ESP_OK) {
        ESP_LOGI(TAG,
                 "Stepper ready IN1=%d IN2=%d IN3=%d IN4=%d",
                 PLATFORM_STEPPER_IN1_PIN,
                 PLATFORM_STEPPER_IN2_PIN,
                 PLATFORM_STEPPER_IN3_PIN,
                 PLATFORM_STEPPER_IN4_PIN);
    }
    return err;
}

esp_err_t platform_stepper_driver_rotate_angle_with_accel(float angle,
                                                          int target_delay_us,
                                                          int timeout_ms)
{
    if (!ensure_stepper_motion_mutex()) {
        return ESP_ERR_INVALID_STATE;
    }
    if (fabsf(angle) > PLATFORM_STEPPER_MAX_ROTATION_DEG) {
        ESP_LOGE(TAG,
                 "rotation rejected: angle %.1f exceeds %.1f deg limit",
                 angle,
                 PLATFORM_STEPPER_MAX_ROTATION_DEG);
        return ESP_ERR_INVALID_ARG;
    }

    int64_t start_ms = esp_timer_get_time() / 1000;
    TickType_t lock_ticks = (timeout_ms > 0) ? pdMS_TO_TICKS(timeout_ms) : portMAX_DELAY;
    if (timeout_ms > 0 && lock_ticks == 0) {
        lock_ticks = 1;
    }
    if (xSemaphoreTake(s_stepper_motion_mutex, lock_ticks) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }

    int steps = (int)lroundf(angle * PLATFORM_STEPPER_STEPS_PER_REV / 360.0f);
    if (steps == 0) {
        (void)xSemaphoreGive(s_stepper_motion_mutex);
        return ESP_OK;
    }

    if (target_delay_us <= 0) {
        target_delay_us = PLATFORM_STEPPER_SPEED_NORMAL;
    }
    if (target_delay_us < PLATFORM_STEPPER_SPEED_ULTRA_FAST) {
        target_delay_us = PLATFORM_STEPPER_SPEED_ULTRA_FAST;
    }

    int absolute_steps = (steps > 0) ? steps : -steps;
    if (timeout_ms > 0) {
        int elapsed_ms = (int)(esp_timer_get_time() / 1000 - start_ms);
        int remaining_ms = timeout_ms - elapsed_ms;
        if (remaining_ms <= 0) {
            (void)xSemaphoreGive(s_stepper_motion_mutex);
            return ESP_ERR_TIMEOUT;
        }
        int estimated_ms = estimate_rotation_time_ms(absolute_steps, target_delay_us);
        if (estimated_ms > remaining_ms) {
            (void)xSemaphoreGive(s_stepper_motion_mutex);
            return ESP_ERR_TIMEOUT;
        }
    }

    esp_err_t err = ESP_OK;
    err = rotate_with_accel(absolute_steps,
                            steps > 0,
                            target_delay_us,
                            timeout_ms,
                            start_ms);
    if (err != ESP_OK) {
        portENTER_CRITICAL(&s_stepper_mux);
        stepper_release_locked();
        portEXIT_CRITICAL(&s_stepper_mux);
    }

    (void)xSemaphoreGive(s_stepper_motion_mutex);
    return err;
}

void platform_stepper_driver_power_off(void)
{
    if (!ensure_stepper_motion_mutex()) {
        return;
    }

    TickType_t lock_ticks = pdMS_TO_TICKS(PLATFORM_STEPPER_POWER_OFF_WAIT_MS);
    if (lock_ticks == 0) {
        lock_ticks = 1;
    }
    if (xSemaphoreTake(s_stepper_motion_mutex, lock_ticks) != pdTRUE) {
        return;
    }

    portENTER_CRITICAL(&s_stepper_mux);
    stepper_release_locked();
    portEXIT_CRITICAL(&s_stepper_mux);

    (void)xSemaphoreGive(s_stepper_motion_mutex);
}