/*
 * Developed by TComLab
 * Version: v0.1
 * Date: 2025-12-15
 */

#include "uart_mux.h"

#include <string.h>
#include <limits.h>

#include "dcdc_controller.h"
#include "driver/gpio.h"
#include "driver/uart.h"
#include "esp_check.h"
#include "esp_log.h"
#include "esp_timer.h"
#include "freertos/FreeRTOS.h"
#include "freertos/semphr.h"
#include "freertos/task.h"
#include "sdkconfig.h"
#include "ws2812_status.h"

#ifndef CONFIG_WATCH_UART_MUX_CHANNELS
#define CONFIG_WATCH_UART_MUX_CHANNELS 5
#endif
#ifndef CONFIG_WATCH_UART_HEARTBEAT_TIMEOUT_SEC
#define CONFIG_WATCH_UART_HEARTBEAT_TIMEOUT_SEC 60
#endif

#if CONFIG_WATCH_UART_MUX_ENABLED

#define UART_MUX_MAX_CHANNELS CONFIG_WATCH_UART_MUX_CHANNELS

static const char *TAG = "uart_mux";

static const gpio_num_t s_select_pins[3] = {
    CONFIG_WATCH_UART_MUX_SEL_A0,
    CONFIG_WATCH_UART_MUX_SEL_A1,
    CONFIG_WATCH_UART_MUX_SEL_A2,
};

static SemaphoreHandle_t s_mutex;
static size_t s_active_channel = SIZE_MAX;
static bool s_initialized;
static int64_t s_last_heartbeat_us[UART_MUX_MAX_CHANNELS];
static TaskHandle_t s_watchdog_task;
static uint8_t s_consecutive_miss[UART_MUX_MAX_CHANNELS];

// Перемикає апаратний мультиплексор на вказаний канал під захистом мьютекса,
// оновлюючи таймстемп останнього heartbeat для контролю watchdog.
static esp_err_t uart_mux_select_locked(size_t channel)
{
    if (channel >= UART_MUX_MAX_CHANNELS) {
        return ESP_ERR_INVALID_ARG;
    }
    if (channel == s_active_channel) {
        return ESP_OK;
    }
    for (int bit = 0; bit < 3; ++bit) {
        int level = (channel >> bit) & 0x1;
        ESP_RETURN_ON_ERROR(gpio_set_level(s_select_pins[bit], level), TAG,
                            "GPIO set failed");
    }
    s_active_channel = channel;
    s_last_heartbeat_us[channel] = esp_timer_get_time();
    return ESP_OK;
}

