#include "uart_trans.h"

#include <string.h>

#include "FreeRTOS.h"
#include "task.h"
#include "usart.h"

#include "app_runtime.h"
#include "config.h"
#include "debug_log.h"
#include "route_msg.h"

#define UART_MUX_SYNC  0x7Eu
#define UART_MUX_TAIL  0x7Fu

#define UART_NOTIFY_RX_U0   (1UL << 0)
#define UART_NOTIFY_RX_U1   (1UL << 1)
#define UART_NOTIFY_TX_U0   (1UL << 8)
#define UART_NOTIFY_TX_U1   (1UL << 9)

typedef struct {
    UART_HandleTypeDef *huart;
    uint8_t rx_dma_buffer[UART_RX_DMA_BUFFER_SIZE];
    uint8_t tx_dma_buffer[UART_TX_DMA_BUFFER_SIZE];
    uint8_t rx_ring[UART_RX_RING_BUFFER_SIZE];
    uint8_t tx_ring[UART_TX_RING_BUFFER_SIZE];
    volatile uint16_t rx_dma_read_index;
    volatile uint16_t rx_head;
    volatile uint16_t rx_tail;
    volatile uint16_t tx_head;
    volatile uint16_t tx_tail;
    volatile uint16_t tx_dma_len;
    volatile uint8_t tx_busy;
} uart_channel_ctx_t;

static uart_channel_ctx_t g_channels[UART_CHANNEL_MAX];

static uint16_t ring_used(uint16_t head, uint16_t tail, uint16_t size)
{
    return (head >= tail) ? (head - tail) : (size - tail + head);
}

static uint16_t ring_free(uint16_t head, uint16_t tail, uint16_t size)
{
    return (uint16_t)(size - ring_used(head, tail, size) - 1u);
}

static void process_rx_snapshot(uart_channel_t channel)
{
    uart_channel_ctx_t *ctx = &g_channels[channel];
    uint16_t dma_write_index = (uint16_t)(UART_RX_DMA_BUFFER_SIZE - __HAL_DMA_GET_COUNTER(ctx->huart->hdmarx));

    if (dma_write_index >= UART_RX_DMA_BUFFER_SIZE) {
        dma_write_index = 0u;
    }

    while (ctx->rx_dma_read_index != dma_write_index) {
        uint16_t next_head = (uint16_t)((ctx->rx_head + 1u) % UART_RX_RING_BUFFER_SIZE);
        if (next_head == ctx->rx_tail) {
            break;
        }
        ctx->rx_ring[ctx->rx_head] = ctx->rx_dma_buffer[ctx->rx_dma_read_index];
        ctx->rx_head = next_head;
        ctx->rx_dma_read_index = (uint16_t)((ctx->rx_dma_read_index + 1u) % UART_RX_DMA_BUFFER_SIZE);
    }
}

static void kick_tx(uart_channel_t channel)
{
    uart_channel_ctx_t *ctx = &g_channels[channel];
    uint16_t available;
    uint16_t chunk;
    uint16_t i;

    if (ctx->tx_busy != 0u) {
        return;
    }

    available = ring_used(ctx->tx_head, ctx->tx_tail, UART_TX_RING_BUFFER_SIZE);
    if (available == 0u) {
        return;
    }

    chunk = available;
    if (chunk > UART_TX_DMA_BUFFER_SIZE) {
        chunk = UART_TX_DMA_BUFFER_SIZE;
    }

    for (i = 0; i < chunk; ++i) {
        ctx->tx_dma_buffer[i] = ctx->tx_ring[ctx->tx_tail];
        ctx->tx_tail = (uint16_t)((ctx->tx_tail + 1u) % UART_TX_RING_BUFFER_SIZE);
    }

    ctx->tx_dma_len = chunk;
    ctx->tx_busy = 1u;
    if (HAL_UART_Transmit_DMA(ctx->huart, ctx->tx_dma_buffer, chunk) != HAL_OK) {
        ctx->tx_busy = 0u;
    }
}

static uint16_t uart_ring_available(uart_channel_t channel)
{
    return ring_used(g_channels[channel].rx_head, g_channels[channel].rx_tail, UART_RX_RING_BUFFER_SIZE);
}

static uint16_t uart_ring_read(uart_channel_t channel, uint8_t *data, uint16_t max_len)
{
    uart_channel_ctx_t *ctx = &g_channels[channel];
    uint16_t copied = 0u;

    while (copied < max_len && ctx->rx_tail != ctx->rx_head) {
        data[copied++] = ctx->rx_ring[ctx->rx_tail];
        ctx->rx_tail = (uint16_t)((ctx->rx_tail + 1u) % UART_RX_RING_BUFFER_SIZE);
    }
    return copied;
}

