#include <atomic>
#include "esp_netif.h"
#include "lwip/api.h"
#include "lwip/tcpip.h"
#include "lwip/tcp.h"
#include "lwip/priv/tcp_priv.h"
extern "C" {
#include "lwip/priv/sockets_priv.h"
}
#include <cerrno>
#include <cstring>
#include "lwip/netdb.h"
#include "lwip/inet.h"
#include "../../.cache/speech-diagnostics/ip_choice.h"
#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 std::atomic<int> trace_fd{-1};
static std::atomic<int> rx_live{0},rx_max{0},rx_total{0},rx_freed{0};
extern "C" void speech_diag_net_socket(int fd) {trace_fd=fd;}
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<8192) 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=8192",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);
}
static void alloc_failed(size_t size,uint32_t caps,const char *) {
    speech_diag_net_record(7,esp_timer_get_time(),0,(int)size,(int)caps,0,0);
}
static void tcp_snapshot(void *) {
    if(!capturing)return;
    const int fd=trace_fd.load();if(fd<0)return;
    const auto *socket=lwip_socket_dbg_get_socket(fd);
    if(!socket || !socket->conn || NETCONNTYPE_GROUP(socket->conn->type)!=NETCONN_TCP)return;
    const auto *pcb=socket->conn->pcb.tcp;
    if(!pcb || pcb->remote_port!=443)return;
    unsigned oo_count=0,oo_bytes=0;
#if TCP_QUEUE_OOSEQ
    for(const tcp_seg *seg=pcb->ooseq;seg && oo_count<128;seg=seg->next){++oo_count;oo_bytes+=seg->len;}
#endif
    int available=-1;
#if LWIP_SO_RCVBUF
    available=socket->conn->recv_avail;
#endif
    const int64_t now=esp_timer_get_time();
    speech_diag_net_record(12,now,0,rx_live.load(),rx_max.load(),rx_total.load(),rx_freed.load());
    speech_diag_net_record(8,now,0,fd,(int)pcb->rcv_wnd,available,(int)socket->rcvevent);
    speech_diag_net_record(9,now,0,(int)pcb->rcv_nxt,(int)oo_count,(int)oo_bytes,(int)pcb->rcv_ann_wnd);
    speech_diag_net_record(10,now,0,(int)heap_caps_get_free_size(MALLOC_CAP_INTERNAL),(int)heap_caps_get_largest_free_block(MALLOC_CAP_INTERNAL),(int)pcb->state,(int)pcb->flags);
}
static void request_snapshot(void *) {if(capturing && trace_fd>=0)tcpip_try_callback(tcp_snapshot,nullptr);}
extern "C" void speech_diag_net_init() {
    records=(NetRecord*)heap_caps_calloc(8192,sizeof(NetRecord),MALLOC_CAP_SPIRAM);
    ESP_ERROR_CHECK(heap_caps_register_failed_alloc_callback(alloc_failed));
    ESP_LOGI("NETDIAG","[DIAG_NET] init psram=%d dump_after_ms=45000",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,45000000));
    esp_timer_handle_t sample_timer;
    const esp_timer_create_args_t sample_args={.callback=request_snapshot,.arg=nullptr,.dispatch_method=ESP_TIMER_TASK,.name="tcpdiag",.skip_unhandled_events=true};
    ESP_ERROR_CHECK(esp_timer_create(&sample_args,&sample_timer));
    ESP_ERROR_CHECK(esp_timer_start_periodic(sample_timer,100000));
}
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;
}

extern "C" int __real_lwip_getaddrinfo(const char*,const char*,const struct addrinfo*,struct addrinfo**);
extern "C" int __wrap_lwip_getaddrinfo(const char *name,const char *service,const struct addrinfo *hints,struct addrinfo **result) {
    const bool original=name && strcmp(name,"gw.pangtiao.men")==0;
    const char *choice=original && DIAG_FIXED_IP[0] ? DIAG_FIXED_IP : name;
    const int ret=__real_lwip_getaddrinfo(choice,service,hints,result);
    if(original && ret==0) {
        unsigned n=0;
        for(auto *p=*result;p;p=p->ai_next) {
            if(p->ai_family==AF_INET) {
                const auto *address=(const struct sockaddr_in*)p->ai_addr;
                char ip[16];inet_ntop(AF_INET,&address->sin_addr,ip,sizeof(ip));
                ESP_LOGI("NETDIAG","[DNS_DIAG] index=%u ip=%s forced=%d",n++,ip,DIAG_FIXED_IP[0]!=0);
            }
        }
    }
    return ret;
}

extern "C" void __real_tcp_segs_free(struct tcp_seg *);
extern "C" void __wrap_tcp_segs_free(struct tcp_seg *segments) {
    unsigned n=0,bytes=0;
    for(const tcp_seg *seg=segments;seg && n<128;seg=seg->next){++n;bytes+=seg->len;}
    speech_diag_net_record(11,esp_timer_get_time(),0,(int)(uintptr_t)__builtin_return_address(0),(int)n,(int)bytes,0);
    __real_tcp_segs_free(segments);
}

extern "C" esp_err_t __real_esp_netif_receive(esp_netif_t*,void*,size_t,void*);
extern "C" esp_err_t __wrap_esp_netif_receive(esp_netif_t *netif,void *buffer,size_t len,void *eb) {
    if(eb){const int live=++rx_live;++rx_total;int max=rx_max.load();while(live>max && !rx_max.compare_exchange_weak(max,live)){} }
    return __real_esp_netif_receive(netif,buffer,len,eb);
}
extern "C" void __real_esp_netif_free_rx_buffer(void*,void*);
extern "C" void __wrap_esp_netif_free_rx_buffer(void *netif,void *buffer) {
    if(buffer){--rx_live;++rx_freed;}
    __real_esp_netif_free_rx_buffer(netif,buffer);
}