package probe

import (
	"path/filepath"
	"testing"
	"time"
)

func newTestStore(t *testing.T) *Store {
	t.Helper()
	st, err := Open(filepath.Join(t.TempDir(), "probe.db"))
	if err != nil {
		t.Fatalf("open: %v", err)
	}
	t.Cleanup(func() { st.Close() })
	return st
}

func book(t *testing.T, st *Store, agent int64, boot string, rx, tx int64, at time.Time, booted *int64, lastSeen int64) Traffic {
	t.Helper()
	tr, err := st.accumulate(agent, boot, rx, tx, at, booted, lastSeen)
	if err != nil {
		t.Fatalf("accumulate: %v", err)
	}
	return tr
}

func i64p(v int64) *int64 { return &v }

func TestPeriodStart(t *testing.T) {
	d := func(y int, m time.Month, day int) time.Time {
		return time.Date(y, m, day, 0, 0, 0, 0, time.Local)
	}
	cases := []struct {
		today time.Time
		reset int64
		want  time.Time
	}{
		{d(2026, 10, 23), 1, d(2026, 10, 1)},
		{d(2026, 10, 23), 23, d(2026, 10, 23)},
		{d(2026, 10, 23), 24, d(2026, 9, 24)},
		{d(2026, 2, 15), 31, d(2026, 1, 31)}, // Jan has 31
		{d(2026, 2, 28), 31, d(2026, 2, 28)}, // short month: reset lands on last day
		{d(2026, 3, 1), 31, d(2026, 2, 28)},  // Mar 1 with reset 31 is still in Feb's period
		{d(2026, 3, 31), 31, d(2026, 3, 31)},
		{d(2026, 1, 1), 15, d(2025, 12, 15)},  // wraparound into previous year
		{d(2026, 10, 23), 0, d(2026, 10, 1)},  // clamp low
		{d(2026, 10, 23), 99, d(2026, 9, 30)}, // clamp high: Sep has 30
	}
	for _, c := range cases {
		if got := PeriodStart(c.today, c.reset); !got.Equal(c.want) {
			t.Errorf("PeriodStart(%v, %d) = %v, want %v", c.today, c.reset, got, c.want)
		}
	}
}

func TestAccumulateFirstReadingAlignsBaseline(t *testing.T) {
	st := newTestStore(t)
	now := time.Now()
	tr := book(t, st, 1, "k1/fp", 1_000_000, 2_000_000, now, nil, 0)
	if tr.TotalRx != 0 || tr.TotalTx != 0 || tr.MonthRx != 0 || tr.DayRx != 0 {
		t.Fatalf("first reading must book 0, got %+v", tr)
	}
	// Same boot, counters climb: the delta books.
	tr = book(t, st, 1, "k1/fp", 1_000_600, 2_000_900, now.Add(time.Minute), nil, now.Unix())
	if tr.TotalRx != 600 || tr.TotalTx != 900 || tr.MonthRx != 600 || tr.DayTx != 900 {
		t.Fatalf("delta not booked: %+v", tr)
	}
}

func TestAccumulateCounterShrinkSameBoot(t *testing.T) {
	st := newTestStore(t)
	now := time.Now()
	book(t, st, 1, "k1/fp", 5000, 5000, now, nil, 0)
	book(t, st, 1, "k1/fp", 6000, 6000, now.Add(time.Minute), nil, now.Unix())
	// A counted interface vanishes: the sum shrinks. Booking it again
	// would double-count the remainder, so the delta is 0 and the
	// baseline re-aligns.
	tr := book(t, st, 1, "k1/fp", 2000, 2000, now.Add(2*time.Minute), nil, now.Unix())
	if tr.TotalRx != 1000 {
		t.Fatalf("shrink must book 0, total=%d", tr.TotalRx)
	}
	tr = book(t, st, 1, "k1/fp", 2500, 2600, now.Add(3*time.Minute), nil, now.Unix())
	if tr.TotalRx != 1500 || tr.TotalTx != 1600 {
		t.Fatalf("post-shrink delta wrong: %+v", tr)
	}
}

func TestAccumulateBootChangeZeroDelta(t *testing.T) {
	st := newTestStore(t)
	now := time.Now()
	book(t, st, 1, "k1/fpA", 1000, 1000, now, nil, 0)
	book(t, st, 1, "k1/fpA", 2000, 2000, now.Add(time.Minute), nil, now.Unix())
	// Same kernel, different interface digest (a counted iface joined):
	// epoch changed, delta must be 0 even though booted looks recent.
	booted := now.Add(30 * time.Second).Unix()
	tr := book(t, st, 1, "k1/fpB", 9000, 9000, now.Add(2*time.Minute), &booted, now.Add(time.Minute).Unix())
	if tr.TotalRx != 1000 {
		t.Fatalf("epoch change must book 0, total=%d", tr.TotalRx)
	}
}

