#include "websocket_client.h"
#include "sensors.h"
#include "protocol.h"
#include "commands.h"
#include <ArduinoWebsockets.h>
#include <WiFi.h>

using namespace websockets;

WiFiSettings wifiSettings;  // Global runtime instance

// ============================================================================
// Outbound client (connects to Radxa for FACE/ALIV)
// ============================================================================

static WebsocketsClient client;
static bool s_connected = false;
static unsigned long lastReconnectAttempt = 0;
constexpr unsigned long RECONNECT_INTERVAL = 5000;

static String buildWsUrl() {
  return "ws://" + String(wifiSettings.host) + ":" + String(wifiSettings.port) + String(wifiSettings.path);
}

// ============================================================================
// Inbound server (accepts commands from computer)
// ============================================================================

static WebsocketsServer* server = nullptr;
constexpr uint8_t MAX_SERVER_CLIENTS = 2;
static WebsocketsClient serverClients[MAX_SERVER_CLIENTS];
static bool serverClientActive[MAX_SERVER_CLIENTS] = {false};

// ============================================================================
// Shared protocol packet parser
// ============================================================================

// Parse protocol-framed binary data, call handler for each valid packet
static void parseProtocolPackets(const uint8_t* data, size_t len,
    void (*handler)(const char tag[4], const uint8_t* payload, uint16_t payloadLen))
{
  size_t pos = 0;
  
  while (pos + 12 <= len) {
    if (data[pos] != 0xA5 || data[pos + 1] != 0x5A) {
      pos++;
      continue;
    }
    
    char tag[4];
    memcpy(tag, &data[pos + 2], 4);
    
    uint16_t payloadLen = data[pos + 6] | (data[pos + 7] << 8);
    
    size_t totalPacketLen = 2 + 4 + 2 + 2 + payloadLen + 2;
    if (pos + totalPacketLen > len) break;
    
    const uint8_t* payload = &data[pos + 10];
    
    // Verify CRC
    const uint8_t* crcData = &data[pos + 2];
    uint16_t crcDataLen = 4 + 2 + 2 + payloadLen;
    uint16_t computed = crc16Compute(crcData, crcDataLen);
    uint16_t received = data[pos + 10 + payloadLen] | (data[pos + 10 + payloadLen + 1] << 8);
    
    if (computed == received) {
      handler(tag, payload, payloadLen);
    }
    
    pos += totalPacketLen;
  }
}

// ============================================================================
// Client packet handler (FACE/ALIV from Radxa)
// ============================================================================

static void handleClientPacket(const char tag[4], const uint8_t* payload, uint16_t len) {
  if (memcmp(tag, Tag::FACE, 4) == 0) {
    faceDetect.feedPayload(payload, len);
  }
  else if (memcmp(tag, Tag::ALIV, 4) == 0) {
    if (len >= 2) {
      uint8_t componentId = payload[0];
      uint8_t alive = payload[1];
      if (componentId == 3) {
        faceDetect.setAlive(alive != 0);
      }
    }
    sendPacket(Tag::ALIV, payload, len);
  }
}

// ============================================================================
// Server packet handler (any command from computer)
// ============================================================================

static void handleServerPacket(const char tag[4], const uint8_t* payload, uint16_t len) {
  // Route to the standard command dispatcher - handles VSME, MSET, BHVR, SSET, etc.
  dispatchCommand(tag, payload, len);
}

// ============================================================================
// Client callbacks
// ============================================================================

static void onClientMessage(WebsocketsMessage message) {
  std::string raw = message.rawData();
  if (raw.empty()) return;
  parseProtocolPackets(reinterpret_cast<const uint8_t*>(raw.data()), raw.size(), handleClientPacket);
}

static void onClientEvent(WebsocketsEvent event, String data) {
  switch (event) {
    case WebsocketsEvent::ConnectionOpened:
      s_connected = true;
      Serial.println("[WS Client] Connected");
      break;
    case WebsocketsEvent::ConnectionClosed:
      s_connected = false;
      Serial.println("[WS Client] Disconnected");
      break;
    default: break;
  }
}

// ============================================================================
// Server client callbacks (created per-client via lambda in loop)
// ============================================================================

