pangolin/server/internal/usage/usage_integration_test.go

//go:build integration

package usage_test

import (
	"context"
	"database/sql"
	"fmt"
	"sync"
	"testing"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/redis/go-redis/v9"
	"github.com/testcontainers/testcontainers-go"
	tcmysql "github.com/testcontainers/testcontainers-go/modules/mysql"
	tcredis "github.com/testcontainers/testcontainers-go/modules/redis"

	"github.com/wangjia/pangolin/server/internal/usage"
)

// --------------------------------------------------------------------------
// Container setup
// --------------------------------------------------------------------------

func setupMySQL(t *testing.T) *sql.DB {
	t.Helper()
	ctx := context.Background()

	ctr, err := tcmysql.Run(ctx, "mysql:8.0",
		tcmysql.WithDatabase("pangolin_test"),
		tcmysql.WithUsername("root"),
		tcmysql.WithPassword("test"),
	)
	testcontainers.CleanupContainer(t, ctr)
	if err != nil {
		t.Fatalf("mysql container: %v", err)
	}
	dsn, err := ctr.ConnectionString(ctx, "parseTime=true", "loc=UTC", "time_zone=%27%2B00%3A00%27")
	if err != nil {
		t.Fatalf("mysql dsn: %v", err)
	}
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		t.Fatalf("open mysql: %v", err)
	}
	t.Cleanup(func() { db.Close() })
	if err := applySchema(db); err != nil {
		t.Fatalf("schema: %v", err)
	}
	return db
}

func setupRedis(t *testing.T) *redis.Client {
	t.Helper()
	ctx := context.Background()
	ctr, err := tcredis.Run(ctx, "redis:7-alpine")
	testcontainers.CleanupContainer(t, ctr)
	if err != nil {
		t.Fatalf("redis container: %v", err)
	}
	addr, err := ctr.ConnectionString(ctx)
	if err != nil {
		t.Fatalf("redis addr: %v", err)
	}
	for _, p := range []string{"redis://", "rediss://"} {
		if len(addr) > len(p) && addr[:len(p)] == p {
			addr = addr[len(p):]
		}
	}
	rdb := redis.NewClient(&redis.Options{Addr: addr})
	t.Cleanup(func() { rdb.Close() })
	return rdb
}

