#include "speech_transport_trace.h"

#include <atomic>
#include <cerrno>
#include <cstdint>
#include "esp_log.h"
#include "esp_timer.h"
#include "esp_tls.h"
#include "esp_transport.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "lwip/sockets.h"
#include "mbedtls/net_sockets.h"

namespace {
constexpr uint32_t WINDOW_MS = 120000, SLOW_MS = 250, SAMPLE_EVERY = 32;
constexpr unsigned RING_SIZE = 16;
enum Stage : unsigned { POLL, TRANSPORT, TLS, SOCKET, STAGE_COUNT };
constexpr const char *NAMES[] = {"poll", "transport_read", "tls_read", "socket_recv"};
struct Counters {
    uint32_t calls = 0, slow = 0, max_ms = 0;
    uint64_t bytes = 0;
    int64_t last_finish_us = 0;
};
struct Record {
    uint32_t epoch, op, parent, stage, calls, slow, max_ms;
    uint64_t bytes;
    int64_t start_us, elapsed_us, gap_us;
    uintptr_t handle;
    int requested, result, error, fd, socket_before, socket_after, tls_before;
    unsigned priority, core;
};

// One producer (the WebSocket receive task), one consumer (input_sender).
// Fixed storage; no allocation, payload copy, or logging in a read wrapper.
Record s_records[RING_SIZE];
std::atomic<unsigned> s_write{0}, s_read{0}, s_dropped{0};
std::atomic<uint32_t> s_epoch{0}, s_until_ms{0};
std::atomic<TaskHandle_t> s_receive_task{nullptr};
Counters s_counters[STAGE_COUNT];
uint32_t s_seen_epoch = 0, s_op = 0, s_parent = 0;

int socket_available(int fd)
{
    int bytes = -1;
    if (fd >= 0 && ioctl(fd, FIONREAD, &bytes) != 0) bytes = -1;
    return bytes;
}

void enqueue(const Record &record)
{
    const unsigned write = s_write.load(std::memory_order_relaxed);
    const unsigned read = s_read.load(std::memory_order_acquire);
    if (write - read >= RING_SIZE) {
        s_dropped.fetch_add(1, std::memory_order_relaxed);
        return;
    }
    s_records[write % RING_SIZE] = record;
    s_write.store(write + 1, std::memory_order_release);
}

class TimedRead {
public:
    TimedRead(Stage stage, uintptr_t handle, int requested, int fd, int tls_before = -1)
    {
        const int saved_errno = errno;
        const uint32_t now_ms = static_cast<uint32_t>(esp_timer_get_time() / 1000);
        const uint32_t until = s_until_ms.load(std::memory_order_acquire);
        if (xTaskGetCurrentTaskHandle() != s_receive_task.load(std::memory_order_acquire) ||
            !until || static_cast<int32_t>(until - now_ms) <= 0) {
            errno = saved_errno;
            return;
        }
        active = true;
        r.epoch = s_epoch.load(std::memory_order_acquire);
        if (r.epoch != s_seen_epoch) {
            for (auto &counter : s_counters) counter = {};
            s_seen_epoch = r.epoch;
        }
        r.op = ++s_op;
        r.parent = s_parent;
        s_parent = r.op;
        r.stage = stage;
        r.handle = handle;
        r.requested = requested;
        r.fd = fd;
        r.tls_before = tls_before;
        r.socket_before = socket_available(fd);
        r.priority = uxTaskPriorityGet(nullptr);
        r.core = xPortGetCoreID();
        r.start_us = esp_timer_get_time();
        const auto previous = s_counters[stage].last_finish_us;
        r.gap_us = previous ? r.start_us - previous : 0;
        errno = saved_errno;
    }

    void finish(int result)
    {
        if (!active) return;
        const int saved_errno = errno;
        const int64_t finish_us = esp_timer_get_time();
        r.elapsed_us = finish_us - r.start_us;
        r.result = result;
        r.error = result < 0 ? saved_errno : 0;
        r.socket_after = socket_available(r.fd);
        auto &counter = s_counters[r.stage];
        ++counter.calls;
        if (result > 0 && r.stage != POLL) counter.bytes += result;
        if (r.elapsed_us >= SLOW_MS * 1000) ++counter.slow;
        const uint32_t elapsed_ms = static_cast<uint32_t>(r.elapsed_us / 1000);
        if (elapsed_ms > counter.max_ms) counter.max_ms = elapsed_ms;
        counter.last_finish_us = finish_us;
        r.calls = counter.calls;
        r.bytes = counter.bytes;
        r.max_ms = counter.max_ms;
        r.slow = counter.slow;
        if (r.elapsed_us >= SLOW_MS * 1000 || r.gap_us >= SLOW_MS * 1000 ||
            result < 0 || counter.calls == 1 || counter.calls % SAMPLE_EVERY == 0)
            enqueue(r);
        s_parent = r.parent;
        errno = saved_errno;
    }
private:
    bool active = false;
    Record r{};
};
}

