#include "ble_host.h"
#include "esp_log.h"
#include "esp_timer.h"
#include "host_transport.h"
#include "nvs_flash.h"
#include "ota_espnow.h"
#include "freertos/FreeRTOS.h"
#include "freertos/semphr.h"
#include "freertos/task.h"

#include "nimble/nimble_port.h"
#include "nimble/nimble_port_freertos.h"
#include "host/ble_att.h"
#include "host/ble_hs.h"
#include "host/util/util.h"
#include "services/gap/ble_svc_gap.h"
#include "services/gatt/ble_svc_gatt.h"

#include <string.h>

static const char *TAG = "[BLE_HOST]";
static const char *DEVICE_NAME = "Powerpod";

#define BLE_CONN_ITVL_MIN_UNITS 6 /* 7.5 ms */
#define BLE_CONN_ITVL_MAX_UNITS 6
#define BLE_CONN_TUNE_MAX_TRIES 6
/** Prefer large ATT MTU so CLIENT_INFO / BATTERY fit in fewer notify chunks. */
#define BLE_PREFERRED_ATT_MTU 517
#define BLE_NOTIFY_CHUNK_RETRIES 24
/** During ESP-NOW OTA dist, fail notify sooner so WiFi/ESP-NOW keeps airtime. */
#define BLE_NOTIFY_CHUNK_RETRIES_OTA_DIST 4
#define BLE_NOTIFY_CHUNK_RETRY_MS 5
/** Pace multi-chunk notifies so the central parser sees contiguous frame bytes. */
#define BLE_NOTIFY_CHUNK_GAP_MS 8
#define BLE_NOTIFY_CHUNK_GAP_OTA_DIST_MS 2

/*
 * Nordic UART style:
 *   Service:        6e400001-b5a3-f393-e0a9-e50e24dcca9e
 *   Characteristic: 6e400002-b5a3-f393-e0a9-e50e24dcca9e
 * Write = host→ESP (framed), Notify = ESP→host (framed).
 */
static const ble_uuid128_t ble_svc_uuid =
    BLE_UUID128_INIT(0x9e, 0xca, 0xdc, 0x24, 0x0e, 0xe5, 0xa9, 0xe0, 0x93, 0xf3,
                     0xa3, 0xb5, 0x01, 0x00, 0x40, 0x6e);
static const ble_uuid128_t ble_chr_uuid =
    BLE_UUID128_INIT(0x9e, 0xca, 0xdc, 0x24, 0x0e, 0xe5, 0xa9, 0xe0, 0x93, 0xf3,
                     0xa3, 0xb5, 0x02, 0x00, 0x40, 0x6e);

static uint8_t own_addr_type;
static uint16_t conn_handle = BLE_HS_CONN_HANDLE_NONE;
static uint16_t ble_chr_val_handle;
static bool notify_enabled;
static uint16_t att_mtu = 23;
static esp_timer_handle_t conn_tune_timer;
static uint8_t conn_tune_tries;
static SemaphoreHandle_t s_tx_mu;

static int ble_gap_event(struct ble_gap_event *event, void *arg);
static void ble_request_fast_conn_params(uint16_t handle);
static void ble_optimize_link(uint16_t handle);

bool ble_host_ready(void) {
  return notify_enabled && conn_handle != BLE_HS_CONN_HANDLE_NONE;
}

uint16_t ble_host_att_mtu(void) { return att_mtu; }

uint16_t ble_host_max_payload(void) {
  if (att_mtu < 4) {
    return 20;
  }
  return (uint16_t)(att_mtu - 3);
}

/**
 * One ATT notification (≤ MTU-3). Retries on transient NimBLE buffer pressure.
 */