static void onServerMessage(WebsocketsMessage message) {
  std::string raw = message.rawData();
  if (raw.empty()) return;
  parseProtocolPackets(reinterpret_cast<const uint8_t*>(raw.data()), raw.size(), handleServerPacket);
}

static void onServerEvent(WebsocketsEvent event, String data) {
  // Connection events are handled in the accept logic
}

// ============================================================================
// Public API
// ============================================================================

void websocketSetup() {
  Serial.print("[WiFi] SSID: '");
  Serial.print(wifiSettings.ssid);
  Serial.println("'");
  
  WiFi.mode(WIFI_STA);
  WiFi.begin(wifiSettings.ssid, wifiSettings.password);
  Serial.print("[WiFi] Connecting");
  int attempts = 0;
  while (WiFi.status() != WL_CONNECTED && attempts < 30) {
    delay(500);
    Serial.print(".");
    attempts++;
  }
  Serial.println();
  if (WiFi.status() != WL_CONNECTED) {
    Serial.println("[WiFi] Failed to connect");
    return;
  }
  Serial.print("[WiFi] OK ");
  Serial.println(WiFi.localIP());

  // --- Start outbound client (to Radxa) ---
  client.onMessage(onClientMessage);
  client.onEvent(onClientEvent);

  String url = buildWsUrl();
  Serial.println("[WS Client] Connecting to " + url);
  if (client.connect(url)) {
    s_connected = true;
    Serial.println("[WS Client] Connected");
  } else {
    Serial.println("[WS Client] Connect failed (will retry)");
  }

  // --- Start inbound server (allocated after WiFi to avoid early init issues) ---
  server = new WebsocketsServer();
  server->listen(WS_SERVER_PORT);
  if (server->available()) {
    Serial.print("[WS Server] Listening on port ");
    Serial.println(WS_SERVER_PORT);
  } else {
    Serial.println("[WS Server] Failed to start");
  }
}

void websocketLoop() {
  if (WiFi.status() != WL_CONNECTED) {
    s_connected = false;
    return;
  }

  // --- Outbound client ---
  if (!s_connected && !client.available()) {
    unsigned long now = millis();
    if (now - lastReconnectAttempt >= RECONNECT_INTERVAL) {
      lastReconnectAttempt = now;
      String url = buildWsUrl();
      if (client.connect(url)) {
        s_connected = true;
        Serial.println("[WS Client] Reconnected");
      }
    }
  }

  if (client.available()) {
    client.poll();
  }

  // --- Inbound server: accept new connections ---
  if (server && server->available() && server->poll()) {
    WebsocketsClient newClient = server->accept();
    if (newClient.available()) {
      // Find an empty slot
      bool accepted = false;
      for (uint8_t i = 0; i < MAX_SERVER_CLIENTS; i++) {
        if (!serverClientActive[i] || !serverClients[i].available()) {
          serverClients[i] = std::move(newClient);
          serverClients[i].onMessage(onServerMessage);
          serverClients[i].onEvent(onServerEvent);
          serverClientActive[i] = true;
          Serial.print("[WS Server] Client connected (slot ");
          Serial.print(i);
          Serial.println(")");
          accepted = true;
          break;
        }
      }
      if (!accepted) {
        Serial.println("[WS Server] Max clients reached, rejecting");
        newClient.close();
      }
    }
  }

  // --- Inbound server: poll existing clients ---
  for (uint8_t i = 0; i < MAX_SERVER_CLIENTS; i++) {
    if (serverClientActive[i]) {
      if (serverClients[i].available()) {
        serverClients[i].poll();
      } else {
        serverClientActive[i] = false;
        Serial.print("[WS Server] Client disconnected (slot ");
        Serial.print(i);
        Serial.println(")");
      }
    }
  }
}

bool websocketConnected() {
  return s_connected && client.available();
}

void websocketReconnect() {
  client.close();
  s_connected = false;
  
  WiFi.disconnect();
  WiFi.begin(wifiSettings.ssid, wifiSettings.password);
  Serial.print("[WiFi] Reconnecting to ");
  Serial.println(wifiSettings.ssid);
  
  lastReconnectAttempt = 0;
}

uint8_t websocketServerClientCount() {
  uint8_t count = 0;
  for (uint8_t i = 0; i < MAX_SERVER_CLIENTS; i++) {
    if (serverClientActive[i] && serverClients[i].available()) count++;
  }
  return count;
}