#include "TBeamSupremeBleInterface.h"

#include <BLE2902.h>
#include <BLEService.h>
#include <Cryptography/Random.h>
#include <Identity.h>
#include <Log.h>

#include <algorithm>
#include <string.h>

using namespace RNS;

#ifndef NODE_LABEL
#define NODE_LABEL "?"
#endif

static TBeamSupremeBleInterface* active_ble_interface = nullptr;

class TBeamSupremeBleInterface::ServerCallbacks : public BLEServerCallbacks {
public:
  explicit ServerCallbacks(TBeamSupremeBleInterface* owner) : _owner(owner) {}

  void onConnect(BLEServer* server) override {
    (void)server;
    _owner->_connected = true;
    _owner->_server_handshake_received = false;
    Serial.println("BLE peripheral: central connected");
    INFO("BLE central connected");
  }

  void onDisconnect(BLEServer* server) override {
    (void)server;
    _owner->_connected = false;
    _owner->_server_handshake_received = false;
    _owner->reset_reassembly();
    Serial.println("BLE peripheral: central disconnected; restarting advertising");
    INFO("BLE central disconnected; restarting advertising");
    BLEDevice::startAdvertising();
  }

private:
  TBeamSupremeBleInterface* _owner;
};

class TBeamSupremeBleInterface::RxCallbacks : public BLECharacteristicCallbacks {
public:
  explicit RxCallbacks(TBeamSupremeBleInterface* owner) : _owner(owner) {}

  void onWrite(BLECharacteristic* characteristic) override {
    std::string value = characteristic->getValue();
    if (value.empty()) {
      return;
    }

    if (!_owner->_server_handshake_received && value.size() == 16) {
      _owner->_server_handshake_received = true;
      Serial.println("BLE peripheral: identity handshake received");
      INFOF("BLE identity handshake received: %d bytes", (int)value.size());
      return;
    }

    _owner->handle_fragment(reinterpret_cast<const uint8_t*>(value.data()), value.size());
  }

private:
  TBeamSupremeBleInterface* _owner;
};

class TBeamSupremeBleInterface::AdvertisedDeviceCallbacks : public BLEAdvertisedDeviceCallbacks {
public:
  explicit AdvertisedDeviceCallbacks(TBeamSupremeBleInterface* owner) : _owner(owner) {}

  void onResult(BLEAdvertisedDevice advertised_device) override {
    bool has_service = advertised_device.haveServiceUUID() &&
        advertised_device.isAdvertisingService(BLEUUID(TBeamSupremeBleInterface::SERVICE_UUID));
    bool has_rns_name = advertised_device.haveName() &&
        advertised_device.getName().rfind("RNS-", 0) == 0;

    if ((has_service || has_rns_name) || _owner->_scan_report_count < 5) {
      Serial.printf("BLE scan: addr=%s name=%s service=%s rssi=%d\r\n",
                    advertised_device.getAddress().toString().c_str(),
                    advertised_device.haveName() ? advertised_device.getName().c_str() : "",
                    has_service ? "yes" : "no",
                    advertised_device.getRSSI());
      _owner->_scan_report_count++;
    }

    if (!has_service && !has_rns_name) {
      return;
    }

    BLEDevice::getScan()->stop();
    delete _owner->_advertised_device;
    _owner->_advertised_device = new BLEAdvertisedDevice(advertised_device);
    _owner->_do_connect = true;
    _owner->_scanning = false;
    Serial.printf("BLE central: peer candidate found: %s\r\n",
                  advertised_device.getAddress().toString().c_str());
    INFOF("BLE peer found: %s", advertised_device.getAddress().toString().c_str());
  }

private:
  TBeamSupremeBleInterface* _owner;
};

TBeamSupremeBleInterface::TBeamSupremeBleInterface(const char* name) : InterfaceImpl(name) {
  _IN = true;
  _OUT = true;
  _bitrate = 1000000;
  _HW_MTU = BLE_PAYLOAD_SIZE;
#if defined(BLE_ROLE_CENTRAL)
  _central_role = true;
#else
  _central_role = false;
#endif
}

TBeamSupremeBleInterface::~TBeamSupremeBleInterface() {
  stop();
  delete _advertised_device;
}

bool TBeamSupremeBleInterface::start() {
  if (_started) {
    return true;
  }

  active_ble_interface = this;
  BLEDevice::init(std::string("RNS-") + NODE_LABEL);
  BLEDevice::setMTU(BLE_MTU);

  if (_central_role) {
    Serial.println("BLE central: starting scanner");
    INFO("BLE starting as central/client");
    start_central_scan();
  } else {
    Serial.println("BLE peripheral: starting advertiser");
    INFO("BLE starting as peripheral/server");
    start_peripheral();
  }

  _started = true;
  _online = true;
  return true;
}

void TBeamSupremeBleInterface::stop() {
  _online = false;
  _connected = false;
  if (_client && _client->isConnected()) {
    _client->disconnect();
  }
}

