#include "simple_ble.hpp"
#include "sdkconfig.h"

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

#include <nimble/ble.h>
#include <host/ble_att.h>
#include <host/ble_gap.h>
#include <host/ble_gatt.h>
#include <host/ble_hs.h>
#include <host/ble_hs_mbuf.h>
#include <host/ble_sm.h>
#include <host/ble_store.h>
#include <host/ble_uuid.h>
#include <host/util/util.h>
#include <nimble/nimble_port.h>
#include <nimble/nimble_port_freertos.h>
#include <services/gap/ble_svc_gap.h>
#include <services/gatt/ble_svc_gatt.h>

#include <cstdio>
#include <cstring>

extern "C" void ble_store_config_init(void);

namespace microbu::ble {
namespace {

const char* TAG = "cits_ble";
constexpr uint32_t FIXED_PASSKEY = 123456;

// Base UUID: 0000xxxx-ba5e-4c17-8000-00805f9b34fb
// Little-endian byte order for NimBLE (16 bytes, index 15 down to 0):
#define CITS_UUID_BASE(id_lo, id_hi) \
    0xfb, 0x34, 0x9b, 0x5f, 0x80, 0x00, 0x00, 0x80, \
    0x17, 0x4c, 0x5e, 0xba, (id_lo), (id_hi), 0x00, 0x00

static const ble_uuid128_t s_svc_uuid         = BLE_UUID128_INIT(CITS_UUID_BASE(0x75, 0xc1));
static const ble_uuid128_t s_chr_btp_req_uuid = BLE_UUID128_INIT(CITS_UUID_BASE(0x76, 0xc1));
static const ble_uuid128_t s_chr_btp_ind_uuid = BLE_UUID128_INIT(CITS_UUID_BASE(0x77, 0xc1));
static const ble_uuid128_t s_chr_poti_uuid    = BLE_UUID128_INIT(CITS_UUID_BASE(0x78, 0xc1));
static const ble_uuid128_t s_chr_status_uuid  = BLE_UUID128_INIT(CITS_UUID_BASE(0x79, 0xc1));
static const ble_uuid128_t s_chr_id_uuid      = BLE_UUID128_INIT(CITS_UUID_BASE(0x7a, 0xc1));
static const ble_uuid128_t s_chr_config_uuid  = BLE_UUID128_INIT(CITS_UUID_BASE(0x7b, 0xc1));
static const ble_uuid128_t s_chr_result_uuid  = BLE_UUID128_INIT(CITS_UUID_BASE(0x7c, 0xc1));

enum class ChrId : uintptr_t {
    BtpRequest = 1,
    BtpIndication = 2,
    Poti = 3,
    Status = 4,
    IdChange = 5,
    Configure = 6,
    Result = 7,
};

Receiver s_receiver;
Counters statistics;
char s_dev_name[32] = "micrOBU";
uint16_t s_btp_ind_val_handle = 0;
uint16_t s_status_val_handle = 0;
uint16_t s_id_val_handle = 0;
uint16_t s_result_val_handle = 0;
volatile uint16_t s_conn_handle = BLE_HS_CONN_HANDLE_NONE;
int s_notify_count = 0;   // number of characteristics with active CCCD subscription
bool s_subscribed = false;
uint8_t s_own_addr_type = 0;
volatile bool s_synced = false;
volatile bool s_adv_allowed = true;

// GATT characteristic read/write callback, shared by all characteristics (arg identifies which).
int chr_access(uint16_t conn_handle, uint16_t attr_handle,
               struct ble_gatt_access_ctxt *ctxt, void *arg) {
    uintptr_t id = reinterpret_cast<uintptr_t>(arg);
    if (ctxt->op == BLE_GATT_ACCESS_OP_WRITE_CHR) {
        uint16_t total = OS_MBUF_PKTLEN(ctxt->om);
        if (total == 0 || total > link::maximum_message) {
            ESP_LOGW(TAG, "rx invalid length: %u", (unsigned)total);
            ++statistics.malformed;
            return BLE_ATT_ERR_INVALID_ATTR_VALUE_LEN;
        }
        link::Bytes buf(total);
        uint16_t copied = 0;
        int rc = ble_hs_mbuf_to_flat(ctxt->om, buf.data(), total, &copied);
        if (rc != 0 || copied != total) {
            ++statistics.malformed;
            return BLE_ATT_ERR_UNLIKELY;
        }
        ++statistics.rx_messages;
        if (s_receiver) {
            // Require valid link header (opcode + flags + seq).
            // Host->station opcodes are 0x01..0x0F; station->host are 0x80+.
            if (total >= link::header_size && buf[0] >= 0x01 && buf[0] <= 0x0F) {
                ESP_LOGD(TAG, "rx chr %u op=0x%02x len=%u", (unsigned)id, buf[0], (unsigned)total);
                s_receiver(std::move(buf));
            } else {
                ESP_LOGW(TAG, "rx chr %u unframed or invalid op=0x%02x len=%u", (unsigned)id, buf[0], (unsigned)total);
                ++statistics.malformed;
                return BLE_ATT_ERR_INVALID_ATTR_VALUE_LEN;
            }
        }
        return 0;
    } else if (ctxt->op == BLE_GATT_ACCESS_OP_READ_CHR) {
        if (id == static_cast<uintptr_t>(ChrId::Status)) {
            return 0;  // empty read; status is pushed via notify
        }
    }
    return BLE_ATT_ERR_READ_NOT_PERMITTED;
}


#define CHR_FLAG_RW_ENC      (BLE_GATT_CHR_F_READ_ENC | BLE_GATT_CHR_F_READ_AUTHEN | \
                              BLE_GATT_CHR_F_WRITE_ENC | BLE_GATT_CHR_F_WRITE_AUTHEN)
#define CHR_FLAG_NOTIFY_ENC  (BLE_GATT_CHR_F_NOTIFY_INDICATE_ENC | \
                              BLE_GATT_CHR_F_NOTIFY_INDICATE_AUTHEN)

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wmissing-field-initializers"
static struct ble_gatt_chr_def s_chr_defs[] = {
    {
        // NF-SAP: BTP-DATA.request (Central -> Station Write)
        .uuid = &s_chr_btp_req_uuid.u,
        .access_cb = chr_access,
        .arg = reinterpret_cast<void*>(ChrId::BtpRequest),
        .flags = BLE_GATT_CHR_F_WRITE | CHR_FLAG_RW_ENC,
    },
    {
        // NF-SAP: BTP-DATA.indication (Station -> Central Notify)
        .uuid = &s_chr_btp_ind_uuid.u,
        .access_cb = chr_access,
        .arg = reinterpret_cast<void*>(ChrId::BtpIndication),
        .flags = BLE_GATT_CHR_F_NOTIFY | CHR_FLAG_NOTIFY_ENC,
        .val_handle = &s_btp_ind_val_handle,
    },
    {
        // PoTi: Position & Time Fix Update (Central -> Station Write)
        .uuid = &s_chr_poti_uuid.u,
        .access_cb = chr_access,
        .arg = reinterpret_cast<void*>(ChrId::Poti),
        .flags = BLE_GATT_CHR_F_WRITE | CHR_FLAG_RW_ENC,
    },
    {
        // Station Status (Station -> Central Read/Notify)
        .uuid = &s_chr_status_uuid.u,
        .access_cb = chr_access,
        .arg = reinterpret_cast<void*>(ChrId::Status),
        .flags = BLE_GATT_CHR_F_READ | BLE_GATT_CHR_F_NOTIFY | CHR_FLAG_RW_ENC | CHR_FLAG_NOTIFY_ENC,
        .val_handle = &s_status_val_handle,
    },
    {
        // SF-SAP: Identity Change Event (Station -> Central Notify)
        .uuid = &s_chr_id_uuid.u,
        .access_cb = chr_access,
        .arg = reinterpret_cast<void*>(ChrId::IdChange),
        .flags = BLE_GATT_CHR_F_WRITE | BLE_GATT_CHR_F_NOTIFY | CHR_FLAG_RW_ENC | CHR_FLAG_NOTIFY_ENC,
        .val_handle = &s_id_val_handle,
    },
    {
        // Station Configure / Credentials (Central -> Station Write)
        .uuid = &s_chr_config_uuid.u,
        .access_cb = chr_access,
        .arg = reinterpret_cast<void*>(ChrId::Configure),
        .flags = BLE_GATT_CHR_F_WRITE | CHR_FLAG_RW_ENC,
    },
    {
        // Result / Response (Station -> Central Notify)
        .uuid = &s_chr_result_uuid.u,
        .access_cb = chr_access,
        .arg = reinterpret_cast<void*>(ChrId::Result),
        .flags = BLE_GATT_CHR_F_NOTIFY | CHR_FLAG_NOTIFY_ENC,
        .val_handle = &s_result_val_handle,
    },
    {0}
};