static bool uart_ring_peek_byte(uart_channel_t channel, uint16_t offset, uint8_t *data)
{
    uart_channel_ctx_t *ctx = &g_channels[channel];
    uint16_t available = ring_used(ctx->rx_head, ctx->rx_tail, UART_RX_RING_BUFFER_SIZE);

    if (data == NULL || offset >= available) {
        return false;
    }

    *data = ctx->rx_ring[(ctx->rx_tail + offset) % UART_RX_RING_BUFFER_SIZE];
    return true;
}

static void uart_ring_drop(uart_channel_t channel, uint16_t len)
{
    uart_channel_ctx_t *ctx = &g_channels[channel];

    ctx->rx_tail = (uint16_t)((ctx->rx_tail + len) % UART_RX_RING_BUFFER_SIZE);
}

static void uart_route_raw_channel(uart_channel_t channel)
{
    const device_config_t *cfg = config_get();
    uint8_t buffer[ROUTE_MSG_MAX_PAYLOAD];
    uint16_t len;
    uint8_t uart_endpoint = (channel == UART_CHANNEL_U1) ? ENDPOINT_UART3 : ENDPOINT_UART2;
    uint32_t i;

    len = uart_ring_read(channel, buffer, sizeof(buffer));
    if (len == 0u) {
        return;
    }

    for (i = 0; i < CONFIG_LINK_COUNT; ++i) {
        if (cfg->links[i].enabled == 0u || cfg->links[i].uart != ((channel == UART_CHANNEL_U1) ? LINK_UART_U1 : LINK_UART_U0)) {
            continue;
        }

        (void)route_send(xLinkTxQueues[i],
                         uart_endpoint,
                         config_link_index_to_endpoint((uint8_t)i),
                         (channel == UART_CHANNEL_U1) ? ROUTE_CONN_UART3 : ROUTE_CONN_UART2,
                         buffer,
                         len,
                         pdMS_TO_TICKS(10));
    }
}

static void uart_send_tcp_msg_to_uarts(route_msg_t *msg)
{
    uint8_t frame[ROUTE_MSG_MAX_PAYLOAD + 6u];
    uint16_t frame_len = 0u;

    if ((msg->dst_mask & ENDPOINT_UART2) != 0u) {
        if (config_get()->mux_mode == MUX_MODE_FRAME) {
            if (uart_mux_encode_frame(msg->src_id, ENDPOINT_UART2, msg->data, msg->len, frame, &frame_len, sizeof(frame))) {
                (void)uart_trans_send_buffer(UART_CHANNEL_U0, frame, frame_len);
            }
        } else {
            (void)uart_trans_send_buffer(UART_CHANNEL_U0, msg->data, msg->len);
        }
    }

    if ((msg->dst_mask & ENDPOINT_UART3) != 0u) {
        if (config_get()->mux_mode == MUX_MODE_FRAME) {
            if (uart_mux_encode_frame(msg->src_id, ENDPOINT_UART3, msg->data, msg->len, frame, &frame_len, sizeof(frame))) {
                (void)uart_trans_send_buffer(UART_CHANNEL_U1, frame, frame_len);
            }
        } else {
            (void)uart_trans_send_buffer(UART_CHANNEL_U1, msg->data, msg->len);
        }
    }
}

static void uart_route_mux_frame(uart_channel_t source_channel, const uart_mux_frame_t *frame)
{
    uint32_t i;
    uint8_t source_conn = (source_channel == UART_CHANNEL_U1) ? ROUTE_CONN_UART3 : ROUTE_CONN_UART2;
    uint8_t out_frame[ROUTE_MSG_MAX_PAYLOAD + 6u];
    uint16_t out_frame_len = 0u;

    if (frame->dst_mask == 0u) {
        (void)route_send(xConfigQueue,
                         frame->src_id,
                         0u,
                         source_conn,
                         frame->payload,
                         frame->payload_len,
                         pdMS_TO_TICKS(10));
        return;
    }

    for (i = 0; i < CONFIG_LINK_COUNT; ++i) {
        uint8_t endpoint = config_link_index_to_endpoint((uint8_t)i);
        if ((frame->dst_mask & endpoint) != 0u) {
            (void)route_send(xLinkTxQueues[i], frame->src_id, endpoint, source_conn, frame->payload, frame->payload_len, pdMS_TO_TICKS(10));
        }
    }

    if ((frame->dst_mask & ENDPOINT_UART2) != 0u && source_channel != UART_CHANNEL_U0) {
        if (uart_mux_encode_frame(frame->src_id, ENDPOINT_UART2, frame->payload, frame->payload_len, out_frame, &out_frame_len, sizeof(out_frame))) {
            (void)uart_trans_send_buffer(UART_CHANNEL_U0, out_frame, out_frame_len);
        }
    }

    if ((frame->dst_mask & ENDPOINT_UART3) != 0u && source_channel != UART_CHANNEL_U1) {
        if (uart_mux_encode_frame(frame->src_id, ENDPOINT_UART3, frame->payload, frame->payload_len, out_frame, &out_frame_len, sizeof(out_frame))) {
            (void)uart_trans_send_buffer(UART_CHANNEL_U1, out_frame, out_frame_len);
        }
    }
}