func TestAccumulateUnseenBootBooksWholeReading(t *testing.T) {
	st := newTestStore(t)
	t0 := time.Now().Add(-2 * time.Hour)
	book(t, st, 1, "k1/fp", 500, 500, t0, nil, 0)
	tr := book(t, st, 1, "k1/fp", 800, 900, t0.Add(time.Minute), nil, t0.Unix())
	if tr.TotalRx != 300 {
		t.Fatalf("setup delta wrong: %+v", tr)
	}
	// The machine reboots while the master is not hearing from it:
	// new kernel boot, booted after the master's last-seen stamp, so
	// the whole counters-since-boot book.
	booted := t0.Add(30 * time.Minute).Unix()
	tr = book(t, st, 1, "k2/fp", 200, 100, t0.Add(31*time.Minute), &booted, t0.Add(time.Minute).Unix())
	if tr.TotalRx != 500 || tr.TotalTx != 500 {
		t.Fatalf("unseen boot must book whole reading: %+v", tr)
	}
	// And it continues from the new baseline.
	tr = book(t, st, 1, "k2/fp", 350, 100, t0.Add(32*time.Minute), &booted, t0.Add(31*time.Minute).Unix())
	if tr.TotalRx != 650 {
		t.Fatalf("post-reboot delta wrong: %+v", tr)
	}
}

func TestAccumulateBootChangeSeenRecentlyBooksZero(t *testing.T) {
	st := newTestStore(t)
	t0 := time.Now().Add(-time.Hour)
	book(t, st, 1, "k1/fp", 1000, 1000, t0, nil, 0)
	book(t, st, 1, "k1/fp", 2000, 2000, t0.Add(time.Minute), nil, t0.Unix())
	// A different kernel boot that started BEFORE the master last saw
	// the node (e.g. a second machine sharing the token): not the
	// unseen-boot exception; books 0.
	booted := t0.Add(30 * time.Second).Unix() // < lastSeen t0+60s
	tr := book(t, st, 1, "k2/fp", 5000, 5000, t0.Add(2*time.Minute), &booted, t0.Add(time.Minute).Unix())
	if tr.TotalRx != 1000 {
		t.Fatalf("seen boot change must book 0, total=%d", tr.TotalRx)
	}
}

func TestAccumulateCrossDay(t *testing.T) {
	st := newTestStore(t)
	day1 := dateOnly(time.Now()).Add(-48 * time.Hour) // midnight, two days ago
	r1 := day1.Add(23*time.Hour + 58*time.Minute)     // day D 23:58
	r2 := day1.Add(23*time.Hour + 59*time.Minute)     // day D 23:59
	r3 := day1.Add(24*time.Hour + time.Minute)        // day D+1 00:01
	book(t, st, 1, "k/fp", 1000, 1000, r1, nil, 0)
	tr := book(t, st, 1, "k/fp", 1100, 1100, r2, nil, r1.Unix())
	if tr.DayRx != 100 {
		t.Fatalf("day1 day_rx=%d, want 100", tr.DayRx)
	}
	tr = book(t, st, 1, "k/fp", 1150, 1150, r3, nil, r2.Unix())
	if tr.DayRx != 50 || tr.TotalRx != 150 || tr.MonthRx != 150 {
		t.Fatalf("cross-day split wrong: %+v", tr)
	}
	if tr.DayStart != dateOnly(r3).Format(dateLayout) {
		t.Fatalf("day stamp=%s", tr.DayStart)
	}
}

