#include <stdio.h>
#include <string.h>
#include "mqtt_connect_diag.h"
#include "sockets.h"

/* Only installed during synchronous CONNECT. Never dump the CONNECT payload:
 * it contains credentials. Keep only the first four incoming CONNACK bytes. */
static struct {
    int (*read)(Network *, unsigned char *, int, int);
    int (*write)(Network *, unsigned char *, int, int);
    unsigned int read_calls, write_calls, received, sent;
    unsigned char prefix[4];
    unsigned int prefix_size;
} mqtt_diag;

static unsigned int mqtt_diag_elapsed(unsigned int start)
{
    return (unsigned int)(ql_rtos_get_systicks() - start) * 5U;
}

static int mqtt_diag_read(Network *network, unsigned char *buffer, int size, int timeout)
{
    unsigned int start = ql_rtos_get_systicks();
    int result = mqtt_diag.read(network, buffer, size, timeout);
    int socket_error = result <= 0 ? ql_get_soc_errno() : 0;
    int i;
    ++mqtt_diag.read_calls;
    if (result > 0) {
        mqtt_diag.received += (unsigned int)result;
        for (i = 0; i < result && i < size && mqtt_diag.prefix_size < sizeof(mqtt_diag.prefix); ++i)
            mqtt_diag.prefix[mqtt_diag.prefix_size++] = buffer[i];
    }
    printf("[MQTT-DIAG] read request=%d rc=%d timeout_ms=%d elapsed_ms=%u errno_snapshot=%d\r\n",
           size, result, timeout, mqtt_diag_elapsed(start), socket_error);
    return result;
}

static int mqtt_diag_write(Network *network, unsigned char *buffer, int size, int timeout)
{
    unsigned int start = ql_rtos_get_systicks();
    int result = mqtt_diag.write(network, buffer, size, timeout);
    int socket_error = result <= 0 ? ql_get_soc_errno() : 0;
    ++mqtt_diag.write_calls;
    if (result > 0) mqtt_diag.sent += (unsigned int)result;
    printf("[MQTT-DIAG] write request=%d rc=%d timeout_ms=%d elapsed_ms=%u errno_snapshot=%d\r\n",
           size, result, timeout, mqtt_diag_elapsed(start), socket_error);
    return result;
}

static const char *mqtt_diag_connack_reason(unsigned char code)
{
    switch (code) {
    case 0: return "accepted";
    case 1: return "protocol_version_rejected";
    case 2: return "client_id_rejected";
    case 3: return "server_unavailable";
    case 4: return "bad_username_or_password";
    case 5: return "not_authorized";
    default: return "reserved_code";
    }
}

int talkingq_mqtt_connect_diagnosed(MQTTClient *client, MQTTPacket_connectData *options)
{
    Network *network = client->ipstack;
    MQTTConnackData connack = {255, 255};
    unsigned int start;
    int result;
    int valid_connack;
    memset(&mqtt_diag, 0, sizeof(mqtt_diag));
    mqtt_diag.read = network->mqttread;
    mqtt_diag.write = network->mqttwrite;
    if (!mqtt_diag.read || !mqtt_diag.write) {
        printf("[MQTT-DIAG] missing transport callbacks; calling SDK unchanged\r\n");
        return MQTTConnectWithResults(client, options, &connack);
    }
    printf("[MQTT-DIAG] CONNECT begin version=%u clean=%u keepalive_s=%u timeout_ms=%u already_connected=%d\r\n",
           (unsigned int)options->MQTTVersion, (unsigned int)options->cleansession,
           (unsigned int)options->keepAliveInterval, client->command_timeout_ms, client->isconnected);
    network->mqttread = mqtt_diag_read;
    network->mqttwrite = mqtt_diag_write;
    start = ql_rtos_get_systicks();
    result = MQTTConnectWithResults(client, options, &connack);
    network->mqttread = mqtt_diag.read;
    network->mqttwrite = mqtt_diag.write;
    printf("[MQTT-DIAG] CONNECT end rc=%d elapsed_ms=%u write_calls=%u tx_bytes=%u read_calls=%u rx_bytes=%u\r\n",
           result, mqtt_diag_elapsed(start), mqtt_diag.write_calls, mqtt_diag.sent,
           mqtt_diag.read_calls, mqtt_diag.received);
    valid_connack = mqtt_diag.prefix_size == 4 && mqtt_diag.prefix[0] == 0x20 &&
                    mqtt_diag.prefix[1] == 2 && (mqtt_diag.prefix[2] & 0xfe) == 0 &&
                    mqtt_diag.prefix[3] <= 5 &&
                    (mqtt_diag.prefix[3] == 0 || mqtt_diag.prefix[2] == 0);
    if (valid_connack) {
        printf("[MQTT-DIAG] CONNACK code=%u reason=%s session_present=%u\r\n",
               (unsigned int)mqtt_diag.prefix[3], mqtt_diag_connack_reason(mqtt_diag.prefix[3]),
               (unsigned int)mqtt_diag.prefix[2]);
    } else if (mqtt_diag.received) {
        printf("[MQTT-DIAG] no valid initial CONNACK: received_prefix_bytes=%u\r\n", mqtt_diag.prefix_size);
    } else {
        printf("[MQTT-DIAG] no response bytes returned by SDK read\r\n");
    }
    if (result != 0 && !mqtt_diag.write_calls)
        printf("[MQTT-DIAG] failed before transport write; inspect client state/CONNECT serialization\r\n");
    /* SDK read maps both EOF and some socket errors to -1. The errno snapshot
     * can be stale; do not claim a broker rejection without a valid CONNACK. */
    return result;
}