func applySchema(db *sql.DB) error {
	stmts := []string{
		`CREATE TABLE IF NOT EXISTS users (
			id          BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY,
			uuid        CHAR(36) NOT NULL UNIQUE,
			email       VARCHAR(255) NOT NULL UNIQUE,
			pw_hash     VARCHAR(255) NOT NULL DEFAULT '',
			dp_uuid     CHAR(36) NOT NULL,
			status      ENUM('active','banned') NOT NULL DEFAULT 'active',
			created_at  DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
			INDEX idx_dp (dp_uuid)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,
		`CREATE TABLE IF NOT EXISTS plans (
			id            BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY,
			code          ENUM('free','pro','team') NOT NULL UNIQUE,
			max_devices   INT NOT NULL DEFAULT 1,
			daily_minutes INT NULL,
			ad_gate       BOOLEAN NOT NULL DEFAULT FALSE
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,
		`CREATE TABLE IF NOT EXISTS subscriptions (
			id         BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY,
			user_id    BIGINT UNSIGNED NOT NULL,
			plan_id    BIGINT UNSIGNED NOT NULL,
			expires_at DATETIME(6) NOT NULL,
			source     ENUM('trial','code') NOT NULL,
			created_at DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
			INDEX idx_user_exp (user_id, expires_at)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,
		`CREATE TABLE IF NOT EXISTS usage_daily (
			user_id        BIGINT UNSIGNED NOT NULL,
			date           DATE NOT NULL,
			bytes_up       BIGINT UNSIGNED NOT NULL DEFAULT 0,
			bytes_down     BIGINT UNSIGNED NOT NULL DEFAULT 0,
			minutes_used   INT NOT NULL DEFAULT 0,
			ad_unlocked_at DATETIME(6) NULL,
			PRIMARY KEY (user_id, date)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,
		`INSERT IGNORE INTO plans (code, max_devices, daily_minutes, ad_gate)
		 VALUES ('free',1,10,TRUE), ('pro',5,NULL,FALSE), ('team',10,NULL,FALSE)`,
	}
	for _, s := range stmts {
		if _, err := db.Exec(s); err != nil {
			return fmt.Errorf("exec: %w\nSQL: %s", err, s)
		}
	}
	return nil
}

// createUser inserts a user with the given dp_uuid and returns its id.
func createUser(t *testing.T, db *sql.DB, dpUUID string) int64 {
	t.Helper()
	uuid := fmt.Sprintf("u-%d", time.Now().UnixNano())
	res, err := db.Exec(
		`INSERT INTO users (uuid, email, pw_hash, dp_uuid) VALUES (?,?,?,?)`,
		uuid, uuid+"@example.com", "x", dpUUID)
	if err != nil {
		t.Fatalf("createUser: %v", err)
	}
	id, _ := res.LastInsertId()
	return id
}

// giveSubscription grants userID an active subscription on the given plan.
func giveSubscription(t *testing.T, db *sql.DB, userID int64, plan string) {
	t.Helper()
	var planID int64
	if err := db.QueryRow(`SELECT id FROM plans WHERE code=?`, plan).Scan(&planID); err != nil {
		t.Fatalf("plan lookup: %v", err)
	}
	if _, err := db.Exec(
		`INSERT INTO subscriptions (user_id, plan_id, expires_at, source)
		 VALUES (?,?,?,'code')`,
		userID, planID, time.Now().UTC().AddDate(0, 0, 30)); err != nil {
		t.Fatalf("giveSubscription: %v", err)
	}
}

// fakeVerifier is an AdVerifier test double.
type fakeVerifier struct{ ok bool }

func (f fakeVerifier) Provider() string { return "fake" }
func (f fakeVerifier) Verify(context.Context, usage.AdVerifyRequest) error {
	if f.ok {
		return nil
	}
	return fmt.Errorf("fake verifier: rejected")
}

// --------------------------------------------------------------------------
// Aggregation
// --------------------------------------------------------------------------

func TestIntAggregateConcurrentMultiNode(t *testing.T) {
	db := setupMySQL(t)
	rdb := setupRedis(t)
	store := usage.NewStore(db)
	agg := usage.NewAggregator(store, rdb, time.Minute, time.Minute)

	uid := createUser(t, db, "dp-agg-1")

	const goroutines = 25
	const each = 4
	var wg sync.WaitGroup
	wg.Add(goroutines)
	for g := 0; g < goroutines; g++ {
		go func(g int) {
			defer wg.Done()
			for i := 0; i < each; i++ {
				err := agg.ReportUsage(context.Background(), usage.Report{
					DPUUID:    "dp-agg-1",
					BytesUp:   100,
					BytesDown: 200,
					Minutes:   1,
					BatchID:   fmt.Sprintf("b-%d-%d", g, i), // unique → all counted
				})
				if err != nil {
					t.Errorf("report: %v", err)
				}
			}
		}(g)
	}
	wg.Wait()

	var up, down uint64
	var minutes int
	db.QueryRow(`SELECT bytes_up, bytes_down, minutes_used FROM usage_daily WHERE user_id=?`, uid).
		Scan(&up, &down, &minutes)

	wantUp := uint64(goroutines * each * 100)
	wantMin := goroutines * each
	if up != wantUp || down != wantUp*2 || minutes != wantMin {
		t.Errorf("got up=%d down=%d min=%d, want up=%d down=%d min=%d",
			up, down, minutes, wantUp, wantUp*2, wantMin)
	}
}

func TestIntAggregateReplayDedup(t *testing.T) {
	db := setupMySQL(t)
	rdb := setupRedis(t)
	store := usage.NewStore(db)
	agg := usage.NewAggregator(store, rdb, time.Minute, time.Minute)

	uid := createUser(t, db, "dp-replay")

	r := usage.Report{DPUUID: "dp-replay", BytesUp: 500, Minutes: 5, BatchID: "same-batch"}
	for i := 0; i < 4; i++ {
		if err := agg.ReportUsage(context.Background(), r); err != nil {
			t.Fatalf("report %d: %v", i, err)
		}
	}

	var up uint64
	var minutes int
	db.QueryRow(`SELECT bytes_up, minutes_used FROM usage_daily WHERE user_id=?`, uid).Scan(&up, &minutes)
	if up != 500 || minutes != 5 {
		t.Errorf("replay not deduped: up=%d minutes=%d, want 500/5", up, minutes)
	}
}

func TestIntAggregateCrossUTCDayBucketing(t *testing.T) {
	db := setupMySQL(t)
	store := usage.NewStore(db)
	agg := usage.NewAggregator(store, nil, time.Minute, time.Minute)

	uid := createUser(t, db, "dp-days")

	day1 := time.Date(2026, 6, 12, 23, 30, 0, 0, time.UTC)
	day2 := time.Date(2026, 6, 13, 0, 15, 0, 0, time.UTC)

	must := func(at time.Time, minutes int) {
		if err := agg.ReportUsage(context.Background(), usage.Report{
			DPUUID: "dp-days", BytesUp: 10, Minutes: minutes, At: at,
		}); err != nil {
			t.Fatalf("report: %v", err)
		}
	}
	must(day1, 3)
	must(day1.Add(10*time.Minute), 2) // still 06-12
	must(day2, 7)                     // 06-13

	check := func(date string, wantMin int) {
		var m int
		err := db.QueryRow(`SELECT minutes_used FROM usage_daily WHERE user_id=? AND date=?`, uid, date).Scan(&m)
		if err != nil {
			t.Fatalf("query %s: %v", date, err)
		}
		if m != wantMin {
			t.Errorf("date %s minutes=%d, want %d", date, m, wantMin)
		}
	}
	check("2026-06-12", 5)
	check("2026-06-13", 7)
}

func TestIntAggregateUnknownDPUUIDDropped(t *testing.T) {
	db := setupMySQL(t)
	store := usage.NewStore(db)
	agg := usage.NewAggregator(store, nil, time.Minute, time.Minute)

	err := agg.ReportUsage(context.Background(), usage.Report{
		DPUUID: "does-not-exist", BytesUp: 1, Minutes: 1,
	})
	if err != nil {
		t.Fatalf("unknown dp_uuid should be a silent no-op, got %v", err)
	}
	var n int
	db.QueryRow(`SELECT COUNT(*) FROM usage_daily`).Scan(&n)
	if n != 0 {
		t.Errorf("expected no rows for unknown dp_uuid, got %d", n)
	}
}

// --------------------------------------------------------------------------
// Quota (CheckFreeConnect)
// --------------------------------------------------------------------------

func TestIntQuotaFreeNotUnlocked(t *testing.T) {
	db := setupMySQL(t)
	svc := usage.NewService(usage.NewStore(db), nil, fakeVerifier{ok: true}, time.Hour)

	uid := createUser(t, db, "dp-q1") // no subscription → free plan

	_, apiErr := svc.CheckFreeConnect(context.Background(), uid)
	if apiErr == nil || apiErr.Code != "AD_NOT_UNLOCKED" {
		t.Fatalf("expected AD_NOT_UNLOCKED, got %v", apiErr)
	}
	if apiErr.MessageZH == "" || apiErr.MessageEn == "" {
		t.Error("expected bilingual error messages")
	}
}

func TestIntQuotaFreeUnlockedRemainingDecrements(t *testing.T) {
	db := setupMySQL(t)
	store := usage.NewStore(db)
	agg := usage.NewAggregator(store, nil, time.Minute, time.Minute)
	svc := usage.NewService(store, nil, fakeVerifier{ok: true}, time.Hour)

	uid := createUser(t, db, "dp-q2")

	if _, err := store.MarkAdUnlocked(context.Background(), uid, time.Now().UTC()); err != nil {
		t.Fatalf("unlock: %v", err)
	}

	rem, apiErr := svc.CheckFreeConnect(context.Background(), uid)
	if apiErr != nil {
		t.Fatalf("unexpected err: %v", apiErr)
	}
	if rem != 10 {
		t.Errorf("remaining=%d, want 10", rem)
	}

	// Consume 4 minutes.
	agg.ReportUsage(context.Background(), usage.Report{DPUUID: "dp-q2", Minutes: 4})
	rem, _ = svc.CheckFreeConnect(context.Background(), uid)
	if rem != 6 {
		t.Errorf("after 4 min: remaining=%d, want 6", rem)
	}
}

func TestIntQuotaFreeExhausted(t *testing.T) {
	db := setupMySQL(t)
	store := usage.NewStore(db)
	svc := usage.NewService(store, nil, fakeVerifier{ok: true}, time.Hour)

	uid := createUser(t, db, "dp-q3")
	store.MarkAdUnlocked(context.Background(), uid, time.Now().UTC())
	store.AggregateUsage(context.Background(), uid, time.Now().UTC(), 0, 0, 10)

	_, apiErr := svc.CheckFreeConnect(context.Background(), uid)
	if apiErr == nil || apiErr.Code != "QUOTA_EXHAUSTED" {
		t.Fatalf("expected QUOTA_EXHAUSTED, got %v", apiErr)
	}
}

func TestIntQuotaPaidUnlimited(t *testing.T) {
	db := setupMySQL(t)
	store := usage.NewStore(db)
	svc := usage.NewService(store, nil, fakeVerifier{ok: true}, time.Hour)

	for _, plan := range []string{"pro", "team"} {
		uid := createUser(t, db, "dp-"+plan)
		giveSubscription(t, db, uid, plan)
		// Even with lots of minutes used, paid plans are not gated.
		store.AggregateUsage(context.Background(), uid, time.Now().UTC(), 0, 0, 9999)

		rem, apiErr := svc.CheckFreeConnect(context.Background(), uid)
		if apiErr != nil {
			t.Fatalf("%s: unexpected err %v", plan, apiErr)
		}
		if rem != usage.Unlimited {
			t.Errorf("%s: remaining=%d, want Unlimited(%d)", plan, rem, usage.Unlimited)
		}
	}
}

// --------------------------------------------------------------------------
// Ads unlock
// --------------------------------------------------------------------------

func TestIntAdsUnlockForgedRejected(t *testing.T) {
	db := setupMySQL(t)
	rdb := setupRedis(t)
	svc := usage.NewService(usage.NewStore(db), rdb, fakeVerifier{ok: false}, time.Hour)

	uid := createUser(t, db, "dp-ad1")
	_, apiErr := svc.UnlockAd(context.Background(), uid, "device-1", "forged-token")
	if apiErr == nil || apiErr.Code != "AD_VERIFY_FAILED" {
		t.Fatalf("expected AD_VERIFY_FAILED, got %v", apiErr)
	}
}

func TestIntAdsUnlockReplayRejected(t *testing.T) {
	db := setupMySQL(t)
	rdb := setupRedis(t)
	svc := usage.NewService(usage.NewStore(db), rdb, fakeVerifier{ok: true}, time.Hour)

	uid := createUser(t, db, "dp-ad2")
	already, apiErr := svc.UnlockAd(context.Background(), uid, "device-1", "token-xyz")
	if apiErr != nil || already {
		t.Fatalf("first unlock: already=%v err=%v", already, apiErr)
	}
	// Same token again → replay.
	_, apiErr = svc.UnlockAd(context.Background(), uid, "device-1", "token-xyz")
	if apiErr == nil || apiErr.Code != "AD_TOKEN_REPLAY" {
		t.Fatalf("expected AD_TOKEN_REPLAY, got %v", apiErr)
	}
}

func TestIntAdsUnlockSecondTimeIdempotent(t *testing.T) {
	db := setupMySQL(t)
	rdb := setupRedis(t)
	svc := usage.NewService(usage.NewStore(db), rdb, fakeVerifier{ok: true}, time.Hour)

	uid := createUser(t, db, "dp-ad3")
	if _, apiErr := svc.UnlockAd(context.Background(), uid, "d", "tok-1"); apiErr != nil {
		t.Fatalf("first unlock: %v", apiErr)
	}
	// Different (valid) token, same day → idempotent success.
	already, apiErr := svc.UnlockAd(context.Background(), uid, "d", "tok-2")
	if apiErr != nil {
		t.Fatalf("second unlock err: %v", apiErr)
	}
	if !already {
		t.Error("expected already-unlocked idempotent success")
	}

	// Exactly one unlock timestamp.
	var n int
	db.QueryRow(`SELECT COUNT(*) FROM usage_daily WHERE user_id=? AND ad_unlocked_at IS NOT NULL`, uid).Scan(&n)
	if n != 1 {
		t.Errorf("expected 1 unlocked day, got %d", n)
	}
}

// --------------------------------------------------------------------------
// Usage curve & today summary
// --------------------------------------------------------------------------

func TestIntUsageCurveZeroFill(t *testing.T) {
	db := setupMySQL(t)
	store := usage.NewStore(db)
	svc := usage.NewService(store, nil, nil, time.Hour)

	uid := createUser(t, db, "dp-curve")
	today := time.Now().UTC()
	epochHour := func(tm time.Time) int64 { return tm.UTC().Unix() / 3600 }
	// Only two of the last 7 days have data (curve reads hourly buckets now).
	db.Exec(`INSERT INTO usage_hourly (user_id, hour, bytes_up, bytes_down, minutes_used) VALUES (?,?,?,?,?)`,
		uid, epochHour(today), 1000, 2000, 5)
	db.Exec(`INSERT INTO usage_hourly (user_id, hour, bytes_up, bytes_down, minutes_used) VALUES (?,?,?,?,?)`,
		uid, epochHour(today.AddDate(0, 0, -3)), 50, 60, 2)

	points, apiErr := svc.UsageCurve(context.Background(), uid, 7, 0)
	if apiErr != nil {
		t.Fatalf("curve: %v", apiErr)
	}
	if len(points) != 7 {
		t.Fatalf("expected 7 points, got %d", len(points))
	}
	// Oldest first, dates contiguous.
	for i := 1; i < len(points); i++ {
		if points[i].Date <= points[i-1].Date {
			t.Errorf("points not ascending: %s then %s", points[i-1].Date, points[i].Date)
		}
	}
	// Last point = today with the data.
	last := points[len(points)-1]
	if last.BytesUp != 1000 || last.BytesDown != 2000 || last.MinutesUsed != 5 {
		t.Errorf("today point wrong: %+v", last)
	}
	// Days with no data are zero-filled.
	zeros := 0
	for _, p := range points {
		if p.BytesUp == 0 && p.BytesDown == 0 && p.MinutesUsed == 0 {
			zeros++
		}
	}
	if zeros != 5 {
		t.Errorf("expected 5 zero-filled days, got %d", zeros)
	}
}

func TestIntUsageCurveOnlyCurrentUser(t *testing.T) {
	db := setupMySQL(t)
	store := usage.NewStore(db)
	svc := usage.NewService(store, nil, nil, time.Hour)

	uA := createUser(t, db, "dp-A")
	uB := createUser(t, db, "dp-B")
	h := time.Now().UTC().Unix() / 3600
	db.Exec(`INSERT INTO usage_hourly (user_id, hour, bytes_up, bytes_down, minutes_used) VALUES (?,?,?,?,?)`, uA, h, 111, 0, 1)
	db.Exec(`INSERT INTO usage_hourly (user_id, hour, bytes_up, bytes_down, minutes_used) VALUES (?,?,?,?,?)`, uB, h, 999, 0, 9)

	points, _ := svc.UsageCurve(context.Background(), uA, 1, 0)
	if len(points) != 1 || points[0].BytesUp != 111 {
		t.Errorf("user A curve leaked other user data: %+v", points)
	}
}

func TestIntTodaySummary(t *testing.T) {
	db := setupMySQL(t)
	store := usage.NewStore(db)
	svc := usage.NewService(store, nil, nil, time.Hour)

	// Free user: limited + ad flag.
	uFree := createUser(t, db, "dp-sf")
	store.MarkAdUnlocked(context.Background(), uFree, time.Now().UTC())
	store.AggregateUsage(context.Background(), uFree, time.Now().UTC(), 0, 0, 3)

	sum, apiErr := svc.TodaySummary(context.Background(), uFree)
	if apiErr != nil {
		t.Fatalf("summary: %v", apiErr)
	}
	if sum.MinutesUsed != 3 || sum.MinutesRemaining == nil || *sum.MinutesRemaining != 7 || !sum.AdUnlocked {
		t.Errorf("free summary wrong: %+v (remaining=%v)", sum, sum.MinutesRemaining)
	}

	// Pro user: unlimited (nil remaining).
	uPro := createUser(t, db, "dp-sp")
	giveSubscription(t, db, uPro, "pro")
	sum2, _ := svc.TodaySummary(context.Background(), uPro)
	if sum2.MinutesRemaining != nil {
		t.Errorf("pro remaining should be nil(unlimited), got %v", *sum2.MinutesRemaining)
	}
}