func TestAccumulateCrossPeriodWithResetDay(t *testing.T) {
	st := newTestStore(t)
	// reset day 15 on the meta row.
	if err := st.ensureRows(7); err != nil {
		t.Fatal(err)
	}
	if _, err := st.db.Exec(`UPDATE probe_meta SET traffic_reset_day = 15 WHERE agent_id = 7`); err != nil {
		t.Fatal(err)
	}
	now := time.Now()
	// Find a date in the current period and one in the next, both <= now
	// when possible: use period arithmetic from PeriodStart itself.
	periodStart := PeriodStart(dateOnly(now), 15)
	inPeriod := periodStart.Add(24 * time.Hour)
	if inPeriod.After(now) {
		inPeriod = periodStart
	}
	book(t, st, 7, "k/fp", 1000, 1000, inPeriod, nil, 0)
	tr := book(t, st, 7, "k/fp", 1300, 1300, inPeriod.Add(time.Minute), nil, inPeriod.Unix())
	if tr.MonthRx != 300 {
		t.Fatalf("month before boundary=%d, want 300", tr.MonthRx)
	}
	nextPeriod := PeriodStart(periodStart.AddDate(0, 2, 0), 15)
	at := nextPeriod.Add(time.Minute)
	tr = book(t, st, 7, "k/fp", 1500, 1500, at, nil, inPeriod.Unix())
	// Delta 200 lands in the new period; the month counter restarts at
	// the delta alone, the total keeps everything. (When the booking is
	// future-dated relative to a held-back stamp, accumulate's dated()
	// guard applies; both stamps here precede the booking date.)
	if tr.TotalRx != 500 {
		t.Fatalf("total=%d, want 500", tr.TotalRx)
	}
	if tr.MonthRx != 200 {
		t.Fatalf("month after boundary=%d, want 200", tr.MonthRx)
	}
	if tr.MonthStart != nextPeriod.Format(dateLayout) {
		t.Fatalf("month stamp=%s, want %s", tr.MonthStart, nextPeriod.Format(dateLayout))
	}
}

func TestMonthUsedModes(t *testing.T) {
	tr := Traffic{MonthRx: 100, MonthTx: 250}
	if got := tr.MonthUsed("sum"); got != 350 {
		t.Errorf("sum=%d", got)
	}
	if got := tr.MonthUsed("up"); got != 250 {
		t.Errorf("up=%d", got)
	}
	if got := tr.MonthUsed("down"); got != 100 {
		t.Errorf("down=%d", got)
	}
	if got := tr.MonthUsed("max"); got != 250 {
		t.Errorf("max=%d", got)
	}
	if got := tr.MonthUsed(""); got != 350 {
		t.Errorf("default=%d", got)
	}
}

func TestTrafficReadGating(t *testing.T) {
	st := newTestStore(t)
	// Book in an old period (two months ago): month figures belong to
	// that period and must read back as zero, totals intact.
	old := dateOnly(time.Now()).AddDate(0, -2, 0)
	book(t, st, 3, "k/fp", 1000, 1000, old, nil, 0)
	book(t, st, 3, "k/fp", 1400, 1500, old.Add(time.Minute), nil, old.Unix())
	got, ok, err := st.TrafficFor(3)
	if err != nil || !ok {
		t.Fatalf("TrafficFor: ok=%v err=%v", ok, err)
	}
	if got.TotalRx != 400 || got.TotalTx != 500 {
		t.Fatalf("totals must survive gating: %+v", got)
	}
	if got.MonthRx != 0 || got.MonthTx != 0 || got.DayRx != 0 {
		t.Fatalf("stale period must gate to zero: %+v", got)
	}
	all, err := st.AllTraffic()
	if err != nil {
		t.Fatal(err)
	}
	if all[3].TotalRx != 400 || all[3].MonthRx != 0 {
		t.Fatalf("AllTraffic gating wrong: %+v", all[3])
	}
	if _, ok, _ := st.TrafficFor(999); ok {
		t.Fatal("unknown agent must report not-found")
	}
}

func TestSetTrafficStampsCurrentPeriod(t *testing.T) {
	st := newTestStore(t)
	now := time.Now()
	book(t, st, 5, "k/fp", 1000, 1000, now, nil, 0)
	ok, err := st.SetTraffic(5, i64p(10_000), nil, i64p(700), nil)
	if err != nil || !ok {
		t.Fatalf("SetTraffic: ok=%v err=%v", ok, err)
	}
	got, _, _ := st.TrafficFor(5)
	if got.TotalRx != 10_000 || got.MonthRx != 700 {
		t.Fatalf("correction not visible: %+v", got)
	}
	// The next booking must not discard the correction: month is
	// stamped with the current period, so the delta adds onto it.
	tr := book(t, st, 5, "k/fp", 1100, 1000, now.Add(time.Minute), nil, now.Unix())
	if tr.MonthRx != 800 || tr.TotalRx != 10_100 {
		t.Fatalf("booking after correction wrong: %+v", tr)
	}
	if ok, _ := st.SetTraffic(4242, i64p(1), nil, nil, nil); ok {
		t.Fatal("SetTraffic on unknown agent must be false")
	}
}