#include "../src/etcp_bbr.h"
#include <string.h>
#include <stdio.h>
#include <stdlib.h>

#define MSEG(s) do { test_total++; printf("TEST %d: %-50s ", test_total, s); fflush(stdout); } while(0)
#define PASS()  do { puts("PASS"); test_passed++; } while(0)
#define FAIL(m) do { printf("FAIL: %s\n", m); test_failed++; return -1; } while(0)

static int test_total   = 0;
static int test_passed  = 0;
static int test_failed  = 0;

static uint32_t g_cwnd;
static uint32_t g_pacing;

static uint32_t run_ack_ex(struct bbr* bbr, uint32_t acked, uint32_t rtt_us,
                            int app, int lost, uint32_t inflight, uint32_t prior_del)
{
    bbr->now_tb += rtt_us / 100;
    struct bbr_rate_sample rs = {
        .delivered = acked, .interval_us = rtt_us, .rtt_us = rtt_us,
        .acked_sacked = acked, .prior_delivered = prior_del,
        .tx_in_flight = inflight > 0 ? inflight : acked,
        .lost = lost, .is_app_limited = app,
    };
    bbr_main(bbr, &rs, &g_cwnd, &g_pacing, 1400, inflight, inflight >= g_cwnd);
    return g_cwnd;
}

static uint32_t run_ack(struct bbr* bbr, uint32_t acked, uint32_t rtt_us,
                         int app, int lost, uint32_t inflight)
{ return run_ack_ex(bbr, acked, rtt_us, app, lost, inflight, bbr->delivered); }

/* --- tests --- */

static int test_init(void)
{
    MSEG("bbr_init — startup mode");
    struct bbr s; memset(&s, 0xFF, sizeof(s)); bbr_init(&s);
    if (s.mode != BBR_STARTUP) FAIL("mode");
    if (s.cycle_idx != 0) FAIL("cycle_idx");
    if (s.min_rtt_us != ~0U) FAIL("min_rtt_us");
    if (s.bw_lo != ~0U) FAIL("bw_lo");
    if (s.inflight_lo != ~0U) FAIL("inflight_lo");
    if (s.inflight_hi != ~0U) FAIL("inflight_hi");
    if (s.initialized != 1) FAIL("initialized");
    if (s.full_bw_reached) FAIL("full_bw_reached");
    PASS(); return 0;
}

static int test_startup_growth(void)
{
    MSEG("startup — cwnd grows");
    struct bbr s; memset(&s, 0, sizeof(s)); s.now_tb = 10000; bbr_init(&s);
    s.init_cwnd = 4; g_cwnd = 5600; g_pacing = 0;
    uint32_t prev = g_cwnd;
    for (int r = 0; r < 4; r++) {
        s.now_tb += 10000;
        for (int a = 0; a < 5; a++) run_ack(&s, 1400, 10000, 0, 0, prev);
        if (g_cwnd <= prev) FAIL("cwnd did not grow");
        prev = g_cwnd;
    }
    PASS(); return 0;
}

static int test_full_bw(void)
{
    MSEG("full_bw → DRAIN → PROBE_BW");
    struct bbr s; memset(&s, 0, sizeof(s)); s.now_tb = 10000; bbr_init(&s);
    s.init_cwnd = 4; g_cwnd = 5600; g_pacing = 0;
   // 4 rounds: first resets full_bw (sample_bw >= 0*1.25), next 3 count up
    for (int r = 0; r < 4; r++) {
        s.now_tb += 10000;
        uint32_t rd = s.delivered;
        for (int a = 0; a < 4; a++) run_ack_ex(&s, 1400, 10000, 0, 0, g_cwnd, rd);
    }
    if (!s.full_bw_reached) FAIL("full_bw_reached");
   // Mode transitions STARTUP→DRAIN→PROBE_BW might all happen in one ACK
    if (s.mode != BBR_DRAIN && s.mode != BBR_PROBE_BW) FAIL("mode not DRAIN/PROBE_BW after full_bw");
    PASS(); return 0;
    PASS(); return 0;
}

