#include <driver/uart.h>
#include <driver/gpio.h>
#include <esp_log.h>
#include <esp_timer.h>
#include <esp_random.h>

#include <freertos/FreeRTOS.h>
#include <freertos/task.h>

#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

static const char* TAG = "test";

static const gpio_num_t tx = GPIO_NUM_11;
static const gpio_num_t rx = GPIO_NUM_10;
static const gpio_num_t rts = GPIO_NUM_8;
static const gpio_num_t cts = GPIO_NUM_9;

static const uart_port_t uart_num = UART_NUM_2;

void config(void) {
    uart_config_t uart_config = {
        .baud_rate = 115200,
        .data_bits = UART_DATA_8_BITS,
        .parity = UART_PARITY_DISABLE,
        .stop_bits = UART_STOP_BITS_1,
        .flow_ctrl = UART_HW_FLOWCTRL_DISABLE,
    };

    ESP_ERROR_CHECK(uart_driver_install(uart_num, 1024 /* rx_buf */, 0 /* tx_buf */, 0 /* EQ size */, NULL /* EQ handle */, 0 /* intr_flags */));

    ESP_ERROR_CHECK(uart_param_config(uart_num, &uart_config));

    ESP_ERROR_CHECK(uart_set_pin(uart_num, tx, rx, rts, cts));

    ESP_ERROR_CHECK(uart_set_mode(uart_num, UART_MODE_RS485_COLLISION_DETECT));

    gpio_reset_pin(rts);
    gpio_set_direction(rts, GPIO_MODE_OUTPUT);

    gpio_reset_pin(cts);
    gpio_set_direction(cts, GPIO_MODE_OUTPUT);

    // Enable reading
    gpio_set_level(cts, 0);
}

ssize_t rs485_write(uart_port_t uart, void* buffer, size_t size) {
    // Enable write
    ESP_ERROR_CHECK(gpio_set_level(rts, 1));
    
    ssize_t written = uart_write_bytes(uart, buffer, size);
    ESP_ERROR_CHECK(uart_wait_tx_done(uart_num, 1000));

    // Disable write
    ESP_ERROR_CHECK(gpio_set_level(rts, 0));

    return written;
}

bool rs485_collision(uart_port_t uart) {
    bool ret = false;

    ESP_ERROR_CHECK(uart_get_collision_flag(uart, &ret));

    return ret;
}

typedef enum {
    PROTOCOL_IDLE,
    PROTOCOL_PENDING_WRITE,
} protocol_state_t;

typedef enum : uint64_t {
    PROTOCOL_PRIORITY_0 = 1000,
    PROTOCOL_PRIORITY_1 = 3000,
    PROTOCOL_PRIORITY_2 = 10000,
} protocol_priority_t;

typedef struct {
    protocol_state_t state;
    uart_port_t uart;

    uint8_t* buffer;
    size_t buffer_capacity;

    size_t size;
    protocol_priority_t priority;
    uint64_t timeout_delta;
    uint64_t last_activity;
} protocol_context_t;

typedef protocol_context_t* protocol_handle_t;

uint32_t protocol_random_timeout() {
    return esp_random() & 511;
}

esp_err_t protocol_create(protocol_handle_t* out_handle, size_t buffer_size, uart_port_t uart) {
    protocol_handle_t ctx = (protocol_handle_t) calloc(1, sizeof(protocol_context_t));
    if (ctx == NULL)
        return ESP_ERR_NO_MEM;

    *ctx = (protocol_context_t) {
        .state = PROTOCOL_IDLE,
        .uart = uart,
        .buffer = NULL,
        .buffer_capacity = buffer_size,
        .size = 0,
        .priority = PROTOCOL_PRIORITY_0,
        .timeout_delta = 0,
    };

    ctx->buffer = (uint8_t*) calloc(buffer_size, 1);
    if (ctx->buffer == NULL) {
        free(ctx);
        return ESP_ERR_NO_MEM;
    }

    *out_handle = ctx;

    return ESP_OK;
}

esp_err_t protocol_destroy(protocol_handle_t handle) {
    if (handle->state != PROTOCOL_IDLE)
        return ESP_ERR_INVALID_ARG;

    free(handle->buffer);
    free(handle);
    return ESP_OK;
}

esp_err_t protocol_write(protocol_handle_t handle, protocol_priority_t priority, uint8_t* buffer, size_t size) {
    if (handle->state != PROTOCOL_IDLE)
        return ESP_ERR_INVALID_STATE;

    if (handle->buffer_capacity < size)
        return ESP_ERR_INVALID_ARG;

    handle->size = size;
    memcpy(handle->buffer, buffer, size);

    handle->state = PROTOCOL_PENDING_WRITE;
    handle->timeout_delta = priority + protocol_random_timeout();

    return ESP_OK;
}

size_t protocol_tick(protocol_handle_t handle, uint8_t* buffer, size_t buffer_capacity) {
    uint64_t now = esp_timer_get_time();
    int read = uart_read_bytes(handle->uart, buffer, buffer_capacity, 0);

    assert(read >= 0);

    if (read > 0) {
        handle->last_activity = esp_timer_get_time();
        return read;
    }
    
    if (handle->state == PROTOCOL_IDLE)
        return 0;

    assert(handle->state == PROTOCOL_PENDING_WRITE);

    if (handle->last_activity + handle->timeout_delta > now)
        return 0;

    rs485_write(handle->uart, handle->buffer, handle->size);
    
    if (rs485_collision(uart_num)) {
        handle->timeout_delta = handle->priority + protocol_random_timeout();
        ESP_LOGW(TAG, "Collision occured");
    }

    handle->state = PROTOCOL_IDLE;

    return 0;
}

bool protocol_write_finished(protocol_handle_t handle) {
    return handle->state == PROTOCOL_IDLE;
}

void app_main(void)
{
    uint8_t* data = malloc(128);
    uint8_t* in_data = malloc(1024);

    for (size_t i = 0; i < 26; i++) {
        data[i] = 'a' + i % 26;
    }

    config();

    protocol_handle_t protocol;

    ESP_ERROR_CHECK(protocol_create(&protocol, 1024, uart_num));

    int counter = 0;

    while (true) {
        size_t read = protocol_tick(protocol, in_data, 1024);
        if (read) {
            in_data[read] = 0;
            ESP_LOGI(TAG, "Read %u bytes: %s", read, in_data);
        }

        if (protocol_write_finished(protocol))
            ESP_ERROR_CHECK(protocol_write(protocol, PROTOCOL_PRIORITY_0, data, 26));

        vTaskDelay(2);
    }
}