// Періодично опитує всі канали, щоб зчитати heartbeat та перезапускає DCDC,
// якщо канал «мовчить» довше за CONFIG_WATCH_UART_HEARTBEAT_TIMEOUT_SEC.
static void uart_mux_watchdog_task(void *arg)
{
    const TickType_t poll_interval = pdMS_TO_TICKS(10000);
    const TickType_t read_timeout = pdMS_TO_TICKS(10);
    const int64_t timeout_us = (int64_t)CONFIG_WATCH_UART_HEARTBEAT_TIMEOUT_SEC * 1000000LL;
    uint8_t buffer[CONFIG_WATCH_UART_MUX_DEFAULT_READ_LEN];

    while (true) {
        vTaskDelay(poll_interval);
        int64_t now = esp_timer_get_time();
        for (size_t ch = 0; ch < UART_MUX_MAX_CHANNELS; ++ch) {
            if (xSemaphoreTake(s_mutex, pdMS_TO_TICKS(20)) == pdTRUE) {
                ws2812_status_indicate_polling(ch, 2000);
                if (uart_mux_select_locked(ch) == ESP_OK) {
                    int read = uart_read_bytes(CONFIG_WATCH_UART_PORT,
                                               buffer,
                                               sizeof(buffer),
                                               read_timeout);
                    if (read > 0) {
                        s_last_heartbeat_us[ch] = now;
                        s_consecutive_miss[ch] = 0;
                        ws2812_status_set_ack_state(ch, true);
                    } else if (s_consecutive_miss[ch] < UINT8_MAX) {
                        s_consecutive_miss[ch]++;
                        ws2812_status_set_ack_state(ch, false);
                    }
                }
                xSemaphoreGive(s_mutex);
            }

            if (dcdc_get_state(ch) && s_consecutive_miss[ch] >= 3) {
                ESP_LOGW(TAG, "CH%u: немає відповіді 3 рази поспіль, перезапуск живлення", (unsigned)ch);
                ws2812_status_set_ack_state(ch, false);
                dcdc_disable(ch);
                vTaskDelay(pdMS_TO_TICKS(2000));
                dcdc_enable(ch);
                s_consecutive_miss[ch] = 0;
                s_last_heartbeat_us[ch] = esp_timer_get_time();
            }

            if (dcdc_get_state(ch) && s_last_heartbeat_us[ch] > 0 &&
                (now - s_last_heartbeat_us[ch]) > timeout_us) {
                ESP_LOGW(TAG, "Heartbeat каналу %u втрачено, перезавантаження...", (unsigned)ch);
                ws2812_status_set_ack_state(ch, false);
                dcdc_disable(ch);
                vTaskDelay(pdMS_TO_TICKS(3000));
                dcdc_enable(ch);
                s_last_heartbeat_us[ch] = esp_timer_get_time();
            }
        }
    }
}

#endif // CONFIG_WATCH_UART_MUX_ENABLED

// Налаштовує GPIO-вибірники, драйвер UART та створює watchdog-задачу для
// мультиплексора; повторний виклик просто повертає ESP_OK.
esp_err_t uart_mux_init(void)
{
#if !CONFIG_WATCH_UART_MUX_ENABLED
    return ESP_ERR_NOT_SUPPORTED;
#else
    if (s_initialized) {
        return ESP_OK;
    }

    gpio_config_t io_conf = {
        .intr_type = GPIO_INTR_DISABLE,
        .mode = GPIO_MODE_OUTPUT,
        .pull_down_en = GPIO_PULLDOWN_DISABLE,
        .pull_up_en = GPIO_PULLUP_DISABLE,
        .pin_bit_mask = 0,
    };
    for (int i = 0; i < 3; ++i) {
        io_conf.pin_bit_mask = 1ULL << s_select_pins[i];
        ESP_RETURN_ON_ERROR(gpio_config(&io_conf), TAG, "GPIO config failed");
        ESP_RETURN_ON_ERROR(gpio_set_level(s_select_pins[i], 0), TAG, "GPIO init level failed");
    }
    s_active_channel = 0;

    uart_config_t uart_cfg = {
        .baud_rate = CONFIG_WATCH_UART_BAUD,
        .data_bits = UART_DATA_8_BITS,
        .parity = UART_PARITY_DISABLE,
        .stop_bits = UART_STOP_BITS_1,
        .flow_ctrl = UART_HW_FLOWCTRL_DISABLE,
        .source_clk = UART_SCLK_APB,
    };
    ESP_RETURN_ON_ERROR(uart_param_config(CONFIG_WATCH_UART_PORT, &uart_cfg), TAG, "UART config failed");
    ESP_RETURN_ON_ERROR(uart_set_pin(CONFIG_WATCH_UART_PORT,
                                     CONFIG_WATCH_UART_TX_GPIO,
                                     CONFIG_WATCH_UART_RX_GPIO,
                                     UART_PIN_NO_CHANGE,
                                     UART_PIN_NO_CHANGE),
                        TAG, "UART pin assign failed");
    ESP_RETURN_ON_ERROR(uart_driver_install(CONFIG_WATCH_UART_PORT,
                                            CONFIG_WATCH_UART_MUX_DEFAULT_READ_LEN * 2,
                                            CONFIG_WATCH_UART_MUX_DEFAULT_READ_LEN * 2,
                                            0, NULL, 0),
                        TAG, "UART driver install failed");

    s_mutex = xSemaphoreCreateMutex();
    ESP_RETURN_ON_FALSE(s_mutex, ESP_ERR_NO_MEM, TAG, "mutex alloc failed");

    int64_t now = esp_timer_get_time();
    for (size_t ch = 0; ch < UART_MUX_MAX_CHANNELS; ++ch) {
        s_last_heartbeat_us[ch] = now;
        s_consecutive_miss[ch] = 0;
    }

    if (xTaskCreate(uart_mux_watchdog_task, "uart_mux_wd", 4096, NULL, 5, &s_watchdog_task) != pdPASS) {
        return ESP_ERR_NO_MEM;
    }

    s_initialized = true;
    ESP_LOGI(TAG, "UART мультиплексор активовано, каналів: %d", UART_MUX_MAX_CHANNELS);
    return ESP_OK;
#endif
}