static struct ble_gatt_svc_def s_svc_defs[] = {
    {
        .type = BLE_GATT_SVC_TYPE_PRIMARY,
        .uuid = &s_svc_uuid.u,
        .characteristics = s_chr_defs,
    },
    {0}
};
#pragma GCC diagnostic pop

int start_advertising();

// NimBLE GAP callback: connect/disconnect, pairing (fixed-passkey), and notify subscriptions.
int gap_event(struct ble_gap_event *event, void *arg) {
    struct ble_gap_conn_desc desc;
    switch (event->type) {
    case BLE_GAP_EVENT_CONNECT:
        ESP_LOGI(TAG, "connect status=%d handle=%d", event->connect.status, event->connect.conn_handle);
        if (event->connect.status == 0) {
            s_conn_handle = event->connect.conn_handle;
            s_subscribed = false;
            struct ble_gap_upd_params params = {};
            params.itvl_min = 12;              // 15 ms
            params.itvl_max = 16;              // 20 ms
            params.latency = 0;
            params.supervision_timeout = 400;  // 4 s
            params.min_ce_len = 0;
            params.max_ce_len = 0;
            ble_gap_update_params(event->connect.conn_handle, &params);
            ble_gap_security_initiate(event->connect.conn_handle);
        } else {
            start_advertising();
        }
        return 0;

    case BLE_GAP_EVENT_DISCONNECT:
        ESP_LOGI(TAG, "disconnect reason=%d", event->disconnect.reason);
        s_conn_handle = BLE_HS_CONN_HANDLE_NONE;
        s_subscribed = false;
        s_notify_count = 0;
        start_advertising();
        return 0;

    case BLE_GAP_EVENT_REPEAT_PAIRING:
        // MicrOBU: logged, since a phone that pairs again while we hold its bond means one side
        // lost the keys; the phone's side then usually needs "forget device" as well.
        ESP_LOGW(TAG, "phone started pairing again although bonded: dropping the old bond");
        if (ble_gap_conn_find(event->repeat_pairing.conn_handle, &desc) == 0) {
            ble_store_util_delete_peer(&desc.peer_id_addr);
        }
        return BLE_GAP_REPEAT_PAIRING_RETRY;

    case BLE_GAP_EVENT_ENC_CHANGE:
        // MicrOBU: the one place the board side says whether the link got secure. 0 = encrypted.
        if (ble_gap_conn_find(event->enc_change.conn_handle, &desc) == 0) {
            ESP_LOGI(TAG, "encryption change status=%d encrypted=%d authenticated=%d bonded=%d",
                     event->enc_change.status, desc.sec_state.encrypted, desc.sec_state.authenticated,
                     desc.sec_state.bonded);
        } else {
            ESP_LOGI(TAG, "encryption change status=%d", event->enc_change.status);
        }
        return 0;

    case BLE_GAP_EVENT_MTU:
        ESP_LOGI(TAG, "ATT MTU %u", unsigned(event->mtu.value));
        return 0;

    case BLE_GAP_EVENT_PASSKEY_ACTION:
        if (event->passkey.params.action == BLE_SM_IOACT_DISP) {
            struct ble_sm_io pkey = {};
            pkey.action  = BLE_SM_IOACT_DISP;
            pkey.passkey = FIXED_PASSKEY;
            ESP_LOGW(TAG, "BLE PAIRING FIXED PASSKEY: %06lu", static_cast<unsigned long>(FIXED_PASSKEY));
            int rc = ble_sm_inject_io(event->passkey.conn_handle, &pkey);
            if (rc != 0) {
                ESP_LOGW(TAG, "ble_sm_inject_io rc=%d", rc);
            }
        } else {
            ESP_LOGW(TAG, "unhandled passkey action %d", event->passkey.params.action);
        }
        return 0;

    case BLE_GAP_EVENT_SUBSCRIBE:
        if (event->subscribe.attr_handle == s_btp_ind_val_handle ||
            event->subscribe.attr_handle == s_status_val_handle ||
            event->subscribe.attr_handle == s_id_val_handle ||
            event->subscribe.attr_handle == s_result_val_handle) {
            if (event->subscribe.cur_notify && !event->subscribe.prev_notify)
                ++s_notify_count;
            else if (!event->subscribe.cur_notify && event->subscribe.prev_notify)
                --s_notify_count;
            if (s_notify_count < 0) s_notify_count = 0;
            s_subscribed = (s_notify_count > 0);
            ESP_LOGI(TAG, "subscription handle %u changed: notify=%d count=%d",
                     event->subscribe.attr_handle, (int)event->subscribe.cur_notify, s_notify_count);
        }
        return 0;

    case BLE_GAP_EVENT_ADV_COMPLETE:
        start_advertising();
        return 0;

    default:
        return 0;
    }
}