int uart_trans_init(void)
{
    memset(g_channels, 0, sizeof(g_channels));
    g_channels[UART_CHANNEL_U0].huart = &huart2;
    g_channels[UART_CHANNEL_U1].huart = &huart3;
    debug_log_printf("[UART] init u0=%p u1=%p\r\n", (void *)g_channels[UART_CHANNEL_U0].huart, (void *)g_channels[UART_CHANNEL_U1].huart);
    return 0;
}

int uart_trans_config(uint8_t uart_index, uint32_t baudrate)
{
    UART_HandleTypeDef *huart = (uart_index == LINK_UART_U1) ? &huart3 : &huart2;
    huart->Init.BaudRate = baudrate;
    return (HAL_UART_Init(huart) == HAL_OK) ? 0 : -1;
}

int uart_trans_start_all(void)
{
    uint32_t i;
    for (i = 0; i < UART_CHANNEL_MAX; ++i) {
        if (g_channels[i].huart == NULL) {
            debug_log_printf("[UART] start fail null handle ch=%lu\r\n", (unsigned long)i);
            return -1;
        }
        g_channels[i].rx_dma_read_index = 0u;
        g_channels[i].rx_head = 0u;
        g_channels[i].rx_tail = 0u;
        g_channels[i].tx_head = 0u;
        g_channels[i].tx_tail = 0u;
        g_channels[i].tx_dma_len = 0u;
        g_channels[i].tx_busy = 0u;
        __HAL_UART_ENABLE_IT(g_channels[i].huart, UART_IT_IDLE);
        if (HAL_UART_Receive_DMA(g_channels[i].huart, g_channels[i].rx_dma_buffer, UART_RX_DMA_BUFFER_SIZE) != HAL_OK) {
            debug_log_printf("[UART] dma start fail ch=%lu\r\n", (unsigned long)i);
            return -1;
        }
    }
    debug_log_write("[UART] rx dma started\r\n");
    return 0;
}

bool uart_trans_send_buffer(uart_channel_t channel, const uint8_t *data, uint16_t len)
{
    uart_channel_ctx_t *ctx = &g_channels[channel];
    uint16_t written = 0u;

    if (data == NULL || len == 0u) {
        return false;
    }

    while (written < len && ring_free(ctx->tx_head, ctx->tx_tail, UART_TX_RING_BUFFER_SIZE) > 0u) {
        ctx->tx_ring[ctx->tx_head] = data[written++];
        ctx->tx_head = (uint16_t)((ctx->tx_head + 1u) % UART_TX_RING_BUFFER_SIZE);
    }

    kick_tx(channel);
    return (written == len);
}

void uart_trans_notify_rx_from_isr(uart_channel_t channel, BaseType_t *xHigherPriorityTaskWoken)
{
    uint32_t notify = (channel == UART_CHANNEL_U1) ? UART_NOTIFY_RX_U1 : UART_NOTIFY_RX_U0;
    if (xUartRxTaskHandle != NULL) {
        xTaskNotifyFromISR(xUartRxTaskHandle, notify, eSetBits, xHigherPriorityTaskWoken);
    }
}

void uart_trans_tx_cplt_handler(uart_channel_t channel)
{
    uint32_t notify = (channel == UART_CHANNEL_U1) ? UART_NOTIFY_TX_U1 : UART_NOTIFY_TX_U0;
    BaseType_t xHigherPriorityTaskWoken = pdFALSE;
    if (xUartRxTaskHandle != NULL) {
        xTaskNotifyFromISR(xUartRxTaskHandle, notify, eSetBits, &xHigherPriorityTaskWoken);
        portYIELD_FROM_ISR(xHigherPriorityTaskWoken);
    }
}