void TBeamSupremeBleInterface::loop() {
  if (!_online) {
    return;
  }

  if (_central_role) {
    if (_do_connect && _advertised_device) {
      _do_connect = false;
      if (!connect_to_advertised_device(_advertised_device)) {
        _connected = false;
        _next_scan_ms = millis() + SCAN_RETRY_MS;
      }
    }

    if (!_connected && !_scanning && (int32_t)(millis() - _next_scan_ms) >= 0) {
      start_central_scan();
    }
  }

  if (_reassembly_started_ms != 0 && millis() - _reassembly_started_ms > REASSEMBLY_TIMEOUT_MS) {
    WARNING("BLE reassembly timeout; dropping partial packet");
    reset_reassembly();
  }

  RNS::Bytes packet({RNS::Type::NONE});
  while (dequeue_packet(packet)) {
    InterfaceImpl::handle_incoming(packet);
  }
}

void TBeamSupremeBleInterface::start_peripheral() {
  _server = BLEDevice::createServer();
  _server->setCallbacks(new ServerCallbacks(this));

  BLEService* service = _server->createService(SERVICE_UUID);
  _tx_characteristic = service->createCharacteristic(
      TX_UUID,
      BLECharacteristic::PROPERTY_READ | BLECharacteristic::PROPERTY_NOTIFY);
  _tx_characteristic->addDescriptor(new BLE2902());

  BLECharacteristic* rx_characteristic = service->createCharacteristic(
      RX_UUID,
      BLECharacteristic::PROPERTY_WRITE | BLECharacteristic::PROPERTY_WRITE_NR);
  rx_characteristic->setCallbacks(new RxCallbacks(this));

  _identity_characteristic = service->createCharacteristic(
      IDENTITY_UUID,
      BLECharacteristic::PROPERTY_READ);
  RNS::Bytes identity = local_identity_hash();
  _identity_characteristic->setValue((uint8_t*)identity.data(), identity.size());

  service->start();

  BLEAdvertising* advertising = BLEDevice::getAdvertising();
  advertising->addServiceUUID(SERVICE_UUID);
  advertising->setScanResponse(true);
  advertising->setMinPreferred(0x00);
  BLEDevice::startAdvertising();
  Serial.println("BLE peripheral: advertising Reticulum service");
  INFO("BLE advertising Reticulum service");
}

void TBeamSupremeBleInterface::start_central_scan() {
  _scanning = true;
  _next_scan_ms = millis() + SCAN_RETRY_MS;
  _scan_report_count = 0;

  BLEScan* scan = BLEDevice::getScan();
  scan->setAdvertisedDeviceCallbacks(new AdvertisedDeviceCallbacks(this), true);
  scan->setInterval(1349);
  scan->setWindow(449);
  scan->setActiveScan(true);
  Serial.println("BLE central: scanning for Reticulum service");
  scan->start(5, false);
  if (!_do_connect && !_connected) {
    _scanning = false;
    _next_scan_ms = millis() + SCAN_RETRY_MS;
    Serial.println("BLE central: scan complete; no peer found");
  }
  scan->clearResults();
  INFO("BLE scanning for Reticulum service");
}

bool TBeamSupremeBleInterface::connect_to_advertised_device(BLEAdvertisedDevice* device) {
  Serial.printf("BLE central: connecting to %s\r\n", device->getAddress().toString().c_str());
  INFOF("BLE connecting to %s", device->getAddress().toString().c_str());

  _client = BLEDevice::createClient();
  if (!_client->connect(device)) {
    Serial.println("BLE central: connect failed");
    WARNING("BLE connect failed");
    return false;
  }

  _client->setMTU(BLE_MTU);
  BLERemoteService* service = _client->getService(BLEUUID(SERVICE_UUID));
  if (!service) {
    Serial.println("BLE central: Reticulum service not found after connect");
    WARNING("BLE Reticulum service not found after connect");
    _client->disconnect();
    return false;
  }

  _remote_rx_characteristic = service->getCharacteristic(BLEUUID(RX_UUID));
  _remote_tx_characteristic = service->getCharacteristic(BLEUUID(TX_UUID));
  BLERemoteCharacteristic* identity_characteristic = service->getCharacteristic(BLEUUID(IDENTITY_UUID));

  if (!_remote_rx_characteristic || !_remote_tx_characteristic) {
    Serial.println("BLE central: RX/TX characteristics not found");
    WARNING("BLE Reticulum RX/TX characteristics not found");
    _client->disconnect();
    return false;
  }

  if (identity_characteristic && identity_characteristic->canRead()) {
    std::string peer_identity = identity_characteristic->readValue();
    Serial.printf("BLE central: peer identity read: %d bytes\r\n", (int)peer_identity.size());
    INFOF("BLE peer identity read: %d bytes", (int)peer_identity.size());
  }

  if (_remote_tx_characteristic->canNotify()) {
    _remote_tx_characteristic->registerForNotify(client_notify_callback);
  } else {
    Serial.println("BLE central: peer TX characteristic does not notify");
  }

  RNS::Bytes identity = local_identity_hash();
  _remote_rx_characteristic->writeValue((uint8_t*)identity.data(), identity.size(), false);
  _connected = true;
  _scanning = false;
  Serial.println("BLE central: connected and identity handshake sent");
  INFO("BLE connected and identity handshake sent");
  return true;
}

