#include <string.h>

/* Private */
#include "nh_ctrl_api.h"
#include "protobuf-c/protobuf-c.h"

/* PB config */
#include "esp_hosted_config.pb-c.h"

#define NH_COMMAND_TIMEOUT_MSEC 30000
#define NH_COMMAND_REQ_MEM_SIZE 256

#define NH_CTRL_API_CREATE_ALLOCATOR(allocator_name, allocator_ctx) \
    ProtobufCAllocator allocator_name = {                           \
        .alloc          = nh_ctrl_api_protobuf_allocate,            \
        .free           = nh_ctrl_api_protobuf_free,                \
        .allocator_data = allocator_ctx,                            \
    }

#define NH_CTRL_API_TAKE_SEM(s_api, timeout)  s_api->osa->sem_take(s_api->osa->user_data, s_api->p_sem_req, timeout)
#define NH_CTRL_API_GIVE_SEM(s_api)           s_api->osa->sem_give(api->osa->user_data, api->p_sem_req)
#define NH_CTRL_API_ALLOC(s_api, ptr, p_size) s_api->osa->buf_allocate(s_api->osa->user_data, ptr, p_size)
#define NH_CTRL_API_FREE(s_api, ptr)          s_api->osa->buf_free(s_api->osa->user_data, ptr)

static nh_ret_t nh_ctrl_api_parse_event(nh_ctrl_api_t *api, CtrlMsgId *event_id, uint8_t *buf, uint32_t buf_len);
static void     nh_ctrl_api_dispatch_event(nh_ctrl_api_t *api, CtrlMsgId event_id);
static nh_ret_t nh_ctrl_api_general_request(nh_ctrl_api_t *api, CtrlMsg *tx_msg);
static nh_ret_t nh_ctrl_api_general_response(nh_ctrl_api_t *api, CtrlMsg **rx_msg, ProtobufCAllocator *allocator);
static void    *nh_ctrl_api_protobuf_allocate(void *handle, size_t size);
static void     nh_ctrl_api_protobuf_free(void *handle, void *data);

nh_ret_t nh_ctrl_api_init(nh_ctrl_api_t *api) {
    nh_ret_t ret = NH_RET_SUCCESS;

    ret = api->osa->sem_create(api->osa->user_data, &api->p_sem_req);
    if (ret != NH_RET_SUCCESS) {
        return ret;
    }

    ret = api->osa->sem_create(api->osa->user_data, &api->p_sem_resp);
    if (ret != NH_RET_SUCCESS) {
        goto free_sem_req;
    }

    ret = api->osa->sem_give(api->osa->user_data, api->p_sem_req);
    if (ret != NH_RET_SUCCESS) {
        goto free_sem_resp;
    }

    return ret;

free_sem_resp:
    api->osa->sem_destroy(api->osa->user_data, api->p_sem_resp);

free_sem_req:
    api->osa->sem_destroy(api->osa->user_data, api->p_sem_req);

    return ret;
}

