#include "trikke_ble_transport.h"

#include <string.h>

#include "host/ble_att.h"
#include "host/ble_gap.h"
#include "host/ble_gatt.h"
#include "host/ble_hs.h"
#include "host/ble_uuid.h"
#include "host/util/util.h"
#include "esp_timer.h"
#include "nimble/nimble_port.h"
#include "nimble/nimble_port_freertos.h"
#include "nvs_flash.h"
#include "os/os_mbuf.h"
#include "services/gap/ble_svc_gap.h"
#include "services/gatt/ble_svc_gatt.h"
#include "trikke_ble_protocol.h"
#include "trikke_protocol.h"

#define TRIKKE_BLE_DEVICE_NAME "TrikkeSensor"
#define TRIKKE_BLE_MAX_ATT_PAYLOAD 253
#define TRIKKE_BLE_FRAGMENTS_PER_POLL 8
#define TRIKKE_BLE_ACK_TIMEOUT_US      1000000

static trikke_ble_transport_t *s_ble;
static uint8_t s_own_address_type;
static uint16_t s_data_value_handle;
static uint16_t s_ack_value_handle;

// 7d2ea000-f75b-4a9b-8fbe-3d4c2a1e9c10 and adjacent characteristic UUIDs.
static const ble_uuid128_t TRIKKE_SERVICE_UUID =
    BLE_UUID128_INIT(0x10, 0x9c, 0x1e, 0x2a, 0x4c, 0x3d, 0xbe, 0x8f,
                     0x9b, 0x4a, 0x5b, 0xf7, 0x00, 0xa0, 0x2e, 0x7d);
static const ble_uuid128_t TRIKKE_DATA_UUID =
    BLE_UUID128_INIT(0x11, 0x9c, 0x1e, 0x2a, 0x4c, 0x3d, 0xbe, 0x8f,
                     0x9b, 0x4a, 0x5b, 0xf7, 0x00, 0xa0, 0x2e, 0x7d);
static const ble_uuid128_t TRIKKE_ACK_UUID =
    BLE_UUID128_INIT(0x12, 0x9c, 0x1e, 0x2a, 0x4c, 0x3d, 0xbe, 0x8f,
                     0x9b, 0x4a, 0x5b, 0xf7, 0x00, 0xa0, 0x2e, 0x7d);

static uint32_t get_u32_le(const uint8_t *input)
{
    return (uint32_t)input[0] | ((uint32_t)input[1] << 8) |
           ((uint32_t)input[2] << 16) | ((uint32_t)input[3] << 24);
}

static int data_access(
    uint16_t connection_handle,
    uint16_t attribute_handle,
    struct ble_gatt_access_ctxt *context,
    void *argument)
{
    (void)connection_handle;
    (void)attribute_handle;
    (void)context;
    (void)argument;
    // NimBLE requires every characteristic definition to have an access
    // callback, even though this notify-only value is never client-accessible.
    return BLE_ATT_ERR_UNLIKELY;
}

static int ack_access(
    uint16_t connection_handle,
    uint16_t attribute_handle,
    struct ble_gatt_access_ctxt *context,
    void *argument)
{
    (void)attribute_handle;
    (void)argument;
    trikke_ble_transport_t *ble = s_ble;
    if (ble == NULL || context->op != BLE_GATT_ACCESS_OP_WRITE_CHR) {
        return BLE_ATT_ERR_UNLIKELY;
    }

    uint8_t ack[TRIKKE_BLE_ACK_SIZE] = {0};
    uint16_t ack_size = 0;
    if (OS_MBUF_PKTLEN(context->om) != TRIKKE_BLE_ACK_SIZE ||
        ble_hs_mbuf_to_flat(context->om, ack, sizeof(ack), &ack_size) != 0) {
        portENTER_CRITICAL(&ble->lock);
        ++ble->counters.invalid_ack_count;
        portEXIT_CRITICAL(&ble->lock);
        return BLE_ATT_ERR_INVALID_ATTR_VALUE_LEN;
    }

    uint32_t acknowledged_sequence = 0;
    const bool valid_shape =
        trikke_ble_decode_ack(ack, ack_size, &acknowledged_sequence);
    portENTER_CRITICAL(&ble->lock);
    const bool accepted = valid_shape && ble->frame_active &&
        ble->frame_fully_sent_once &&
        ble->connected && ble->subscribed &&
        ble->connection_handle == connection_handle &&
        ble->frame_epoch == ble->delivery_epoch &&
        acknowledged_sequence == ble->frame_sequence;
    if (accepted) {
        ble->ack_received = true;
    } else {
        ++ble->counters.invalid_ack_count;
    }
    portEXIT_CRITICAL(&ble->lock);
    return accepted ? 0 : BLE_ATT_ERR_UNLIKELY;
}

