#pragma once

/*
 * DataNet.h — SJS DataNet Arduino/ESP32/ESP8266 SDK
 *
 * Dependencies (install via Arduino Library Manager):
 *   - ArduinoJson by Benoit Blanchon, v7+
 *   - WebSockets by Markus Sattler (Links2004), v2.4+
 *   - HTTPClient (built into ESP32/ESP8266 Arduino cores)
 *
 * Platform support: ESP32, ESP8266
 */

#include <Arduino.h>
#include <ArduinoJson.h>
#include <WebSocketsClient.h>

#ifdef ESP32
  #include <WiFi.h>
  #include <HTTPClient.h>
  #include <WiFiClientSecure.h>
#elif defined(ESP8266)
  #include <ESP8266WiFi.h>
  #include <ESP8266HTTPClient.h>
  #include <WiFiClientSecureBearSSL.h>
#else
  #error "DataNet SDK requires ESP32 or ESP8266"
#endif

// ---------------------------------------------------------------------------
// Compile-time configuration — override with -D flags or before including
// ---------------------------------------------------------------------------
#ifndef DATANET_MAX_SUBS
  #define DATANET_MAX_SUBS 8
#endif

#ifndef DATANET_MAX_EVENT_HANDLERS
  #define DATANET_MAX_EVENT_HANDLERS 4
#endif

#ifndef DATANET_JWT_BUF_SIZE
  #define DATANET_JWT_BUF_SIZE 512
#endif

#ifndef DATANET_HEARTBEAT_INTERVAL_MS
  #define DATANET_HEARTBEAT_INTERVAL_MS 30000UL
#endif

#ifndef DATANET_RECONNECT_BASE_MS
  #define DATANET_RECONNECT_BASE_MS 1000UL
#endif

#ifndef DATANET_RECONNECT_MAX_MS
  #define DATANET_RECONNECT_MAX_MS 60000UL
#endif

// After this many reconnects, re-fetch a fresh JWT
#ifndef DATANET_JWT_REFRESH_AFTER_RECONNECTS
  #define DATANET_JWT_REFRESH_AFTER_RECONNECTS 3
#endif

// ---------------------------------------------------------------------------
// Handler typedefs
// ---------------------------------------------------------------------------

// MessageHandler: called when a "pub" envelope arrives on a subscribed channel.
// channel — the channel name string
// data    — the "d" field from the envelope (may be any JSON type)
typedef void (*MessageHandler)(const char* channel, JsonVariant data);

// EventHandler: called for lifecycle events "connect", "disconnect", "error"
// event   — event name string
// info    — optional detail string (error message, disconnect reason, etc.)
typedef void (*EventHandler)(const char* event, const char* info);

// ---------------------------------------------------------------------------
// DataNet
// ---------------------------------------------------------------------------
class DataNet {
public:
    // -----------------------------------------------------------------------
    // Constructor
    // -----------------------------------------------------------------------
    DataNet(
        const char* apiKey,
        const char* apiUrl  = "https://api.datanet.art",
        const char* wsHost  = "ws.datanet.art",
        int         wsPort  = 443
    );

    // -----------------------------------------------------------------------
    // Connection
    // -----------------------------------------------------------------------

    // Fetches JWT via HTTPS POST then opens WSS connection.
    // WiFi must be connected before calling.
    // Returns true on success, false on failure.
    bool connect();

    // Must be called in Arduino loop() — drives WebSocket callbacks and heartbeat.
    void loop();

    // Returns true if the WebSocket is currently connected.
    bool connected();

    // -----------------------------------------------------------------------
    // Pub / Sub
    // -----------------------------------------------------------------------

    // Subscribe to a channel. handler is called when a "pub" arrives.
    // Up to DATANET_MAX_SUBS subscriptions allowed.
    void subscribe(const char* channel, MessageHandler handler);

    // Remove a subscription by channel name.
    void unsubscribe(const char* channel);

