/* Real ETCP lifecycle, AES-CCM and common ingress; deterministic encrypted wire faults. */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "etcp.h"
#include "etcp_api.h"
#include "etcp_connections.h"
#include "etcp_session.h"
#include "config_parser.h"
#include "../lib/debug_config.h"
#include "../lib/mem.h"

#define CHECK(x) do { if (!(x)) { fprintf(stderr,"FAIL %s:%d: %s\n",__func__,__LINE__,#x); exit(1); } } while (0)
struct endpoint { struct UTUN_INSTANCE inst; struct utun_config cfg; struct ETCP_CONN* c; struct ETCP_SOCKET sock; struct ETCP_LINK links[2]; };
struct frame { int from; size_t len; uint8_t bytes[PACKET_DATA_SIZE]; };
static struct endpoint ep[2];
static struct UASYNC* ua;
static struct frame frames[4096], saved[4];
static unsigned head, tail, drops, reinit[2], down[2];
static int drop[2][3], blackhole, dead_primary;

static ssize_t capture(socket_t fd, const void* data, size_t len, const struct sockaddr* addr, socklen_t alen,
                       struct ETCP_LINK* link, void* arg) {
    (void)fd; (void)addr; (void)alen;
    int from = (int)(intptr_t)arg;
    CHECK(len <= sizeof(frames[0].bytes));
    if (blackhole || (dead_primary && link->local_link_id == 1)) { drops++; return len; }
    CHECK(tail-head < 4096);
    struct frame* f = &frames[tail++ % 4096]; f->from = from; f->len = len; memcpy(f->bytes,data,len);
    return len;
}
static void event(struct ETCP_CONN* c, int ev, void* arg) {
    (void)c; int i = (int)(intptr_t)arg;
    if (ev == ETCP_CBK_EVENT_REINIT) reinit[i]++;
    if (ev == ETCP_CBK_EVENT_DOWN) down[i]++;
}
static void deliver(const struct frame* f, int allow_drop) {
    int to = 1-f->from;
    struct ETCP_DGRAM* p = memory_pool_alloc(ep[to].inst.pkt_pool); CHECK(p);
    size_t len = 0;
    CHECK(sc_decrypt(&ep[to].c->crypto_ctx,f->bytes,f->len,(uint8_t*)&p->timestamp,&len) == SC_OK);
    CHECK(len >= 4);
    int kind = p->data[0] - ETCP_SESSION_HELLO;
    if (allow_drop && kind >= 0 && kind < 3 && drop[f->from][kind]) {
        drop[f->from][kind]--; drops++; memory_pool_free(ep[to].inst.pkt_pool,p); return;
    }
    if (kind == 0 && !saved[0].len) saved[0] = *f;
    if (kind == 2 && !saved[1].len) saved[1] = *f;
    CHECK(etcp_packet_decrypted(&ep[to].sock,p,&ep[to].links[0],len) == 0);
}
static void pump(void) {
    uasync_poll(ua,1);
    unsigned end = tail; // callbacks can append replies
    while (head < end) { struct frame f = frames[head++ % 4096]; deliver(&f,1); }
}
static void converge(void) {
    uint64_t deadline = get_time_tb()+40000;
    while (get_time_tb() < deadline) {
        pump();
        if (!ep[0].c->reinit_pending && !ep[1].c->reinit_pending && !ep[0].c->session_timer && !ep[1].c->session_timer && head == tail) break;
    }
    for (int i=0;i<2;i++) {
        CHECK(ep[i].c->initialized && !ep[i].c->session_required && !ep[i].c->reinit_pending);
        CHECK(!ep[i].c->session_timer && !ep[i].c->close_requested && !down[i]);
        CHECK(ep[i].c->peer_reset_id == ep[1-i].c->reset_id);
    }
}
static void setup(int multi) {
    memset(ep,0,sizeof(ep)); memset(drop,0,sizeof(drop)); memset(saved,0,sizeof(saved));
    memset(reinit,0,sizeof(reinit)); memset(down,0,sizeof(down)); head=tail=drops=0; blackhole=dead_primary=0;
    ua = uasync_create(); CHECK(ua);
    for (int i=0;i<2;i++) {
        ep[i].inst.ua=ua; ep[i].inst.node_id=i+1; ep[i].inst.config=&ep[i].cfg;
        ep[i].inst.connections=queue_new(ua,16,0,8,"session_test");
        ep[i].inst.pkt_pool=memory_pool_init(sizeof(struct ETCP_DGRAM)+PACKET_DATA_SIZE,"session_packets");
        CHECK(sc_generate_keypair(&ep[i].inst.my_keys) == SC_OK);
        ep[i].c=etcp_connection_create(&ep[i].inst,"session_test"); CHECK(ep[i].c);
        ep[i].c->peer_node_id=2-i; ep[i].c->links_up=1;
        etcp_conn_ready(ep[i].c);
        etcp_conn_add_cbk(ep[i].c,event,(void*)(intptr_t)i,ETCP_CBK_EVENT_REINIT|ETCP_CBK_EVENT_DOWN);
        CHECK(sc_init_ctx(&ep[i].c->crypto_ctx,&ep[i].inst.my_keys) == SC_OK);
        ep[i].sock.instance=&ep[i].inst;
        for (int j=0;j<=multi;j++) {
            struct ETCP_LINK* l=&ep[i].links[j];
            l->etcp=ep[i].c; l->conn=&ep[i].sock; l->local_link_id=j+1;
            l->initialized=l->link_status=l->recv_keepalive=l->remote_keepalive=1;
            l->link_state=LINK_STATE_CONNECTED; l->mtu=1280; l->remote_addr.ss_family=AF_INET;
            l->send_hook=capture; l->send_hook_ctx=(void*)(intptr_t)i;
        }
        if (multi) ep[i].links[0].next=&ep[i].links[1];
        ep[i].c->links=&ep[i].links[0];
    }
    for (int i=0;i<2;i++) {
        ep[i].c->peer_reset_id=ep[1-i].c->reset_id;
        CHECK(sc_set_peer_public_key(&ep[i].c->crypto_ctx,ep[1-i].inst.my_keys.public_key,0) == SC_OK);
    }
}
static void teardown(void) {
    for (int i=0;i<2;i++) {
        ep[i].c->links=NULL; // fixture owns synthetic authenticated links
        etcp_connection_close(ep[i].c);
    }
    uasync_poll(ua,0);
    for (int i=0;i<2;i++) { queue_free(ep[i].inst.connections); memory_pool_destroy(ep[i].inst.pkt_pool); }
    uasync_poll(ua,0); CHECK(ua->timer_alloc_count == ua->timer_free_count); uasync_destroy(ua,0);
}
static struct frame stream_packet(int data) {
    struct ETCP_DGRAM* p=memory_pool_alloc(ep[0].inst.pkt_pool); CHECK(p);
    p->link=&ep[0].links[0]; p->noencrypt_len=0; p->data_len=8; memset(p->data,0,8); p->data[0]=data ? 0 : 1; if (data) p->data[1]=99;
    CHECK(etcp_encrypt_send(p)>0); memory_pool_free(ep[0].inst.pkt_pool,p);
    CHECK(tail == head+1); struct frame f=frames[head++ % 4096]; return f;
}
int main(void) {
    debug_config_init(); debug_set_level(DEBUG_LEVEL_WARN);
    const char* log=getenv("SESSION_TEST_LOG");
    if (log) { debug_enable_file_output(log,1); debug_set_category_level(DEBUG_CATEGORY_ETCP,DEBUG_LEVEL_DEBUG); }
    size_t baseline=u_get_allocated_count();
    for (int scenario=0;scenario<9;scenario++) {
        setup(scenario == 5);
        saved[2]=stream_packet(0); saved[3]=stream_packet(1);
        if (scenario == 2 || scenario == 3) for (int i=0;i<2;i++) for (int j=0;j<3;j++) drop[i][j]=1;
        if (scenario == 5) dead_primary=1;
        etcp_conn_fatal_reinit(ep[scenario == 1 ? 1 : 0].c,"fault test");
        if (scenario == 3) etcp_conn_fatal_reinit(ep[1].c,"simultaneous reset");
        if (scenario == 4) { pump(); etcp_conn_fatal_reinit(ep[0].c,"second reset before confirmation"); }
        if (scenario == 6) {
            blackhole=1; ep[0].c->session_started=get_time_tb()-100000;
            CHECK(etcp_conn_ref_take(ep[0].c)==0);
            ep[0].c->links=NULL; // timeout closes a real connection, links belong to fixture
            while (ep[0].c->state!=2) pump();
            CHECK(!ep[0].c->session_timer && ep[0].c->delete_complete);
            etcp_conn_ref_free(ep[0].c); ep[0].c=NULL;
            ep[1].c->links=NULL; etcp_connection_close(ep[1].c); uasync_poll(ua,0);
            for (int i=0;i<2;i++) { queue_free(ep[i].inst.connections); memory_pool_destroy(ep[i].inst.pkt_pool); }
            uasync_poll(ua,0); CHECK(ua->timer_alloc_count==ua->timer_free_count); uasync_destroy(ua,0);
        } else {
            converge();
            if (scenario==2 || scenario==3 || scenario==5) CHECK(drops>0);
            uint64_t peer=ep[1].c->peer_reset_id, local=ep[0].c->reset_id;
            unsigned before0=reinit[0], before1=reinit[1];
            ep[1].links[0].last_recv_local_time=123;
            deliver(&saved[2],0); deliver(&saved[3],0); CHECK(ep[1].links[0].last_recv_local_time==123); // stale ACK never reaches stream/liveness
            if (scenario==7) {
                if (saved[0].len) deliver(&saved[0],0);
                if (saved[1].len) { deliver(&saved[1],0); deliver(&saved[1],0); }
                converge();
            }
            if (scenario==8) {
                uint8_t short_frame[25]={ETCP_SESSION_CONFIRM};
                CHECK(etcp_session_receive(ep[0].c,short_frame,1));
                CHECK(!etcp_session_accept_data(ep[0].c,short_frame,0));
                etcp_conn_apply_peer_reset_id(ep[0].c,local); // stale/unproven physical INIT cannot switch peer
                CHECK(ep[0].c->peer_reset_id==ep[1].c->reset_id);
            }
            CHECK(ep[1].c->peer_reset_id==peer && ep[0].c->reset_id==local);
            CHECK(reinit[0]==before0 && reinit[1]==before1);
            teardown();
        }
        CHECK(u_get_allocated_count()==baseline); printf("PASS session scenario %d\n",scenario+1);
    }
    puts("All ETCP session scenarios passed"); return 0;
}