static esp_err_t ble_notify_chunk(const uint8_t *data, uint16_t len) {
  const int max_retries = ota_espnow_distribution_active()
                              ? BLE_NOTIFY_CHUNK_RETRIES_OTA_DIST
                              : BLE_NOTIFY_CHUNK_RETRIES;
  for (int attempt = 0; attempt < max_retries; attempt++) {
    struct os_mbuf *om = ble_hs_mbuf_from_flat(data, len);
    if (om == NULL) {
      vTaskDelay(pdMS_TO_TICKS(BLE_NOTIFY_CHUNK_RETRY_MS));
      continue;
    }

    int rc = ble_gatts_notify_custom(conn_handle, ble_chr_val_handle, om);
    if (rc == 0) {
      return ESP_OK;
    }
    /* Ownership stays with us on failure. */
    os_mbuf_free_chain(om);
    if (!ble_host_ready()) {
      return ESP_ERR_INVALID_STATE;
    }
    ESP_LOGD(TAG, "notify chunk rc=%d attempt=%d len=%u", rc, attempt,
             (unsigned)len);
    vTaskDelay(pdMS_TO_TICKS(BLE_NOTIFY_CHUNK_RETRY_MS));
  }
  ESP_LOGW(TAG, "notify chunk gave up after retries len=%u", (unsigned)len);
  return ESP_ERR_NO_MEM;
}

/**
 * Send a full host frame. If larger than ATT MTU-3, split across consecutive
 * notifies; goTool feeds every notify into the same byte parser, so the frame
 * reassembles without a protocol change. Mutex keeps stream ticks from
 * interleaving mid-frame.
 */
static esp_err_t ble_host_tx(const uint8_t *frame, size_t frame_len) {
  if (!ble_host_ready()) {
    return ESP_ERR_INVALID_STATE;
  }
  if (frame == NULL || frame_len == 0) {
    return ESP_ERR_INVALID_ARG;
  }

  uint16_t max_chunk = ble_host_max_payload();
  if (max_chunk == 0) {
    return ESP_ERR_INVALID_SIZE;
  }

  if (s_tx_mu == NULL ||
      xSemaphoreTake(s_tx_mu, pdMS_TO_TICKS(2000)) != pdTRUE) {
    return ESP_ERR_TIMEOUT;
  }

  esp_err_t result = ESP_OK;
  size_t off = 0;
  while (off < frame_len) {
    if (!ble_host_ready()) {
      result = ESP_ERR_INVALID_STATE;
      break;
    }
    size_t remain = frame_len - off;
    uint16_t n =
        (uint16_t)(remain > max_chunk ? max_chunk : remain);
    result = ble_notify_chunk(&frame[off], n);
    if (result != ESP_OK) {
      ESP_LOGW(TAG, "TX abort at %u/%u (chunk %u, mtu_pay %u)",
               (unsigned)off, (unsigned)frame_len, (unsigned)n,
               (unsigned)max_chunk);
      break;
    }
    off += n;
    if (off < frame_len) {
      const uint32_t gap_ms = ota_espnow_distribution_active()
                                 ? BLE_NOTIFY_CHUNK_GAP_OTA_DIST_MS
                                 : BLE_NOTIFY_CHUNK_GAP_MS;
      vTaskDelay(pdMS_TO_TICKS(gap_ms));
    }
  }

  xSemaphoreGive(s_tx_mu);
  return result;
}

static void ble_conn_tune_timer_cb(void *arg) {
  struct ble_gap_conn_desc desc;

  (void)arg;
  if (conn_handle == BLE_HS_CONN_HANDLE_NONE) {
    return;
  }
  if (ble_gap_conn_find(conn_handle, &desc) != 0) {
    return;
  }
  if (desc.conn_itvl <= BLE_CONN_ITVL_MAX_UNITS) {
    ESP_LOGI(TAG, "conn already at interval_us=%lu",
             (unsigned long)desc.conn_itvl * 1250UL);
    return;
  }
  if (conn_tune_tries >= BLE_CONN_TUNE_MAX_TRIES) {
    ESP_LOGW(TAG, "central kept interval_us=%lu (ESP requested 7500)",
             (unsigned long)desc.conn_itvl * 1250UL);
    return;
  }

  conn_tune_tries++;
  ESP_LOGI(TAG, "retry fast conn params (%u/%u); current interval_us=%lu",
           conn_tune_tries, BLE_CONN_TUNE_MAX_TRIES,
           (unsigned long)desc.conn_itvl * 1250UL);
  ble_request_fast_conn_params(conn_handle);
  esp_timer_start_once(conn_tune_timer, 750 * 1000);
}