static const struct ble_gatt_svc_def TRIKKE_GATT_SERVICES[] = {
    {
        .type = BLE_GATT_SVC_TYPE_PRIMARY,
        .uuid = &TRIKKE_SERVICE_UUID.u,
        .characteristics = (struct ble_gatt_chr_def[]) {
            {
                .uuid = &TRIKKE_DATA_UUID.u,
                .access_cb = data_access,
                .flags = BLE_GATT_CHR_F_NOTIFY,
                .val_handle = &s_data_value_handle,
            },
            {
                .uuid = &TRIKKE_ACK_UUID.u,
                .access_cb = ack_access,
                .flags = BLE_GATT_CHR_F_WRITE | BLE_GATT_CHR_F_WRITE_NO_RSP,
                .val_handle = &s_ack_value_handle,
            },
            {0},
        },
    },
    {0},
};

static int gap_event(struct ble_gap_event *event, void *argument);

static int advertise(void)
{
    struct ble_hs_adv_fields fields = {0};
    fields.flags = BLE_HS_ADV_F_DISC_GEN | BLE_HS_ADV_F_BREDR_UNSUP;
    fields.uuids128 = (ble_uuid128_t *)&TRIKKE_SERVICE_UUID;
    fields.num_uuids128 = 1;
    fields.uuids128_is_complete = 1;
    int result = ble_gap_adv_set_fields(&fields);
    if (result != 0) {
        return result;
    }

    const char *name = ble_svc_gap_device_name();
    struct ble_hs_adv_fields response = {0};
    response.name = (uint8_t *)name;
    response.name_len = strlen(name);
    response.name_is_complete = 1;
    result = ble_gap_adv_rsp_set_fields(&response);
    if (result != 0) {
        return result;
    }

    const struct ble_gap_adv_params parameters = {
        .conn_mode = BLE_GAP_CONN_MODE_UND,
        .disc_mode = BLE_GAP_DISC_MODE_GEN,
    };
    return ble_gap_adv_start(s_own_address_type, NULL, BLE_HS_FOREVER,
                             &parameters, gap_event, NULL);
}

static void on_reset(int reason)
{
    (void)reason;
    trikke_ble_transport_t *ble = s_ble;
    if (ble == NULL) {
        return;
    }
    portENTER_CRITICAL(&ble->lock);
    if (ble->connected) {
        ++ble->counters.disconnect_count;
    }
    ble->connected = false;
    ble->subscribed = false;
    ble->connection_handle = BLE_HS_CONN_HANDLE_NONE;
    ++ble->delivery_epoch;
    portEXIT_CRITICAL(&ble->lock);
}

static void on_sync(void)
{
    if (ble_hs_util_ensure_addr(0) != 0 ||
        ble_hs_id_infer_auto(0, &s_own_address_type) != 0) {
        return;
    }
    (void)advertise();
}

static int gap_event(struct ble_gap_event *event, void *argument)
{
    (void)argument;
    trikke_ble_transport_t *ble = s_ble;
    if (ble == NULL) {
        return 0;
    }

    switch (event->type) {
    case BLE_GAP_EVENT_CONNECT:
        if (event->connect.status == 0) {
            portENTER_CRITICAL(&ble->lock);
            ble->connected = true;
            ble->subscribed = false;
            ble->connection_handle = event->connect.conn_handle;
            ++ble->delivery_epoch;
            portEXIT_CRITICAL(&ble->lock);
        } else {
            (void)advertise();
        }
        return 0;

    case BLE_GAP_EVENT_DISCONNECT:
        portENTER_CRITICAL(&ble->lock);
        if (ble->connected) {
            ++ble->counters.disconnect_count;
        }
        ble->connected = false;
        ble->subscribed = false;
        ble->connection_handle = BLE_HS_CONN_HANDLE_NONE;
        ++ble->delivery_epoch;
        portEXIT_CRITICAL(&ble->lock);
        (void)advertise();
        return 0;

    case BLE_GAP_EVENT_SUBSCRIBE:
        if (event->subscribe.attr_handle == s_data_value_handle) {
            portENTER_CRITICAL(&ble->lock);
            const bool subscribed = event->subscribe.cur_notify != 0;
            if (subscribed != ble->subscribed) {
                ble->subscribed = subscribed;
                ++ble->delivery_epoch;
            }
            portEXIT_CRITICAL(&ble->lock);
        }
        return 0;

    case BLE_GAP_EVENT_ADV_COMPLETE:
        (void)advertise();
        return 0;

    default:
        return 0;
    }
}

static void host_task(void *argument)
{
    (void)argument;
    nimble_port_run();
    nimble_port_freertos_deinit();
}