void TBeamSupremeBleInterface::send_outgoing(const Bytes& data) {
  if (!_online || !_connected) {
    return;
  }

  size_t total = (data.size() + BLE_PAYLOAD_SIZE - 1) / BLE_PAYLOAD_SIZE;
  if (total == 0 || total > 65535) {
    WARNINGF("BLE cannot fragment packet of size %d", (int)data.size());
    return;
  }

  for (size_t i = 0; i < total; ++i) {
    uint8_t fragment[BLE_MTU];
    uint8_t fragment_type = FRAG_CONTINUE;
    if (i == 0) {
      fragment_type = FRAG_START;
    } else if (i == total - 1) {
      fragment_type = FRAG_END;
    }

    size_t offset = i * BLE_PAYLOAD_SIZE;
    size_t chunk = std::min(BLE_PAYLOAD_SIZE, data.size() - offset);
    fragment[0] = fragment_type;
    fragment[1] = (uint8_t)((i >> 8) & 0xFF);
    fragment[2] = (uint8_t)(i & 0xFF);
    fragment[3] = (uint8_t)((total >> 8) & 0xFF);
    fragment[4] = (uint8_t)(total & 0xFF);
    memcpy(fragment + FRAG_HEADER_SIZE, data.data() + offset, chunk);
    send_fragment(fragment, FRAG_HEADER_SIZE + chunk);
    delay(8);
  }

  InterfaceImpl::handle_outgoing(data);
}

void TBeamSupremeBleInterface::send_fragment(const uint8_t* data, size_t len) {
  if (_central_role) {
    if (_remote_rx_characteristic) {
      _remote_rx_characteristic->writeValue((uint8_t*)data, len, false);
    }
  } else {
    if (_tx_characteristic) {
      _tx_characteristic->setValue((uint8_t*)data, len);
      _tx_characteristic->notify();
    }
  }
}

void TBeamSupremeBleInterface::handle_fragment(const uint8_t* data, size_t len) {
  if (len < FRAG_HEADER_SIZE) {
    WARNINGF("BLE fragment too short: %d", (int)len);
    return;
  }

  uint8_t fragment_type = data[0];
  uint16_t sequence = ((uint16_t)data[1] << 8) | data[2];
  uint16_t total = ((uint16_t)data[3] << 8) | data[4];
  const uint8_t* payload = data + FRAG_HEADER_SIZE;
  size_t payload_len = len - FRAG_HEADER_SIZE;

  if ((fragment_type != FRAG_START && fragment_type != FRAG_CONTINUE && fragment_type != FRAG_END) ||
      total == 0 || sequence >= total) {
    WARNING("BLE invalid fragment header");
    reset_reassembly();
    return;
  }

  if (sequence == 0) {
    reset_reassembly();
    _expected_total = total;
    _received_fragments = 0;
    _reassembly_started_ms = millis();
  }

  if (_expected_total != total || sequence != _received_fragments) {
    WARNING("BLE out-of-order fragment; dropping partial packet");
    reset_reassembly();
    return;
  }

  _reassembly_buffer.append(payload, payload_len);
  _received_fragments++;

  if (_received_fragments == _expected_total) {
    enqueue_packet(_reassembly_buffer);
    reset_reassembly();
  }
}

void TBeamSupremeBleInterface::enqueue_packet(const RNS::Bytes& packet) {
  portENTER_CRITICAL(&_queue_mux);
  _incoming_packets.push_back(packet);
  portEXIT_CRITICAL(&_queue_mux);
}

bool TBeamSupremeBleInterface::dequeue_packet(RNS::Bytes& packet) {
  portENTER_CRITICAL(&_queue_mux);
  if (_incoming_packets.empty()) {
    portEXIT_CRITICAL(&_queue_mux);
    return false;
  }
  packet = _incoming_packets.front();
  _incoming_packets.pop_front();
  portEXIT_CRITICAL(&_queue_mux);
  return true;
}

void TBeamSupremeBleInterface::reset_reassembly() {
  _reassembly_buffer.clear();
  _expected_total = 0;
  _received_fragments = 0;
  _reassembly_started_ms = 0;
}

RNS::Bytes TBeamSupremeBleInterface::local_identity_hash() const {
  String material = String("microReticulum BLE ") + NODE_LABEL;
  return Identity::truncated_hash(RNS::bytesFromString(material.c_str()));
}

void TBeamSupremeBleInterface::client_notify_callback(BLERemoteCharacteristic* characteristic,
                                                      uint8_t* data,
                                                      size_t length,
                                                      bool is_notify) {
  (void)characteristic;
  (void)is_notify;
  if (active_ble_interface) {
    active_ble_interface->handle_fragment(data, length);
  }
}