// (Re)starts undirected advertising with the device name and service UUID, unless paused.
int start_advertising() {
    if (!s_adv_allowed || !s_synced || s_conn_handle != BLE_HS_CONN_HANDLE_NONE || ble_gap_adv_active()) return 0;
    const char* name = s_dev_name;
    size_t name_len = strlen(name);

    struct ble_hs_adv_fields adv = {};
    adv.flags = BLE_HS_ADV_F_DISC_GEN | BLE_HS_ADV_F_BREDR_UNSUP;
    adv.tx_pwr_lvl_is_present = 1;
    adv.tx_pwr_lvl = BLE_HS_ADV_TX_PWR_LVL_AUTO;
    adv.name = reinterpret_cast<const uint8_t*>(name);
    adv.name_len = static_cast<uint8_t>(name_len);
    adv.name_is_complete = 1;

    int rc = ble_gap_adv_set_fields(&adv);
    if (rc != 0) {
        ESP_LOGE(TAG, "adv_set_fields rc=%d", rc);
        return rc;
    }

    struct ble_hs_adv_fields rsp = {};
    rsp.uuids128 = &s_svc_uuid;
    rsp.num_uuids128 = 1;
    rsp.uuids128_is_complete = 1;

    rc = ble_gap_adv_rsp_set_fields(&rsp);
    if (rc != 0) {
        ESP_LOGE(TAG, "adv_rsp_set_fields rc=%d", rc);
        return rc;
    }

    struct ble_gap_adv_params params = {};
    params.conn_mode = BLE_GAP_CONN_MODE_UND;
    params.disc_mode = BLE_GAP_DISC_MODE_GEN;

    rc = ble_gap_adv_start(s_own_addr_type, nullptr, BLE_HS_FOREVER,
                           &params, gap_event, nullptr);
    if (rc != 0) {
        ESP_LOGE(TAG, "adv_start rc=%d", rc);
        return rc;
    }
    ESP_LOGI(TAG, "BLE advertising started as '%s'", name);
    return 0;
}