static trikke_transport_status_t ble_begin_packet(
    void *context,
    const uint8_t *packet,
    size_t packet_size)
{
    trikke_ble_transport_t *ble = context;
    if (ble == NULL || packet == NULL ||
        packet_size < TRIKKE_WIRE_HEADER_SIZE ||
        packet_size > TRIKKE_WIRE_MAX_PACKET_SIZE) {
        return TRIKKE_TRANSPORT_FATAL;
    }

    portENTER_CRITICAL(&ble->lock);
    if (!ble->initialized || ble->frame_active) {
        portEXIT_CRITICAL(&ble->lock);
        return TRIKKE_TRANSPORT_FATAL;
    }
    if (!ble->connected || !ble->subscribed) {
        portEXIT_CRITICAL(&ble->lock);
        return TRIKKE_TRANSPORT_RETRY;
    }
    ble->frame_active = true;
    ble->frame_fully_sent_once = false;
    ble->ack_received = false;
    ble->frame = packet;
    ble->frame_size = packet_size;
    ble->next_offset = 0;
    ble->ack_deadline_us = 0;
    ble->frame_sequence = get_u32_le(packet + 12);
    ble->frame_epoch = ble->delivery_epoch;
    portEXIT_CRITICAL(&ble->lock);
    return TRIKKE_TRANSPORT_PENDING;
}

static trikke_transport_status_t ble_poll_once(void *context)
{
    trikke_ble_transport_t *ble = context;
    if (ble == NULL) {
        return TRIKKE_TRANSPORT_FATAL;
    }

    uint16_t connection_handle = BLE_HS_CONN_HANDLE_NONE;
    uint32_t delivery_epoch = 0;
    const uint8_t *frame = NULL;
    size_t frame_size = 0;
    size_t next_offset = 0;

    portENTER_CRITICAL(&ble->lock);
    if (!ble->initialized || !ble->frame_active) {
        portEXIT_CRITICAL(&ble->lock);
        return TRIKKE_TRANSPORT_FATAL;
    }
    if (ble->ack_received) {
        ble->frame_active = false;
        ble->frame_fully_sent_once = false;
        ble->ack_received = false;
        ble->frame = NULL;
        ble->frame_size = 0;
        ble->next_offset = 0;
        ble->ack_deadline_us = 0;
        portEXIT_CRITICAL(&ble->lock);
        return TRIKKE_TRANSPORT_COMPLETE;
    }
    if (!ble->connected || !ble->subscribed) {
        portEXIT_CRITICAL(&ble->lock);
        return TRIKKE_TRANSPORT_PENDING;
    }
    if (ble->frame_epoch != ble->delivery_epoch) {
        if (ble->next_offset != 0) {
            ++ble->counters.replay_count;
        }
        ble->next_offset = 0;
        ble->frame_epoch = ble->delivery_epoch;
        ble->ack_deadline_us = 0;
    }
    if (ble->next_offset == ble->frame_size) {
        if (esp_timer_get_time() >= ble->ack_deadline_us) {
            ble->next_offset = 0;
            ble->ack_deadline_us = 0;
            ++ble->counters.replay_count;
        } else {
            portEXIT_CRITICAL(&ble->lock);
            return TRIKKE_TRANSPORT_PENDING;
        }
    }
    connection_handle = ble->connection_handle;
    delivery_epoch = ble->delivery_epoch;
    frame = ble->frame;
    frame_size = ble->frame_size;
    next_offset = ble->next_offset;
    portEXIT_CRITICAL(&ble->lock);

    const uint16_t mtu = ble_att_mtu(connection_handle);
    if (mtu <= 3 + TRIKKE_BLE_FRAGMENT_HEADER_SIZE) {
        return TRIKKE_TRANSPORT_FATAL;
    }
    size_t att_payload_capacity = mtu - 3;
    if (att_payload_capacity > TRIKKE_BLE_MAX_ATT_PAYLOAD) {
        att_payload_capacity = TRIKKE_BLE_MAX_ATT_PAYLOAD;
    }
    uint8_t fragment[TRIKKE_BLE_MAX_ATT_PAYLOAD] = {0};
    const size_t fragment_size = trikke_ble_encode_fragment(
        fragment, sizeof(fragment), frame, frame_size, next_offset,
        att_payload_capacity);
    if (fragment_size == 0) {
        return TRIKKE_TRANSPORT_FATAL;
    }

    struct os_mbuf *notification =
        ble_hs_mbuf_from_flat(fragment, fragment_size);
    int result = BLE_HS_ENOMEM;
    if (notification != NULL) {
        result = ble_gatts_notify_custom(
            connection_handle, s_data_value_handle, notification);
    }
    if (result != 0) {
        portENTER_CRITICAL(&ble->lock);
        ++ble->counters.send_failure_count;
        portEXIT_CRITICAL(&ble->lock);
        if (result == BLE_HS_ENOMEM || result == BLE_HS_EBUSY ||
            result == BLE_HS_EAGAIN || result == BLE_HS_ENOTCONN) {
            return TRIKKE_TRANSPORT_PENDING;
        }
        return TRIKKE_TRANSPORT_FATAL;
    }

    const size_t sent_data_size =
        fragment_size - TRIKKE_BLE_FRAGMENT_HEADER_SIZE;
    portENTER_CRITICAL(&ble->lock);
    if (ble->frame_active && ble->connected && ble->subscribed &&
        ble->connection_handle == connection_handle &&
        ble->delivery_epoch == delivery_epoch &&
        ble->next_offset == next_offset) {
        ble->next_offset += sent_data_size;
        if (ble->next_offset == ble->frame_size) {
            ble->frame_fully_sent_once = true;
            ble->ack_deadline_us =
                esp_timer_get_time() + TRIKKE_BLE_ACK_TIMEOUT_US;
        }
    }
    portEXIT_CRITICAL(&ble->lock);
    return TRIKKE_TRANSPORT_PENDING;
}

