package db

import (
	"context"
	"fmt"
	"slices"
	"sync"
	"testing"
	"time"

	"tangled.org/core/spindle/quota"
)

const nixCache = "repo-cache"

// far in the future, so every use the tests record is older than it
var anyUse = time.Unix(1<<40, 0)

func nixHash(n int) string { return fmt.Sprintf("%032d", n) }

func nixReq(id quota.Identity, n int, size int64, refs ...string) quota.ReserveRequest {
	hash := nixHash(n)
	return quota.ReserveRequest{
		Kind:      quota.KindNixCache,
		Key:       quota.NixCacheKey(nixCache, hash),
		Identity:  id,
		Resources: quota.Resources{quota.ResourceCacheStorageBytes: size},
		NixObject: &quota.NixObject{Cache: nixCache, StoreHash: hash, StorePath: "/nix/store/" + hash + "-pkg", NarSize: size, References: refs},
	}
}

func nixRef(n int) string { return nixHash(n) + "-dep" }

func putNix(t *testing.T, qs *QuotaStore, id quota.Identity, n int, size int64, refs ...string) {
	t.Helper()
	res, err := qs.Reserve(context.Background(), nixReq(id, n, size, refs...))
	if err != nil {
		t.Fatal(err)
	}
	if !res.Allowed || res.ID == "" {
		t.Fatalf("reservation = %+v, want allowed with an id", res)
	}
	commitReservation(t, qs, res)
}

func count(t *testing.T, database *DB, query string, args ...any) int {
	t.Helper()
	var n int
	if err := database.QueryRow(query, args...).Scan(&n); err != nil {
		t.Fatal(err)
	}
	return n
}

func nixRow(t *testing.T, qs *QuotaStore, n int) NixObjectRow {
	t.Helper()
	var r NixObjectRow
	var lastUsed int64
	if err := qs.db.QueryRow(`SELECT `+nixObjectColumns+` FROM nix_objects WHERE cache = ? AND store_hash = ?`, nixCache, nixHash(n)).
		Scan(&r.Cache, &r.StoreHash, &r.RepoDID, &r.OwnerDID, &r.NarSize, &lastUsed, &r.Gen); err != nil {
		t.Fatal(err)
	}
	r.LastUsedAt = time.Unix(lastUsed, 0)
	return r
}

func nixHashes(rows []NixObjectRow) []string {
	var out []string
	for _, r := range rows {
		out = append(out, r.StoreHash)
	}
	return out
}

func TestNixCommitWritesRowEdgesAndAllocationTogether(t *testing.T) {
	database := newTestDB(t)
	qs := NewQuotaStore(database, nil)

	putNix(t, qs, alice1, 1, 500, nixRef(2), nixRef(3), nixRef(1))

	row := nixRow(t, qs, 1)
	if row.RepoDID != alice1.RepoDID || row.OwnerDID != alice1.OwnerDID || row.NarSize != 500 || row.Gen != 1 {
		t.Fatalf("row = %+v", row)
	}
	var state, storePath string
	if err := database.QueryRow(`SELECT state, store_path FROM nix_objects WHERE store_hash = ?`, nixHash(1)).Scan(&state, &storePath); err != nil {
		t.Fatal(err)
	}
	if state != "live" || storePath != "/nix/store/"+nixHash(1)+"-pkg" {
		t.Fatalf("state = %q, store path = %q", state, storePath)
	}
	if got := count(t, database, `SELECT count(*) FROM nix_refs WHERE cache = ? AND from_hash = ? AND to_hash IN (?, ?)`, nixCache, nixHash(1), nixHash(2), nixHash(3)); got != 2 {
		t.Fatalf("edges to the references = %d, want 2", got)
	}
	if got := count(t, database, `SELECT count(*) FROM nix_refs`); got != 2 {
		t.Fatalf("edges = %d, want the self reference dropped", got)
	}
	if got := count(t, database, `SELECT amount FROM quota_allocations WHERE repo_did = ? AND kind = ? AND key = ?`, alice1.RepoDID, quota.KindNixCache, quota.NixCacheKey(nixCache, nixHash(1))); got != 500 {
		t.Fatalf("allocation = %d, want 500", got)
	}
}