void speech_transport_trace_receive_task()
{
    s_receive_task.store(xTaskGetCurrentTaskHandle(), std::memory_order_release);
}

void speech_transport_trace_arm()
{
    const auto epoch = s_epoch.fetch_add(1, std::memory_order_acq_rel) + 1;
    s_until_ms.store(static_cast<uint32_t>(esp_timer_get_time() / 1000) + WINDOW_MS,
                     std::memory_order_release);
    ESP_LOGI("CLOUD", "[TRANSPORT_RX_ARM] epoch=%lu at_ms=%lld window_ms=%lu",
             static_cast<unsigned long>(epoch), static_cast<long long>(esp_timer_get_time()/1000),
             static_cast<unsigned long>(WINDOW_MS));
}

void speech_transport_trace_flush()
{
    for (unsigned count = 0; count < 4; ++count) {
        const unsigned read = s_read.load(std::memory_order_relaxed);
        if (read == s_write.load(std::memory_order_acquire)) break;
        const Record r = s_records[read % RING_SIZE];
        s_read.store(read + 1, std::memory_order_release);
        ESP_LOGI("CLOUD", "[TRANSPORT_RX] epoch=%lu op=%lu parent=%lu stage=%s handle=%lx start_ms=%lld elapsed_ms=%lld gap_ms=%lld requested=%d result=%d errno=%d fd=%d socket_before=%d socket_after=%d tls_before=%d calls=%lu bytes=%llu slow=%lu max_ms=%lu priority=%u core=%u",
                 static_cast<unsigned long>(r.epoch), static_cast<unsigned long>(r.op),
                 static_cast<unsigned long>(r.parent), NAMES[r.stage], static_cast<unsigned long>(r.handle),
                 static_cast<long long>(r.start_us/1000), static_cast<long long>(r.elapsed_us/1000),
                 static_cast<long long>(r.gap_us/1000), r.requested, r.result, r.error, r.fd,
                 r.socket_before, r.socket_after, r.tls_before, static_cast<unsigned long>(r.calls),
                 static_cast<unsigned long long>(r.bytes), static_cast<unsigned long>(r.slow),
                 static_cast<unsigned long>(r.max_ms), r.priority, r.core);
    }
    const unsigned dropped = s_dropped.exchange(0, std::memory_order_relaxed);
    if (dropped) ESP_LOGW("CLOUD", "[TRANSPORT_RX_DROPPED] metadata_records=%u", dropped);
}

extern "C" {
int __real_esp_transport_read(esp_transport_handle_t, char *, int, int);
int __real_esp_transport_poll_read(esp_transport_handle_t, int);
ssize_t __real_esp_tls_conn_read(esp_tls_t *, void *, size_t);
int __real_mbedtls_net_recv(void *, unsigned char *, size_t);

int __wrap_esp_transport_read(esp_transport_handle_t t, char *buffer, int len, int timeout_ms)
{
    const int saved_errno = errno;
    const int fd = esp_transport_get_socket(t);
    errno = saved_errno;
    TimedRead trace(TRANSPORT, reinterpret_cast<uintptr_t>(t), len, fd);
    const int result = __real_esp_transport_read(t, buffer, len, timeout_ms);
    trace.finish(result);
    return result;
}

int __wrap_esp_transport_poll_read(esp_transport_handle_t t, int timeout_ms)
{
    const int saved_errno = errno;
    const int fd = esp_transport_get_socket(t);
    errno = saved_errno;
    TimedRead trace(POLL, reinterpret_cast<uintptr_t>(t), timeout_ms, fd);
    const int result = __real_esp_transport_poll_read(t, timeout_ms);
    trace.finish(result);
    return result;
}

ssize_t __wrap_esp_tls_conn_read(esp_tls_t *tls, void *buffer, size_t len)
{
    const int saved_errno = errno;
    int fd = -1;
    esp_tls_get_conn_sockfd(tls, &fd);
    const int buffered = static_cast<int>(esp_tls_get_bytes_avail(tls));
    errno = saved_errno;
    TimedRead trace(TLS, reinterpret_cast<uintptr_t>(tls), static_cast<int>(len), fd, buffered);
    const ssize_t result = __real_esp_tls_conn_read(tls, buffer, len);
    trace.finish(static_cast<int>(result));
    return result;
}

int __wrap_mbedtls_net_recv(void *ctx, unsigned char *buffer, size_t len)
{
    const auto *net = static_cast<mbedtls_net_context *>(ctx);
    TimedRead trace(SOCKET, reinterpret_cast<uintptr_t>(ctx), static_cast<int>(len), net ? net->fd : -1);
    const int result = __real_mbedtls_net_recv(ctx, buffer, len);
    trace.finish(result);
    return result;
}
}