static void ble_schedule_conn_tune(void) {
  conn_tune_tries = 0;
  esp_timer_stop(conn_tune_timer);
  esp_timer_start_once(conn_tune_timer, 300 * 1000);
}

static void ble_request_fast_conn_params(uint16_t handle) {
  struct ble_gap_upd_params params = {
      .itvl_min = BLE_CONN_ITVL_MIN_UNITS,
      .itvl_max = BLE_CONN_ITVL_MAX_UNITS,
      .latency = 0,
      .supervision_timeout = 400,
      .min_ce_len = 0,
      .max_ce_len = 24,
  };
  int rc = ble_gap_update_params(handle, &params);
  if (rc != 0) {
    ESP_LOGW(TAG, "conn param update request failed: %d", rc);
  } else {
    ESP_LOGI(TAG, "requested conn interval 7.5 ms");
  }
}

static void ble_optimize_link(uint16_t handle) {
  int rc;

  /*
   * Do NOT request 2M PHY here. With ESP-NOW/WiFi coexistence the controller
   * can assert in llc_phy_upd (e.g. "BLE assert llc_phy_upd.c …") and the
   * host link dies — no further responses. Stay on default 1M.
   */
  ble_request_fast_conn_params(handle);

  rc = ble_gap_set_data_len(handle, 251, 2120);
  if (rc != 0) {
    ESP_LOGW(TAG, "data len request failed: %d", rc);
  }

  ble_schedule_conn_tune();
}

static void ble_advertise(void) {
  struct ble_gap_adv_params adv_params;
  struct ble_hs_adv_fields fields;
  int rc;

  memset(&fields, 0, sizeof(fields));
  fields.flags = BLE_HS_ADV_F_DISC_GEN | BLE_HS_ADV_F_BREDR_UNSUP;
  fields.tx_pwr_lvl_is_present = 1;
  fields.tx_pwr_lvl = BLE_HS_ADV_TX_PWR_LVL_AUTO;
  fields.name = (uint8_t *)DEVICE_NAME;
  fields.name_len = strlen(DEVICE_NAME);
  fields.name_is_complete = 1;

  rc = ble_gap_adv_set_fields(&fields);
  if (rc != 0) {
    ESP_LOGE(TAG, "adv set fields failed: %d", rc);
    return;
  }

  memset(&adv_params, 0, sizeof(adv_params));
  adv_params.conn_mode = BLE_GAP_CONN_MODE_UND;
  adv_params.disc_mode = BLE_GAP_DISC_MODE_GEN;
  adv_params.itvl_min = 0x20;
  adv_params.itvl_max = 0x30;

  rc = ble_gap_adv_start(own_addr_type, NULL, BLE_HS_FOREVER, &adv_params,
                         ble_gap_event, NULL);
  if (rc != 0) {
    ESP_LOGE(TAG, "adv start failed: %d", rc);
    return;
  }

  ESP_LOGI(TAG, "advertising as \"%s\"", DEVICE_NAME);
}

static int ble_chr_access(uint16_t conn_handle_arg, uint16_t attr_handle,
                          struct ble_gatt_access_ctxt *ctxt, void *arg) {
  (void)conn_handle_arg;
  (void)attr_handle;
  (void)arg;

  switch (ctxt->op) {
  case BLE_GATT_ACCESS_OP_READ_CHR:
    return 0;

  case BLE_GATT_ACCESS_OP_WRITE_CHR: {
    uint16_t om_len = OS_MBUF_PKTLEN(ctxt->om);
    uint8_t buf[512];
    uint16_t copied = 0;
    int rc;

    if (om_len == 0 || om_len > sizeof(buf)) {
      return BLE_ATT_ERR_INVALID_ATTR_VALUE_LEN;
    }

    rc = ble_hs_mbuf_to_flat(ctxt->om, buf, sizeof(buf), &copied);
    if (rc != 0) {
      return BLE_ATT_ERR_UNLIKELY;
    }

    host_rx_bytes(buf, copied);
    return 0;
  }

  default:
    return BLE_ATT_ERR_UNLIKELY;
  }
}