func TestNixFailingCommitWritesNothing(t *testing.T) {
	database := newTestDB(t)
	qs := NewQuotaStore(database, nil)
	res, err := qs.Reserve(context.Background(), nixReq(alice1, 1, 500, nixRef(2)))
	if err != nil {
		t.Fatal(err)
	}
	if _, err := database.Exec(`UPDATE quota_reservations SET nix_object = 'not json' WHERE id = ?`, res.ID); err != nil {
		t.Fatal(err)
	}

	if err := qs.Commit(context.Background(), res.ID); err == nil {
		t.Fatal("commit of a corrupt object succeeded")
	}
	for _, table := range []string{"nix_objects", "nix_refs", "quota_allocations"} {
		if got := count(t, database, `SELECT count(*) FROM `+table); got != 0 {
			t.Errorf("%s has %d rows after a failed commit", table, got)
		}
	}
	if got := count(t, database, `SELECT count(*) FROM quota_reservations WHERE id = ?`, res.ID); got != 1 {
		t.Errorf("reservation rows = %d, want it kept for a retry", got)
	}
}

func TestNixReuploadUpsertsWithoutDoubleChargeAndBumpsGen(t *testing.T) {
	database := newTestDB(t)
	qs := NewQuotaStore(database, nil)
	now := time.Unix(1_000, 0)
	qs.now = func() time.Time { return now }

	putNix(t, qs, alice1, 1, 500, nixRef(2))
	now = now.Add(time.Hour)
	putNix(t, qs, alice1, 1, 500, nixRef(3))

	row := nixRow(t, qs, 1)
	if row.Gen != 2 || !row.LastUsedAt.Equal(now) {
		t.Fatalf("row = %+v, want gen 2 used at %v", row, now)
	}
	byRepo, _, total, err := qs.NixCacheUsage(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if byRepo[alice1.RepoDID] != 500 || total != 500 {
		t.Fatalf("usage = %v total %d, want 500", byRepo, total)
	}
	if got := count(t, database, `SELECT count(*) FROM quota_allocations`); got != 1 {
		t.Fatalf("allocations = %d, want 1", got)
	}
	if got := count(t, database, `SELECT count(*) FROM nix_refs WHERE to_hash = ?`, nixHash(2)); got != 0 {
		t.Fatal("the old edge survived the re-upload")
	}
	if got := count(t, database, `SELECT count(*) FROM nix_refs WHERE to_hash = ?`, nixHash(3)); got != 1 {
		t.Fatal("the new edge is missing")
	}
}

func TestNixEvictableSkipsReferencedRecentAndEvicting(t *testing.T) {
	qs := NewQuotaStore(newTestDB(t), nil)
	ctx := context.Background()
	now := time.Unix(10_000, 0)
	qs.now = func() time.Time { return now }

	putNix(t, qs, alice1, 1, 10)
	putNix(t, qs, alice1, 2, 10, nixRef(1))
	now = now.Add(time.Minute)
	putNix(t, qs, alice1, 3, 10)
	now = now.Add(time.Minute)
	putNix(t, qs, alice2, 4, 10)
	now = now.Add(time.Hour)
	putNix(t, qs, alice1, 5, 10)
	if err := qs.RecordNixUse(ctx, nixCache, []string{nixHash(2)}); err != nil {
		t.Fatal(err)
	}

	cutoff := now.Add(-time.Minute)
	rows, err := qs.NixEvictable(ctx, NixEvictFilter{}, cutoff, 10)
	if err != nil {
		t.Fatal(err)
	}
	if got, want := nixHashes(rows), []string{nixHash(3), nixHash(4)}; !slices.Equal(got, want) {
		t.Fatalf("evictable = %v, want %v (oldest first)", got, want)
	}

	if rows, _ := qs.NixEvictable(ctx, NixEvictFilter{RepoDID: alice2.RepoDID}, cutoff, 10); !slices.Equal(nixHashes(rows), []string{nixHash(4)}) {
		t.Fatalf("repo filter = %v", nixHashes(rows))
	}
	if rows, _ := qs.NixEvictable(ctx, NixEvictFilter{OwnerDID: alice1.OwnerDID}, cutoff, 1); len(rows) != 1 {
		t.Fatalf("limit ignored: %v", nixHashes(rows))
	}

	marked, err := qs.MarkNixEvicting(ctx, rows[:1], cutoff)
	if err != nil || len(marked) != 1 {
		t.Fatalf("mark = %v, %v", marked, err)
	}
	rows, _ = qs.NixEvictable(ctx, NixEvictFilter{}, cutoff, 10)
	if len(rows) != 1 || rows[0].StoreHash == marked[0].StoreHash {
		t.Fatalf("evicting row offered again: %v", nixHashes(rows))
	}
}

func TestNixMarkEvictingFlipsARowOnce(t *testing.T) {
	qs := NewQuotaStore(newTestDB(t), nil)
	putNix(t, qs, alice1, 1, 10)
	rows := []NixObjectRow{nixRow(t, qs, 1)}

	var wg sync.WaitGroup
	results := make([][]NixObjectRow, 8)
	for i := range results {
		wg.Add(1)
		go func() {
			defer wg.Done()
			marked, err := qs.MarkNixEvicting(context.Background(), rows, anyUse)
			if err != nil {
				t.Error(err)
			}
			results[i] = marked
		}()
	}
	wg.Wait()

	flipped := 0
	for _, marked := range results {
		flipped += len(marked)
	}
	if flipped != 1 {
		t.Fatalf("%d callers flipped the row, want 1", flipped)
	}
	evicting, err := qs.NixEvicting(context.Background())
	if err != nil || len(evicting) != 1 || evicting[0].Gen != 1 {
		t.Fatalf("evicting = %+v, %v", evicting, err)
	}
}

func TestNixMarkEvictingSkipsAReuploadedRow(t *testing.T) {
	qs := NewQuotaStore(newTestDB(t), nil)
	ctx := context.Background()
	now := time.Unix(10_000, 0)
	qs.now = func() time.Time { return now }
	putNix(t, qs, alice1, 1, 10)
	putNix(t, qs, alice1, 2, 10)
	stale := nixRow(t, qs, 1)
	putNix(t, qs, alice1, 1, 10)
	now = now.Add(time.Hour)
	cutoff := now.Add(-time.Minute)
	used := nixRow(t, qs, 2)
	if err := qs.RecordNixUse(ctx, nixCache, []string{nixHash(2)}); err != nil {
		t.Fatal(err)
	}

	if marked, err := qs.MarkNixEvicting(ctx, []NixObjectRow{stale, used}, cutoff); err != nil || len(marked) != 0 {
		t.Fatalf("mark = %v, %v; want re-uploaded and recently used rows left live", marked, err)
	}
}

func TestNixFinishEvictionDropsRowEdgesAndAllocation(t *testing.T) {
	database := newTestDB(t)
	qs := NewQuotaStore(database, nil)
	ctx := context.Background()
	putNix(t, qs, alice1, 1, 10, nixRef(2))
	putNix(t, qs, alice1, 3, 20)
	marked, err := qs.MarkNixEvicting(ctx, []NixObjectRow{nixRow(t, qs, 1), nixRow(t, qs, 3)}, anyUse)
	if err != nil || len(marked) != 2 {
		t.Fatal(marked, err)
	}
	putNix(t, qs, alice1, 3, 20)

	if err := qs.FinishNixEviction(ctx, nixCache, marked); err != nil {
		t.Fatal(err)
	}
	if got := count(t, database, `SELECT count(*) FROM nix_objects WHERE store_hash = ?`, nixHash(1)); got != 0 {
		t.Fatal("row survived")
	}
	if got := count(t, database, `SELECT count(*) FROM nix_refs`); got != 0 {
		t.Fatal("edges survived")
	}
	if row := nixRow(t, qs, 3); row.Gen != 2 || count(t, database, `SELECT count(*) FROM quota_allocations`) != 1 {
		t.Fatalf("re-uploaded row = %+v", row)
	}
	if _, _, total, _ := qs.NixCacheUsage(ctx); total != 20 {
		t.Fatalf("usage = %d, want 20", total)
	}
}

func TestNixRecordUseOnlyMovesForwardAndIgnoresUnknown(t *testing.T) {
	qs := NewQuotaStore(newTestDB(t), nil)
	ctx := context.Background()
	now := time.Unix(5_000, 0)
	qs.now = func() time.Time { return now }
	putNix(t, qs, alice1, 1, 1)

	now = now.Add(time.Hour)
	if err := qs.RecordNixUse(ctx, nixCache, []string{nixHash(1), nixHash(99)}); err != nil {
		t.Fatal(err)
	}
	if got := nixRow(t, qs, 1).LastUsedAt; !got.Equal(now) {
		t.Fatalf("last used = %v, want %v", got, now)
	}
	now = now.Add(-30 * time.Minute)
	if err := qs.RecordNixUse(ctx, nixCache, []string{nixHash(1)}); err != nil {
		t.Fatal(err)
	}
	if got := nixRow(t, qs, 1).LastUsedAt; got.Equal(now) {
		t.Fatal("an older report moved last_used_at back")
	}
	if err := qs.RecordNixUse(ctx, "other-cache", []string{nixHash(1)}); err != nil {
		t.Fatal(err)
	}
	if count(t, qs.db, `SELECT count(*) FROM nix_objects`) != 1 {
		t.Fatal("unknown hash created a row")
	}
}

func TestDeleteQuotaStateForRepoDropsItsNixIndex(t *testing.T) {
	database := newTestDB(t)
	qs := NewQuotaStore(database, cacheDefaults(0, 100))
	ctx := context.Background()
	if got, err := qs.LimitFor(ctx, quota.ScopeRepo, alice1.RepoDID, quota.ResourceCacheStorageBytes); err != nil || got == nil || *got != 100 {
		t.Fatalf("repo default = %v, %v; want 100", got, err)
	}
	putNix(t, qs, alice1, 1, 10, nixRef(2))
	putNix(t, qs, alice2, 3, 10, nixRef(4))

	if err := database.DeleteQuotaStateForRepo(ctx, alice1.RepoDID); err != nil {
		t.Fatal(err)
	}
	if got := count(t, database, `SELECT count(*) FROM nix_objects`); got != 1 {
		t.Fatalf("objects = %d, want only the other repo's", got)
	}
	if got := count(t, database, `SELECT count(*) FROM nix_refs`); got != 1 {
		t.Fatalf("edges = %d, want only the other repo's", got)
	}
}

func TestNixEvictingReferrerKeepsItsDependency(t *testing.T) {
	qs := NewQuotaStore(newTestDB(t), nil)
	ctx := context.Background()
	now := time.Unix(10_000, 0)
	qs.now = func() time.Time { return now }
	putNix(t, qs, alice1, 1, 10)
	putNix(t, qs, alice1, 2, 10, nixRef(1))
	now = now.Add(time.Hour)
	cutoff := now.Add(-time.Minute)
	dep := nixRow(t, qs, 1)

	// the referrer's delete failed, so it is still in attic
	marked, err := qs.MarkNixEvicting(ctx, []NixObjectRow{nixRow(t, qs, 2)}, anyUse)
	if err != nil || len(marked) != 1 {
		t.Fatalf("mark = %v, %v", marked, err)
	}

	rows, _ := qs.NixEvictable(ctx, NixEvictFilter{}, cutoff, 10)
	if len(rows) != 0 {
		t.Fatalf("evictable = %v, want the dependency held by the evicting referrer", nixHashes(rows))
	}
	if flipped, err := qs.MarkNixEvicting(ctx, []NixObjectRow{dep}, anyUse); err != nil || len(flipped) != 0 {
		t.Fatalf("mark of the held dependency = %v, %v", flipped, err)
	}

	if err := qs.FinishNixEviction(ctx, nixCache, marked); err != nil {
		t.Fatal(err)
	}
	if rows, _ := qs.NixEvictable(ctx, NixEvictFilter{}, cutoff, 10); !slices.Equal(nixHashes(rows), []string{nixHash(1)}) {
		t.Fatalf("evictable after the referrer is gone = %v", nixHashes(rows))
	}
}

func TestNixUploadRevivesADependencyMidEviction(t *testing.T) {
	database := newTestDB(t)
	qs := NewQuotaStore(database, nil)
	ctx := context.Background()
	putNix(t, qs, alice1, 1, 10)
	marked, err := qs.MarkNixEvicting(ctx, []NixObjectRow{nixRow(t, qs, 1)}, anyUse)
	if err != nil || len(marked) != 1 {
		t.Fatal(marked, err)
	}

	putNix(t, qs, alice1, 2, 10, nixRef(1))
	if err := qs.FinishNixEviction(ctx, nixCache, marked); err != nil {
		t.Fatal(err)
	}

	row := nixRow(t, qs, 1)
	if row.Gen != 2 {
		t.Fatalf("dependency = %+v, want it revived", row)
	}
	var state string
	if err := database.QueryRow(`SELECT state FROM nix_objects WHERE store_hash = ?`, nixHash(1)).Scan(&state); err != nil || state != "live" {
		t.Fatalf("state = %q, %v", state, err)
	}
	if got := count(t, database, `SELECT count(*) FROM quota_allocations`); got != 2 {
		t.Fatalf("allocations = %d, want both", got)
	}
}

func TestNixSoftLimitHoldsOffAtTwiceTheLimit(t *testing.T) {
	for _, tc := range []struct {
		name     string
		defaults quota.Defaults
		scope    quota.Scope
		reason   string
	}{
		{"repo", cacheDefaults(0, 100), quota.ScopeRepo, quota.ReasonRepoLimit},
		{"owner", cacheDefaults(100, 0), quota.ScopeUser, quota.ReasonUserLimit},
	} {
		t.Run(tc.name, func(t *testing.T) {
			qs := NewQuotaStore(newTestDB(t), tc.defaults)
			ctx := context.Background()
			if res, err := qs.Reserve(ctx, nixReq(alice1, 9, 101)); err != nil || res.Allowed || res.Temporary || res.Scope != tc.scope || res.Limit != 100 {
				t.Fatalf("reservation = %+v, %v; want a hard %s refusal", res, err, tc.scope)
			}
			putNix(t, qs, alice1, 1, 90)
			putNix(t, qs, alice1, 2, 90)

			res, err := qs.Reserve(ctx, nixReq(alice1, 3, 21))
			if err != nil {
				t.Fatal(err)
			}
			if res.Allowed || !res.Temporary || res.Reason != tc.reason {
				t.Fatalf("reservation = %+v, want a temporary %s denial past 2x", res, tc.reason)
			}
			putNix(t, qs, alice1, 3, 20)
			// a repeat of an object adds nothing, so it still gets in past the ceiling
			putNix(t, qs, alice1, 3, 20)
		})
	}
}

func TestNixDropsRowsInChunksThatFailAlone(t *testing.T) {
	database := newTestDB(t)
	qs := NewQuotaStore(database, nil)
	ctx := context.Background()
	total := nixDropBatch + 200
	var rows []NixObjectRow
	for n := 1; n <= total; n++ {
		putNix(t, qs, alice1, n, 1)
		rows = append(rows, nixRow(t, qs, n))
	}
	marked, err := qs.MarkNixEvicting(ctx, rows, anyUse)
	if err != nil || len(marked) != total {
		t.Fatal(len(marked), err)
	}
	if _, err := database.Exec(`
		CREATE TRIGGER fail_one BEFORE DELETE ON nix_objects WHEN OLD.store_hash = '` + marked[nixDropBatch+10].StoreHash + `'
		BEGIN SELECT RAISE(ABORT, 'boom'); END
	`); err != nil {
		t.Fatal(err)
	}

	if err := qs.FinishNixEviction(ctx, nixCache, marked); err == nil {
		t.Fatal("finish survived the failing row")
	}
	if left := count(t, database, `SELECT count(*) FROM nix_objects`); left != total-nixDropBatch {
		t.Fatalf("%d rows left, want the first chunk gone and the failing one rolled back (%d)", left, total-nixDropBatch)
	}
	if got := count(t, database, `SELECT count(*) FROM quota_allocations`); got != total-nixDropBatch {
		t.Fatalf("allocations = %d, want %d", got, total-nixDropBatch)
	}
}

func TestNixReferenceCyclesPeel(t *testing.T) {
	qs := NewQuotaStore(newTestDB(t), nil)
	ctx := context.Background()
	const n = 3
	for i := 1; i <= n; i++ {
		putNix(t, qs, alice1, i, 10, nixRef(i%n+1))
	}
	for range n {
		rows, err := qs.NixEvictable(ctx, NixEvictFilter{}, anyUse, 10)
		if err != nil || len(rows) == 0 {
			t.Fatalf("evictable = %v, %v; want the cycle to peel", nixHashes(rows), err)
		}
		marked, err := qs.MarkNixEvicting(ctx, rows[:1], anyUse)
		if err != nil || len(marked) != 1 {
			t.Fatalf("mark = %v, %v", marked, err)
		}
		if err := qs.FinishNixEviction(ctx, nixCache, marked); err != nil {
			t.Fatal(err)
		}
	}
	if got := count(t, qs.db, `SELECT count(*) FROM nix_objects`); got != 0 {
		t.Fatalf("%d objects left", got)
	}
}

func TestNixOnlyAnEarlierCommittedDependencyIsHeld(t *testing.T) {
	ctx := context.Background()
	for _, tc := range []struct {
		name  string
		first int
		want  []string
	}{
		{"dependency first", 2, []string{nixHash(1)}},
		{"referrer first", 1, []string{nixHash(1), nixHash(2)}},
	} {
		t.Run(tc.name, func(t *testing.T) {
			qs := NewQuotaStore(newTestDB(t), nil)
			if tc.first == 2 {
				putNix(t, qs, alice1, 2, 10)
			}
			putNix(t, qs, alice1, 1, 10, nixRef(2))
			if tc.first == 1 {
				putNix(t, qs, alice1, 2, 10)
			}
			rows, err := qs.NixEvictable(ctx, NixEvictFilter{}, anyUse, 10)
			if err != nil || !slices.Equal(nixHashes(rows), tc.want) {
				t.Fatalf("evictable = %v, %v; want %v", nixHashes(rows), err, tc.want)
			}
		})
	}
}

func TestRecoverReturnsPublishingAtticUploadsAndDropsTheRest(t *testing.T) {
	ctx := context.Background()
	database := newTestDB(t)
	qs := NewQuotaStore(database, nil)

	reserved, err := qs.Reserve(ctx, nixReq(alice1, 1, 10))
	if err != nil {
		t.Fatal(err)
	}
	publishing, err := qs.Reserve(ctx, nixReq(alice1, 2, 20))
	if err != nil {
		t.Fatal(err)
	}
	if err := qs.BeginCommit(ctx, publishing.ID); err != nil {
		t.Fatal(err)
	}
	plain := reserveBytes(t, qs, alice1, "plain.nar", 30)
	if err := qs.BeginCommit(ctx, plain.ID); err != nil {
		t.Fatal(err)
	}

	got, err := qs.Recover(ctx, nil)
	if err != nil {
		t.Fatal(err)
	}
	if len(got) != 1 || got[0].ReservationID != publishing.ID || got[0].Cache != nixCache || got[0].StoreHash != nixHash(2) || time.Since(got[0].CreatedAt) > time.Minute {
		t.Fatalf("recovered = %+v, want only the publishing attic upload", got)
	}
	if count(t, database, `SELECT count(*) FROM quota_reservations`) != 1 || count(t, database, `SELECT count(*) FROM quota_reservations WHERE id = ?`, publishing.ID) != 1 {
		t.Errorf("want only %s left (reserved %s and plain %s dropped)", publishing.ID, reserved.ID, plain.ID)
	}
	if n := count(t, database, `SELECT count(*) FROM quota_allocations`); n != 0 {
		t.Errorf("allocations = %d, want none until the upload is resolved", n)
	}
}