// NimBLE host reset callback.
void on_reset(int reason) {
    ESP_LOGE(TAG, "NimBLE reset, reason=%d", reason);
}

// NimBLE host sync callback: resolves our address and starts advertising.
void on_sync() {
    int rc = ble_hs_util_ensure_addr(0);
    assert(rc == 0);
    rc = ble_hs_id_infer_auto(0, &s_own_addr_type);
    if (rc != 0) {
        ESP_LOGE(TAG, "ble_hs_id_infer_auto rc=%d", rc);
        return;
    }
    s_synced = true;
    // MicrOBU: a bond that did not persist (full NVS) looks like "the phone has to pair every time".
    int bonds = 0;
    if (ble_store_util_count(BLE_STORE_OBJ_TYPE_OUR_SEC, &bonds) == 0)
        ESP_LOGI(TAG, "%d bonded phone(s) in NVS", bonds);
    start_advertising();
}

void nimble_host_task(void* param) {
    ESP_LOGI(TAG, "NimBLE host task running");
    nimble_port_run();
    nimble_port_freertos_deinit();
}

} // namespace

bool start(Receiver receiver) {
    s_receiver = std::move(receiver);

    uint8_t mac[6] = {};
    if (esp_read_mac(mac, ESP_MAC_BT) == ESP_OK) {
        std::snprintf(s_dev_name, sizeof(s_dev_name), "micrOBU-%02X%02X", mac[4], mac[5]);
    }

    esp_err_t err = nimble_port_init();
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "nimble_port_init failed: %d", err);
        return false;
    }

    ble_hs_cfg.reset_cb = on_reset;
    ble_hs_cfg.sync_cb = on_sync;
    ble_hs_cfg.store_status_cb = ble_store_util_status_rr;

    // Fixed passkey LE Secure Connections + Bonding
    ble_hs_cfg.sm_io_cap = BLE_HS_IO_DISPLAY_ONLY;
    ble_hs_cfg.sm_sc = 1;
    ble_hs_cfg.sm_bonding = 1;
    ble_hs_cfg.sm_mitm = 1;
    ble_hs_cfg.sm_our_key_dist = BLE_SM_PAIR_KEY_DIST_ENC | BLE_SM_PAIR_KEY_DIST_ID;
    ble_hs_cfg.sm_their_key_dist = BLE_SM_PAIR_KEY_DIST_ENC | BLE_SM_PAIR_KEY_DIST_ID;

    ble_svc_gap_init();
    ble_svc_gatt_init();
    ble_svc_gap_device_name_set(s_dev_name);

    int rc = ble_gatts_count_cfg(s_svc_defs);
    if (rc != 0) {
        ESP_LOGE(TAG, "ble_gatts_count_cfg failed: %d", rc);
        nimble_port_deinit();
        return false;
    }
    rc = ble_gatts_add_svcs(s_svc_defs);
    if (rc != 0) {
        ESP_LOGE(TAG, "ble_gatts_add_svcs failed: %d", rc);
        nimble_port_deinit();
        return false;
    }

    ble_store_config_init();
    nimble_port_freertos_init(nimble_host_task);
    ESP_LOGI(TAG, "simple_ble initialized (fixed passkey %06lu)", static_cast<unsigned long>(FIXED_PASSKEY));
    return true;
}