static int test_loss_startup(void)
{
    MSEG("loss events in STARTUP → early full_bw");
    struct bbr s; memset(&s, 0, sizeof(s)); s.now_tb = 10000; bbr_init(&s);
    s.init_cwnd = 4; g_cwnd = 5600; g_pacing = 0; s.loss_round_start = 1;
    for (int i = 0; i < 6; i++) { s.now_tb += 1000; run_ack(&s, 1400, 5000, 0, 1, 5600); }
    if (!s.full_bw_reached) FAIL("full_bw_reached not set");
    PASS(); return 0;
}

static int test_cwnd_bounded(void)
{
    MSEG("cwnd capped by inflight_hi");
    struct bbr s; memset(&s, 0, sizeof(s)); s.now_tb = 10000; bbr_init(&s);
    s.mode = BBR_PROBE_BW; s.full_bw_reached = 1; s.cycle_idx = BBR_BW_PROBE_UP;
    s.inflight_hi = 4000; s.inflight_lo = 4000; g_cwnd = 5600; g_pacing = 0;
    s.now_tb += 5000; run_ack(&s, 1400, 5000, 0, 0, 5000);
    if (g_cwnd > 4000) FAIL("cwnd > inflight_hi");
    PASS(); return 0;
}

static int test_loss_cut(void)
{
    MSEG("loss cuts inflight_lo");
    struct bbr s; memset(&s, 0, sizeof(s)); s.now_tb = 10000; bbr_init(&s);
    s.mode = BBR_PROBE_BW; s.full_bw_reached = 1; s.cycle_idx = BBR_BW_PROBE_CRUISE;
    s.inflight_lo = 5000; s.bw_lo = 1000;
    s.bw_latest = 500; s.inflight_latest = 3000;
    s.delivered = 0; s.loss_round_delivered = 0;
    s.cycle_mstamp_tb = s.now_tb;                    // prevent phase advance
    s.probe_wait_us = 2000000;
    uint32_t prev = s.inflight_lo;
    s.now_tb += 5000; run_ack_ex(&s, 1400, 5000, 0, 1, 5000, 0);
    if (s.inflight_lo >= prev) FAIL("inflight_lo not reduced");
    PASS(); return 0;
}

static int test_fast_path(void)
{
    MSEG("app-limited ACK triggers fast_path");
    struct bbr s; memset(&s, 0, sizeof(s)); s.now_tb = 10000; bbr_init(&s);
    s.init_cwnd = 4; s.full_bw_reached = 1;
    s.mode = BBR_PROBE_BW; s.cycle_idx = BBR_BW_PROBE_CRUISE;
    s.try_fast_path = 1; s.bw_hi[0] = 10000; s.bw_lo = 5000;
    g_cwnd = 5600; g_pacing = 0;
    s.now_tb += 1000; run_ack(&s, 1400, 5000, 1, 0, 2000);
    PASS(); return 0;
}

static int test_probe_rtt(void)
{
    MSEG("PROBE_RTT entry after 5s");
    struct bbr s; memset(&s, 0, sizeof(s)); s.now_tb = 10000; bbr_init(&s);
    s.mode = BBR_PROBE_BW; s.full_bw_reached = 1; s.cycle_idx = BBR_BW_PROBE_CRUISE;
    s.probe_rtt_min_us = ~0U; s.probe_rtt_min_stamp_tb = 10000;
    s.min_rtt_us = ~0U; s.min_rtt_stamp_tb = 10000;
    g_cwnd = 5600; g_pacing = 0;
    s.now_tb = 10000 + 5000U * 10;  // +5 seconds in 0.1ms
    run_ack(&s, 1400, 1000, 0, 0, 5600);
    if (s.mode != BBR_PROBE_RTT) FAIL("mode not PROBE_RTT");
    if (g_cwnd > BBR_CWND_MIN_TARGET * 1400 + 1400) FAIL("cwnd not capped");
    PASS(); return 0;
}

int main(void)
{
    test_init();
    test_startup_growth();
    test_full_bw();
    test_loss_startup();
    test_cwnd_bounded();
    test_loss_cut();
    test_fast_path();
    test_probe_rtt();
    printf("\n=== Results: %d/%d passed, %d failed ===\n", test_passed, test_total, test_failed);
    return test_failed ? 1 : 0;
}