    // Returns the server-side timestamp (Unix ms) of the last message received
    // on the given channel, or 0 if no message has been received yet.
    // Useful for detecting stale readings or measuring latency.
    uint64_t getLastTimestamp(const char* channel);

    // Publish a JsonVariant as the "d" field of a pub envelope.
    // Returns true if the message was sent.
    bool publish(const char* channel, JsonVariant data);

    // Convenience: publish {"<key>": <value>} as a float
    bool publishFloat(const char* channel, const char* key, float value);

    // Convenience: publish {"<key>": "<value>"} as a string
    bool publishString(const char* channel, const char* key, const char* value);

    // -----------------------------------------------------------------------
    // Event handlers
    // -----------------------------------------------------------------------

    // Register a lifecycle event handler.
    // Supported events: "connect", "disconnect", "error"
    // Up to DATANET_MAX_EVENT_HANDLERS handlers per event type.
    void on(const char* event, EventHandler handler);

private:
    // -----------------------------------------------------------------------
    // Internal types
    // -----------------------------------------------------------------------
    struct Subscription {
        char           channel[64];
        MessageHandler handler;
        bool           active;
        uint64_t       lastTs;   // server-side Unix ms timestamp of last pub
    };

    enum class EventType : uint8_t {
        Connect    = 0,
        Disconnect = 1,
        Error      = 2,
        COUNT      = 3
    };

    // -----------------------------------------------------------------------
    // Configuration (stored as pointers to caller-owned strings)
    // -----------------------------------------------------------------------
    const char* _apiKey;
    const char* _apiUrl;
    const char* _wsHost;
    int         _wsPort;

    // -----------------------------------------------------------------------
    // State
    // -----------------------------------------------------------------------
    char        _jwt[DATANET_JWT_BUF_SIZE];
    bool        _wsConnected;
    uint32_t    _lastHeartbeatMs;
    uint32_t    _reconnectAtMs;
    uint8_t     _reconnectCount;
    bool        _reconnectPending;

    // -----------------------------------------------------------------------
    // Subscriptions
    // -----------------------------------------------------------------------
    Subscription _subs[DATANET_MAX_SUBS];

    // -----------------------------------------------------------------------
    // Event handlers storage
    // -----------------------------------------------------------------------
    EventHandler _eventHandlers[static_cast<uint8_t>(EventType::COUNT)][DATANET_MAX_EVENT_HANDLERS];

    // -----------------------------------------------------------------------
    // WebSocket client (Links2004 library)
    // -----------------------------------------------------------------------
    WebSocketsClient _ws;

    // Static trampoline so we can pass a C-style function pointer to the
    // WebSocketsClient library while still dispatching to the instance.
    // Only one DataNet instance is expected per sketch (typical for embedded).
    static DataNet* _instance;
    static void _wsEventCallback(WStype_t type, uint8_t* payload, size_t length);

    // Instance-level WebSocket event handler called from the static trampoline.
    void _handleWsEvent(WStype_t type, uint8_t* payload, size_t length);

    // -----------------------------------------------------------------------
    // Private helpers
    // -----------------------------------------------------------------------

    // Fetch a fresh JWT from the REST API. Returns true on success.
    bool _fetchJwt();

    // Open WSS connection using the stored JWT.
    void _openWebSocket();

    // Handle an incoming WebSocket text frame (raw JSON string).
    void _handleMessage(const char* json, size_t length);

    // Send a raw JSON string over the WebSocket.
    bool _sendJson(const char* json);

    // Send an already-serialised envelope string.
    bool _sendEnvelope(const char* op, const char* channel = nullptr, JsonVariant data = JsonVariant());

    // Re-send "sub" envelopes for all active subscriptions (called on reconnect).
    void _resubscribeAll();

    // Schedule the next reconnect attempt with exponential backoff.
    void _scheduleReconnect();

    // Dispatch a lifecycle event to all registered handlers.
    void _dispatchEvent(EventType type, const char* info = "");

    // Map a string event name to EventType. Returns COUNT on unknown name.
    static EventType _parseEventType(const char* name);
};