// Повертає ознаку ініціалізації модулю UART мультиплексора.
bool uart_mux_ready(void)
{
#if CONFIG_WATCH_UART_MUX_ENABLED
    return s_initialized;
#else
    return false;
#endif
}

// Кількість доступних каналів, визначених у конфігурації.
size_t uart_mux_channel_count(void)
{
    return CONFIG_WATCH_UART_MUX_CHANNELS;
}

// Перемикається на канал, передає буфер даних через загальний UART та
// захищає доступ до шини мьютексом, щоб уникнути гонок між задачами.
esp_err_t uart_mux_write(size_t channel, const uint8_t *data, size_t length, TickType_t timeout)
{
#if !CONFIG_WATCH_UART_MUX_ENABLED
    return ESP_ERR_NOT_SUPPORTED;
#else
    if (!s_initialized) {
        return ESP_ERR_INVALID_STATE;
    }
    if (channel >= UART_MUX_MAX_CHANNELS) {
        return ESP_ERR_INVALID_ARG;
    }
    if (!data || length == 0) {
        return ESP_OK;
    }
    if (xSemaphoreTake(s_mutex, timeout) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }
    esp_err_t err = uart_mux_select_locked(channel);
    if (err == ESP_OK) {
        int written = uart_write_bytes(CONFIG_WATCH_UART_PORT, (const char *)data, length);
        if (written < 0 || (size_t)written != length) {
            err = ESP_FAIL;
        } else {
            if (uart_wait_tx_done(CONFIG_WATCH_UART_PORT, timeout) != ESP_OK) {
                err = ESP_ERR_TIMEOUT;
            }
        }
    }
    xSemaphoreGive(s_mutex);
    return err;
#endif
}

// Читає дані з вказаного каналу, оновлюючи heartbeat під час успішного
// зчитування, і повертає кількість байтів через out_length.
esp_err_t uart_mux_read(size_t channel, uint8_t *buffer, size_t buffer_size, size_t *out_length, TickType_t timeout)
{
#if !CONFIG_WATCH_UART_MUX_ENABLED
    return ESP_ERR_NOT_SUPPORTED;
#else
    if (!s_initialized) {
        return ESP_ERR_INVALID_STATE;
    }
    if (channel >= UART_MUX_MAX_CHANNELS || !buffer || buffer_size == 0) {
        return ESP_ERR_INVALID_ARG;
    }
    if (xSemaphoreTake(s_mutex, timeout) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }
    esp_err_t err = uart_mux_select_locked(channel);
    if (err == ESP_OK) {
        int read = uart_read_bytes(CONFIG_WATCH_UART_PORT, buffer, buffer_size, timeout);
        if (read < 0) {
            err = ESP_FAIL;
        } else {
            if (out_length) {
                *out_length = (size_t)read;
            }
            if (read > 0) {
                s_last_heartbeat_us[channel] = esp_timer_get_time();
            }
        }
    }
    xSemaphoreGive(s_mutex);
    return err;
#endif
}