nh_ret_t nh_ctrl_api_get_mac_address(nh_ctrl_api_t *api, uint8_t *mac_addr, nh_ctrl_wifi_mode_t mode) {
    nh_ret_t ret = NH_RET_SUCCESS;

    NH_CTRL_API_CREATE_ALLOCATOR(c_alloc, api);
    ret = NH_CTRL_API_TAKE_SEM(api, NH_COMMAND_TIMEOUT_MSEC);
    if (ret != NH_RET_SUCCESS) return ret;

    CtrlMsg tx_msg;
    ctrl_msg__init(&tx_msg);

    tx_msg.msg_id       = CTRL_MSG_ID__Req_GetMACAddress;
    tx_msg.payload_case = CTRL_MSG__PAYLOAD_REQ_GET_MAC_ADDRESS;

    ret = NH_CTRL_API_ALLOC(api, (uint8_t **)&tx_msg.req_get_mac_address, sizeof(CtrlMsgReqGetMacAddress));
    if (ret != NH_RET_SUCCESS) goto give_sem_exit;

    ctrl_msg__req__get_mac_address__init(tx_msg.req_get_mac_address);

    tx_msg.req_get_mac_address->mode = (mode == NH_CTRL_WIFI_MODE_AP) ? CTRL__WIFI_MODE__AP : CTRL__WIFI_MODE__STA;

    ret = nh_ctrl_api_general_request(api, &tx_msg);
    if (ret != NH_RET_SUCCESS) goto free_payload_exit;

    CtrlMsg *rx_msg;

    ret = nh_ctrl_api_general_response(api, &rx_msg, &c_alloc);
    if (ret != NH_RET_SUCCESS || rx_msg->msg_id != CTRL_MSG_ID__Resp_GetMACAddress) goto free_msg_exit;

    /* Sanity check */
    if (rx_msg->resp_get_mac_address->mac.len != NH_CTRL_API_MAC_LENGTH) goto free_msg_exit;

    memcpy(mac_addr, rx_msg->resp_get_mac_address->mac.data, rx_msg->resp_get_mac_address->mac.len);
    mac_addr[rx_msg->resp_get_mac_address->mac.len] = '\0';

free_msg_exit:
    ctrl_msg__free_unpacked(rx_msg, &c_alloc);

free_payload_exit:
    NH_CTRL_API_FREE(api, (uint8_t *)tx_msg.req_get_mac_address);

give_sem_exit:
    NH_CTRL_API_GIVE_SEM(api);

    return ret;
}

nh_ret_t nh_ctrl_api_get_ap_scan_list(nh_ctrl_api_t *api, nh_ctrl_api_ap_scan_list_cb_t cb) {
    nh_ret_t ret = NH_RET_SUCCESS;

    NH_CTRL_API_CREATE_ALLOCATOR(c_alloc, api);
    ret = NH_CTRL_API_TAKE_SEM(api, NH_COMMAND_TIMEOUT_MSEC);
    if (ret != NH_RET_SUCCESS) return ret;

    CtrlMsg tx_msg;
    ctrl_msg__init(&tx_msg);

    tx_msg.msg_id       = CTRL_MSG_ID__Req_GetAPScanList;
    tx_msg.payload_case = CTRL_MSG__PAYLOAD_REQ_SCAN_AP_LIST;

    ret = NH_CTRL_API_ALLOC(api, (uint8_t **)&tx_msg.req_scan_ap_list, sizeof(CtrlMsgReqScanResult));
    if (ret != NH_RET_SUCCESS) goto give_sem_exit;

    ctrl_msg__req__scan_result__init(tx_msg.req_scan_ap_list);

    ret = nh_ctrl_api_general_request(api, &tx_msg);
    if (ret != NH_RET_SUCCESS) goto free_payload_exit;

    CtrlMsg *rx_msg;

    ret = nh_ctrl_api_general_response(api, &rx_msg, &c_alloc);
    if (ret != NH_RET_SUCCESS || rx_msg->msg_id != CTRL_MSG_ID__Resp_GetAPScanList) goto free_msg_exit;

    for (uint32_t i = 0; i < rx_msg->resp_scan_ap_list->count; i++) {
        nh_ctrl_ap_scan_item_t item = {
            .rssi            = rx_msg->resp_scan_ap_list->entries[i]->rssi,
            .channel         = rx_msg->resp_scan_ap_list->entries[i]->chnl,
            .encryption_mode = (nh_ctrl_ap_enc_mode_t)rx_msg->resp_scan_ap_list->entries[i]->sec_prot,
        };

        uint32_t ssid_len  = rx_msg->resp_scan_ap_list->entries[i]->ssid.len;
        uint32_t bssid_len = rx_msg->resp_scan_ap_list->entries[i]->bssid.len;

        ret = NH_CTRL_API_ALLOC(api, (uint8_t **)&item.ssid, ssid_len + 1);
        if (ret != NH_RET_SUCCESS) goto free_msg_exit;

        ret = NH_CTRL_API_ALLOC(api, (uint8_t **)&item.bssid, bssid_len + 1);
        if (ret != NH_RET_SUCCESS) {
            NH_CTRL_API_FREE(api, (uint8_t *)item.ssid);

            goto free_msg_exit;
        }

        memcpy(item.ssid, rx_msg->resp_scan_ap_list->entries[i]->ssid.data, ssid_len);
        memcpy(item.bssid, rx_msg->resp_scan_ap_list->entries[i]->bssid.data, bssid_len);

        item.ssid[ssid_len]   = '\0';
        item.bssid[bssid_len] = '\0';

        if (cb) {
            cb(api->user_data, &item);
        }

        NH_CTRL_API_FREE(api, (uint8_t *)item.ssid);
        NH_CTRL_API_FREE(api, (uint8_t *)item.bssid);
    }

free_msg_exit:
    ctrl_msg__free_unpacked(rx_msg, &c_alloc);

free_payload_exit:
    NH_CTRL_API_FREE(api, (uint8_t *)tx_msg.req_get_mac_address);

give_sem_exit:
    NH_CTRL_API_GIVE_SEM(api);

    return ret;
}