bool uart_mux_try_extract_frame(uart_channel_t channel, uart_mux_frame_t *frame)
{
    uint16_t available;
    uint16_t payload_len;
    uint8_t sync_byte;
    uint16_t i;

    if (frame == NULL) {
        return false;
    }

    available = uart_ring_available(channel);
    if (available == 0u) {
        return false;
    }

    if (!uart_ring_peek_byte(channel, 0u, &sync_byte)) {
        return false;
    }

    if (sync_byte != UART_MUX_SYNC) {
        uart_ring_drop(channel, 1u);
        return false;
    }

    if (available < 6u) {
        return false;
    }

    if (!uart_ring_peek_byte(channel, 1u, &sync_byte)) {
        return false;
    }
    payload_len = (uint16_t)((uint16_t)sync_byte << 8);
    if (!uart_ring_peek_byte(channel, 2u, &sync_byte)) {
        return false;
    }
    payload_len = (uint16_t)(payload_len | sync_byte);
    if (payload_len > sizeof(frame->payload)) {
        uart_ring_drop(channel, 1u);
        return false;
    }

    if (available < (uint16_t)(payload_len + 6u)) {
        return false;
    }

    if (!uart_ring_peek_byte(channel, 5u + payload_len, &sync_byte)) {
        return false;
    }
    if (sync_byte != UART_MUX_TAIL) {
        uart_ring_drop(channel, 1u);
        return false;
    }

    if (!uart_ring_peek_byte(channel, 3u, &frame->src_id) ||
        !uart_ring_peek_byte(channel, 4u, &frame->dst_mask)) {
        return false;
    }

    frame->payload_len = payload_len;
    for (i = 0u; i < payload_len; ++i) {
        if (!uart_ring_peek_byte(channel, (uint16_t)(5u + i), &frame->payload[i])) {
            return false;
        }
    }

    uart_ring_drop(channel, (uint16_t)(payload_len + 6u));

    return true;
}

bool uart_mux_encode_frame(uint8_t src_id,
                           uint8_t dst_mask,
                           const uint8_t *payload,
                           uint16_t payload_len,
                           uint8_t *out,
                           uint16_t *out_len,
                           uint16_t out_capacity)
{
    uint16_t frame_len = (uint16_t)(payload_len + 6u);

    if (out == NULL || out_len == NULL || frame_len > out_capacity) {
        return false;
    }

    out[0] = UART_MUX_SYNC;
    out[1] = (uint8_t)(payload_len >> 8);
    out[2] = (uint8_t)(payload_len & 0xFFu);
    out[3] = src_id;
    out[4] = dst_mask;
    if (payload_len > 0u && payload != NULL) {
        memcpy(&out[5], payload, payload_len);
    }
    out[5 + payload_len] = UART_MUX_TAIL;
    *out_len = frame_len;
    return true;
}

void UartRxTask(void *argument)
{
    uint32_t notify_value;
    BaseType_t notified;
    route_msg_t *msg;
    uart_mux_frame_t frame;
    const device_config_t *cfg;

    (void)argument;
    if (uart_trans_start_all() != 0) {
        Debug_TrapWithRttHint("uart-start-fail");
        vTaskSuspend(NULL);
    }
    debug_log_boot("uart-task-started");

    for (;;) {
        notify_value = 0u;
        notified = xTaskNotifyWait(0u, 0xFFFFFFFFu, &notify_value, pdMS_TO_TICKS(10));

        if ((notified == pdTRUE) && ((notify_value & UART_NOTIFY_RX_U0) != 0u)) {
            process_rx_snapshot(UART_CHANNEL_U0);
        }
        if ((notified == pdTRUE) && ((notify_value & UART_NOTIFY_RX_U1) != 0u)) {
            process_rx_snapshot(UART_CHANNEL_U1);
        }
        if ((notified == pdTRUE) && ((notify_value & UART_NOTIFY_TX_U0) != 0u)) {
            g_channels[UART_CHANNEL_U0].tx_busy = 0u;
        }
        if ((notified == pdTRUE) && ((notify_value & UART_NOTIFY_TX_U1) != 0u)) {
            g_channels[UART_CHANNEL_U1].tx_busy = 0u;
        }

        while (xQueueReceive(xTcpRxQueue, &msg, 0) == pdPASS) {
            uart_send_tcp_msg_to_uarts(msg);
            route_msg_free(msg);
        }

        cfg = config_get();
        if (cfg->mux_mode == MUX_MODE_FRAME) {
            while (uart_mux_try_extract_frame(UART_CHANNEL_U0, &frame)) {
                uart_route_mux_frame(UART_CHANNEL_U0, &frame);
            }
            while (uart_mux_try_extract_frame(UART_CHANNEL_U1, &frame)) {
                uart_route_mux_frame(UART_CHANNEL_U1, &frame);
            }
        } else {
            uart_route_raw_channel(UART_CHANNEL_U0);
            uart_route_raw_channel(UART_CHANNEL_U1);
        }

        kick_tx(UART_CHANNEL_U0);
        kick_tx(UART_CHANNEL_U1);
    }
}