bool write(const link::Bytes& message) {
    uint16_t conn_handle = s_conn_handle;
    if (conn_handle == BLE_HS_CONN_HANDLE_NONE) {
        ++statistics.tx_failed;
        return false;
    }
    if (message.empty() || message.size() > link::maximum_message) {
        ++statistics.tx_failed;
        return false;
    }
    uint16_t val_handle = s_result_val_handle;
    if (message.size() >= link::header_size) {
        link::Opcode op = static_cast<link::Opcode>(message[0]);
        switch (op) {
        case link::Opcode::V2X_RX:
        case link::Opcode::BTP_DATA_INDICATION:
            val_handle = s_btp_ind_val_handle;
            break;
        case link::Opcode::SF_IDCHANGE_EVENT:
            val_handle = s_id_val_handle;
            break;
        case link::Opcode::STATUS:
            // Route unsolicited or requested STATUS over the subscribed result/reply characteristic
            val_handle = s_result_val_handle;
            break;
        case link::Opcode::RESULT:
        default:
            val_handle = s_result_val_handle;
            break;
        }
    }
    // MicrOBU: one notification is one whole message (no fragmentation on this GATT layout), and
    // NimBLE cuts a value longer than ATT_MTU - 3 short without an error. The phone asks for an
    // MTU of 517; until it has, drop rather than deliver a truncated message.
    if (message.size() + 3 > ble_att_mtu(conn_handle)) {
        ++statistics.tx_failed;
        return false;
    }
    struct os_mbuf *om = ble_hs_mbuf_from_flat(message.data(), message.size());
    if (!om) {
        ++statistics.tx_failed;
        return false;
    }
    int rc = ble_gatts_notify_custom(conn_handle, val_handle, om);
    if (rc != 0) {
        ESP_LOGW(TAG, "notify failed handle=%u op=0x%02x rc=%d", (unsigned)val_handle, (unsigned)(message.empty() ? 0 : message[0]), rc);
        if (rc == BLE_HS_ENOTSUP) os_mbuf_free_chain(om);
        ++statistics.tx_failed;
        return false;
    }
    ++statistics.tx_messages;
    return true;
}

bool connected() {
    return s_conn_handle != BLE_HS_CONN_HANDLE_NONE && s_subscribed;
}

void set_advertising_allowed(bool allowed) {
    if (s_adv_allowed == allowed) return;
    s_adv_allowed = allowed;
    if (!s_synced) return;
    if (allowed) {
        start_advertising();
        ESP_LOGI(TAG, "USB link idle: BLE advertising resumed");
    } else if (ble_gap_adv_active()) {
        ble_gap_adv_stop();
        ESP_LOGI(TAG, "USB link in use: BLE advertising paused");
    }
}

bool advertising_allowed() {
    return s_adv_allowed;
}

bool is_subscribed() {
    return s_subscribed;
}

Counters counters() {
    return statistics;
}

} // namespace microbu::ble