static const struct ble_gatt_svc_def gatt_svcs[] = {
    {
        .type = BLE_GATT_SVC_TYPE_PRIMARY,
        .uuid = &ble_svc_uuid.u,
        .characteristics =
            (struct ble_gatt_chr_def[]){
                {
                    .uuid = &ble_chr_uuid.u,
                    .access_cb = ble_chr_access,
                    .val_handle = &ble_chr_val_handle,
                    .flags = BLE_GATT_CHR_F_READ | BLE_GATT_CHR_F_WRITE |
                             BLE_GATT_CHR_F_WRITE_NO_RSP | BLE_GATT_CHR_F_NOTIFY,
                },
                {0},
            },
    },
    {0},
};

static int gatt_svr_init(void) {
  int rc;

  ble_svc_gap_init();
  ble_svc_gatt_init();

  rc = ble_gatts_count_cfg(gatt_svcs);
  if (rc != 0) {
    return rc;
  }

  return ble_gatts_add_svcs(gatt_svcs);
}

static int ble_gap_event(struct ble_gap_event *event, void *arg) {
  (void)arg;

  switch (event->type) {
  case BLE_GAP_EVENT_CONNECT:
    ESP_LOGI(TAG, "connection %s; status=%d handle=%d",
             event->connect.status == 0 ? "established" : "failed",
             event->connect.status, event->connect.conn_handle);
    if (event->connect.status != 0) {
      ble_advertise();
    } else {
      conn_handle = event->connect.conn_handle;
      ble_optimize_link(conn_handle);
    }
    return 0;

  case BLE_GAP_EVENT_DISCONNECT:
    ESP_LOGI(TAG, "disconnect; reason=%d", event->disconnect.reason);
    esp_timer_stop(conn_tune_timer);
    conn_tune_tries = 0;
    conn_handle = BLE_HS_CONN_HANDLE_NONE;
    notify_enabled = false;
    att_mtu = 23;
    ble_advertise();
    return 0;

  case BLE_GAP_EVENT_CONN_UPDATE: {
    struct ble_gap_conn_desc desc;
    if (ble_gap_conn_find(event->conn_update.conn_handle, &desc) == 0) {
      ESP_LOGI(TAG, "conn updated; status=%d interval_us=%lu latency=%u",
               event->conn_update.status,
               (unsigned long)desc.conn_itvl * 1250UL, desc.conn_latency);
      if (event->conn_update.status == 0 &&
          desc.conn_itvl > BLE_CONN_ITVL_MAX_UNITS &&
          conn_tune_tries < BLE_CONN_TUNE_MAX_TRIES) {
        esp_timer_stop(conn_tune_timer);
        esp_timer_start_once(conn_tune_timer, 500 * 1000);
      }
    }
    return 0;
  }

  case BLE_GAP_EVENT_PHY_UPDATE_COMPLETE:
    ESP_LOGI(TAG, "phy update; status=%d tx=%u rx=%u",
             event->phy_updated.status, event->phy_updated.tx_phy,
             event->phy_updated.rx_phy);
    return 0;

  case BLE_GAP_EVENT_ADV_COMPLETE:
    ESP_LOGI(TAG, "advertise complete; reason=%d", event->adv_complete.reason);
    ble_advertise();
    return 0;

  case BLE_GAP_EVENT_SUBSCRIBE:
    if (event->subscribe.attr_handle == ble_chr_val_handle) {
      notify_enabled = event->subscribe.cur_notify;
      ESP_LOGI(TAG, "notify %s", notify_enabled ? "enabled" : "disabled");
      if (notify_enabled && conn_handle != BLE_HS_CONN_HANDLE_NONE) {
        ble_request_fast_conn_params(conn_handle);
      }
    }
    return 0;

  case BLE_GAP_EVENT_MTU:
    att_mtu = event->mtu.value;
    ESP_LOGI(TAG, "mtu update; handle=%d mtu=%d", event->mtu.conn_handle,
             event->mtu.value);
    return 0;

  default:
    return 0;
  }
}