nh_ret_t nh_ctrl_api_connect_ap(nh_ctrl_api_t *api, nh_ctrl_ap_conn_params_t *params,
                                nh_ctrl_ap_conn_result_t *result) {
    nh_ret_t ret = NH_RET_SUCCESS;

    NH_CTRL_API_CREATE_ALLOCATOR(c_alloc, api);
    ret = NH_CTRL_API_TAKE_SEM(api, NH_COMMAND_TIMEOUT_MSEC);
    if (ret != NH_RET_SUCCESS) return ret;

    CtrlMsg tx_msg;
    ctrl_msg__init(&tx_msg);

    tx_msg.msg_id       = CTRL_MSG_ID__Req_ConnectAP;
    tx_msg.payload_case = CTRL_MSG__PAYLOAD_REQ_CONNECT_AP;

    ret = NH_CTRL_API_ALLOC(api, (uint8_t **)&tx_msg.req_connect_ap, sizeof(CtrlMsgReqConnectAP));
    if (ret != NH_RET_SUCCESS) goto give_sem_exit;

    ctrl_msg__req__connect_ap__init(tx_msg.req_connect_ap);

    tx_msg.req_connect_ap->ssid              = params->ssid;
    tx_msg.req_connect_ap->bssid             = params->bssid;
    tx_msg.req_connect_ap->pwd               = params->password;
    tx_msg.req_connect_ap->is_wpa3_supported = params->wpa3_supported;
    tx_msg.req_connect_ap->listen_interval   = params->listen_interval;

    ret = nh_ctrl_api_general_request(api, &tx_msg);
    if (ret != NH_RET_SUCCESS) goto free_payload_exit;

    CtrlMsg *rx_msg;

    ret = nh_ctrl_api_general_response(api, &rx_msg, &c_alloc);
    if (ret != NH_RET_SUCCESS || rx_msg->msg_id != CTRL_MSG_ID__Resp_ConnectAP) goto free_msg_exit;

    uint32_t mac_len = rx_msg->resp_connect_ap->mac.len;
    if (mac_len != NH_CTRL_API_MAC_LENGTH) goto free_msg_exit;

    memcpy(result->mac_addr, rx_msg->resp_connect_ap->mac.data, mac_len);
    result->mac_addr[mac_len] = '\0';

    result->status = (nh_ctrl_conn_status_t)rx_msg->resp_connect_ap->resp;

free_msg_exit:
    ctrl_msg__free_unpacked(rx_msg, &c_alloc);

free_payload_exit:
    NH_CTRL_API_FREE(api, (uint8_t *)tx_msg.req_connect_ap);

give_sem_exit:
    NH_CTRL_API_GIVE_SEM(api);

    return ret;
}