static trikke_transport_status_t ble_poll_packet(void *context)
{
    trikke_ble_transport_t *ble = context;
    if (ble == NULL) {
        return TRIKKE_TRANSPORT_FATAL;
    }
    for (unsigned fragment = 0;
         fragment < TRIKKE_BLE_FRAGMENTS_PER_POLL;
         ++fragment) {
        portENTER_CRITICAL(&ble->lock);
        const size_t offset_before_poll = ble->next_offset;
        portEXIT_CRITICAL(&ble->lock);
        const trikke_transport_status_t status = ble_poll_once(context);
        if (status != TRIKKE_TRANSPORT_PENDING) {
            return status;
        }

        // Keep a small-MTU connection useful without turning one poll into an
        // unbounded loop. Stop as soon as the backend is waiting on either the
        // connection/subscription or the receiver's application ACK.
        portENTER_CRITICAL(&ble->lock);
        const bool waiting = !ble->connected || !ble->subscribed ||
            ble->next_offset == ble->frame_size;
        const bool made_progress = ble->next_offset != offset_before_poll;
        portEXIT_CRITICAL(&ble->lock);
        if (waiting || !made_progress) {
            return TRIKKE_TRANSPORT_PENDING;
        }
    }
    return TRIKKE_TRANSPORT_PENDING;
}

esp_err_t trikke_ble_transport_init(
    trikke_ble_transport_t *ble,
    trikke_transport_t *transport)
{
    if (ble == NULL || transport == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    if (s_ble != NULL || ble->initialized) {
        return ESP_ERR_INVALID_STATE;
    }

    memset(ble, 0, sizeof(*ble));
    ble->lock = (portMUX_TYPE)portMUX_INITIALIZER_UNLOCKED;
    ble->connection_handle = BLE_HS_CONN_HANDLE_NONE;
    s_ble = ble;

    esp_err_t error = nvs_flash_init();
    if (error == ESP_ERR_NVS_NO_FREE_PAGES ||
        error == ESP_ERR_NVS_NEW_VERSION_FOUND) {
        error = nvs_flash_erase();
        if (error == ESP_OK) {
            error = nvs_flash_init();
        }
    }
    if (error != ESP_OK) {
        s_ble = NULL;
        return error;
    }
    error = nimble_port_init();
    if (error != ESP_OK) {
        s_ble = NULL;
        return error;
    }

    ble_hs_cfg.reset_cb = on_reset;
    ble_hs_cfg.sync_cb = on_sync;
    ble_svc_gap_init();
    ble_svc_gatt_init();
    int result = ble_gatts_count_cfg(TRIKKE_GATT_SERVICES);
    if (result == 0) {
        result = ble_gatts_add_svcs(TRIKKE_GATT_SERVICES);
    }
    if (result == 0) {
        result = ble_svc_gap_device_name_set(TRIKKE_BLE_DEVICE_NAME);
    }
    if (result != 0) {
        (void)nimble_port_deinit();
        s_ble = NULL;
        return ESP_FAIL;
    }

    ble->initialized = true;
    transport->context = ble;
    transport->begin = ble_begin_packet;
    transport->poll = ble_poll_packet;
    nimble_port_freertos_init(host_task);
    return ESP_OK;
}

void trikke_ble_transport_get_counters(
    trikke_ble_transport_t *ble,
    trikke_ble_transport_counters_t *counters)
{
    if (ble == NULL || counters == NULL) {
        return;
    }
    portENTER_CRITICAL(&ble->lock);
    *counters = ble->counters;
    portEXIT_CRITICAL(&ble->lock);
}