static void ble_on_sync(void) {
  int rc;

  rc = ble_hs_id_infer_auto(0, &own_addr_type);
  if (rc != 0) {
    ESP_LOGE(TAG, "ble_hs_id_infer_auto failed: %d", rc);
    return;
  }

  uint8_t addr_val[6] = {0};
  ble_hs_id_copy_addr(own_addr_type, addr_val, NULL);
  ESP_LOGI(TAG, "device address: %02x:%02x:%02x:%02x:%02x:%02x", addr_val[5],
           addr_val[4], addr_val[3], addr_val[2], addr_val[1], addr_val[0]);

  ble_advertise();
}

static void ble_on_reset(int reason) {
  ESP_LOGE(TAG, "nimble reset; reason=%d", reason);
}

static void ble_host_task(void *param) {
  (void)param;
  ESP_LOGI(TAG, "nimble host task started");
  nimble_port_run();
  nimble_port_freertos_deinit();
}

esp_err_t init_ble_host(QueueHandle_t cmd_queue) {
  esp_err_t ret;

  host_transport_set_cmd_queue(cmd_queue);
  host_transport_set_tx(ble_host_tx);

  s_tx_mu = xSemaphoreCreateMutex();
  if (s_tx_mu == NULL) {
    ESP_LOGE(TAG, "failed to create TX mutex");
    return ESP_ERR_NO_MEM;
  }

  /* NVS may already be init via pod_settings; ignore ALREADY_INIT. */
  ret = nvs_flash_init();
  if (ret == ESP_ERR_NVS_NO_FREE_PAGES ||
      ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
    ESP_ERROR_CHECK(nvs_flash_erase());
    ret = nvs_flash_init();
  }
  if (ret != ESP_OK && ret != ESP_ERR_INVALID_STATE) {
    ESP_LOGE(TAG, "nvs_flash_init failed: %s", esp_err_to_name(ret));
    return ret;
  }

  const esp_timer_create_args_t tune_args = {
      .callback = &ble_conn_tune_timer_cb,
      .name = "ble_conn_tune",
  };
  ESP_ERROR_CHECK(esp_timer_create(&tune_args, &conn_tune_timer));

  ret = nimble_port_init();
  if (ret != ESP_OK) {
    ESP_LOGE(TAG, "nimble_port_init failed: %s", esp_err_to_name(ret));
    return ret;
  }

  ble_hs_cfg.sync_cb = ble_on_sync;
  ble_hs_cfg.reset_cb = ble_on_reset;

  int rc = gatt_svr_init();
  if (rc != 0) {
    ESP_LOGE(TAG, "gatt_svr_init failed: %d", rc);
    return ESP_FAIL;
  }

  rc = ble_svc_gap_device_name_set(DEVICE_NAME);
  if (rc != 0) {
    ESP_LOGE(TAG, "device name set failed: %d", rc);
    return ESP_FAIL;
  }

  rc = ble_att_set_preferred_mtu(BLE_PREFERRED_ATT_MTU);
  if (rc != 0) {
    ESP_LOGW(TAG, "preferred MTU set failed: %d", rc);
  } else {
    ESP_LOGI(TAG, "preferred ATT MTU %u", (unsigned)BLE_PREFERRED_ATT_MTU);
  }

  nimble_port_freertos_init(ble_host_task);
  ESP_LOGI(TAG, "BLE host transport started");
  return ESP_OK;
}