void nh_ctrl_api_task(nh_ctrl_api_t *api) {
    nh_ret_t ret = NH_RET_SUCCESS;

    uint8_t            *payload;
    uint32_t            payload_size;
    nh_serial_ep_type_t ep_type;

    ret = nh_shared_if_serial_if_recv(api->shared_if, &ep_type, &payload, &payload_size, NH_COMMAND_TIMEOUT_MSEC);
    if (ret != NH_RET_SUCCESS) {
        return;
    }

    if (ep_type == NH_SERIAL_EP_RESP) {
        api->p_buf_rx     = payload;
        api->p_buf_rx_len = payload_size;

        ret = api->osa->sem_give(api->osa->user_data, api->p_sem_resp);
        if (ret != NH_RET_SUCCESS) {
            nh_shared_if_serial_if_free(api->shared_if, payload);
            return;
        }
    } else if (ep_type == NH_SERIAL_EP_EVENT) {
        CtrlMsgId event_id;
        ret = nh_ctrl_api_parse_event(api, &event_id, payload, payload_size);
        if (ret == NH_RET_SUCCESS) {
            nh_ctrl_api_dispatch_event(api, event_id);
        }

        nh_shared_if_serial_if_free(api->shared_if, payload);
    }
}

static nh_ret_t nh_ctrl_api_parse_event(nh_ctrl_api_t *api, CtrlMsgId *event_id, uint8_t *buf, uint32_t buf_len) {
    NH_CTRL_API_CREATE_ALLOCATOR(c_alloc, api);

    CtrlMsg *msg = ctrl_msg__unpack(&c_alloc, buf_len, buf);
    if (msg == NULL) {
        return NH_RET_FAIL;
    }

    *event_id = msg->msg_id;

    ctrl_msg__free_unpacked(msg, &c_alloc);
    return NH_RET_SUCCESS;
}

static void nh_ctrl_api_dispatch_event(nh_ctrl_api_t *api, CtrlMsgId event_id) {
    switch (event_id) {
        case CTRL_MSG_ID__Event_ESPInit: {
            if (api->cb.init) {
                api->cb.init(api->user_data);
            }

            break;
        }

        default:
            break;
    }
}

static nh_ret_t nh_ctrl_api_general_request(nh_ctrl_api_t *api, CtrlMsg *tx_msg) {
    nh_ret_t ret = NH_RET_SUCCESS;

    size_t packed_size = ctrl_msg__get_packed_size(tx_msg);
    if (packed_size == 0) {
        return NH_RET_FAIL;
    }

    uint8_t *req_buf;

    ret = NH_CTRL_API_ALLOC(api, &req_buf, packed_size);
    if (ret != NH_RET_SUCCESS) {
        return ret;
    }

    if (ctrl_msg__pack(tx_msg, req_buf) != packed_size) {
        ret = NH_RET_FAIL;
        goto free_tx_exit;
    }

    ret = nh_shared_if_serial_if_send(api->shared_if, req_buf, packed_size, NH_COMMAND_TIMEOUT_MSEC);

free_tx_exit:
    NH_CTRL_API_FREE(api, req_buf);

    return ret;
}

static nh_ret_t nh_ctrl_api_general_response(nh_ctrl_api_t *api, CtrlMsg **rx_msg, ProtobufCAllocator *allocator) {
    nh_ret_t ret = NH_RET_SUCCESS;

    ret = api->osa->sem_take(api->osa->user_data, api->p_sem_resp, NH_COMMAND_TIMEOUT_MSEC);
    if (ret != NH_RET_SUCCESS) {
        return ret;
    }

    *rx_msg = ctrl_msg__unpack(allocator, api->p_buf_rx_len, api->p_buf_rx);
    if (*rx_msg == NULL) {
        ret = NH_RET_FAIL;
    }

    nh_shared_if_serial_if_free(api->shared_if, api->p_buf_rx);

    return ret;
}

static void *nh_ctrl_api_protobuf_allocate(void *handle, size_t size) {
    nh_ctrl_api_t *api = handle;
    uint8_t       *buf = NULL;

    if (NH_CTRL_API_ALLOC(api, &buf, size) != NH_RET_SUCCESS) {
        return NULL;
    }

    return buf;
}

static void nh_ctrl_api_protobuf_free(void *handle, void *data) {
    nh_ctrl_api_t *api = handle;

    NH_CTRL_API_FREE(api, data);
}