#include <cerrno>
#include "mbedtls/net_sockets.h"
#include "lwip/sockets.h"
#include "esp_heap_caps.h"
#include "esp_timer.h"
#include "esp_log.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

struct NetRecord { int64_t us, duration; int kind, a, b, c, d; };
static NetRecord *records;
static unsigned count;
static bool capturing=true;
static portMUX_TYPE trace_lock=portMUX_INITIALIZER_UNLOCKED;
extern "C" void speech_diag_net_record(int kind, int64_t us, int64_t duration, int a,int b,int c,int d) {
    if (!records || !capturing) return;
    portENTER_CRITICAL(&trace_lock);
    if(count<2048) records[count++]={us,duration,kind,a,b,c,d};
    portEXIT_CRITICAL(&trace_lock);
}
static void dump(void *) {
    portENTER_CRITICAL(&trace_lock);capturing=false;unsigned n=count;portEXIT_CRITICAL(&trace_lock);
    ESP_LOGI("NETDIAG","[DIAG_NET] begin count=%u capacity=2048",n);
    for(unsigned i=0;i<n;++i) {
        const auto &r=records[i];
        ESP_LOGI("NETDIAG","[DIAG_NET] n=%u kind=%d us=%lld dur_us=%lld a=%d b=%d c=%d d=%d",
            i,r.kind,(long long)r.us,(long long)r.duration,r.a,r.b,r.c,r.d);
    }
    ESP_LOGI("NETDIAG","[DIAG_NET] end count=%u",n);
}
extern "C" void speech_diag_net_init() {
    records=(NetRecord*)heap_caps_calloc(2048,sizeof(NetRecord),MALLOC_CAP_SPIRAM);
    ESP_LOGI("NETDIAG","[DIAG_NET] init psram=%d dump_after_ms=115000",records!=nullptr);
    esp_timer_handle_t timer;
    const esp_timer_create_args_t args={.callback=dump,.arg=nullptr,.dispatch_method=ESP_TIMER_TASK,.name="netdiag",.skip_unhandled_events=false};
    ESP_ERROR_CHECK(esp_timer_create(&args,&timer));
    ESP_ERROR_CHECK(esp_timer_start_once(timer,115000000));
}
extern "C" int __real_mbedtls_net_recv(void*,unsigned char*,size_t);
extern "C" int __wrap_mbedtls_net_recv(void *context,unsigned char *buffer,size_t len) {
    int pending=-1;
    if(context) ioctl(((mbedtls_net_context*)context)->fd,FIONREAD,&pending);
    const int64_t start=esp_timer_get_time();
    const int ret=__real_mbedtls_net_recv(context,buffer,len);
    const int saved_errno=errno;
    speech_diag_net_record(0,start,esp_timer_get_time()-start,(int)len,ret,pending,ret<0?saved_errno:0);
    errno=saved_errno;
    return ret;
}