package db

import (
	"context"
	"database/sql"
	"errors"
	"sync"
	"testing"

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

var (
	alice1 = quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"}
	alice2 = quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo2"}
)

func cacheDefaults(user, repo int64) quota.Defaults {
	defaults := quota.Defaults{}
	if user != 0 {
		defaults[quota.ScopeUser] = quota.Resources{quota.ResourceCacheStorageBytes: user}
	}
	if repo != 0 {
		defaults[quota.ScopeRepo] = quota.Resources{quota.ResourceCacheStorageBytes: repo}
	}
	return defaults
}

func reserveBytes(t *testing.T, qs *QuotaStore, id quota.Identity, key string, amount int64) quota.Reservation {
	t.Helper()
	res, err := qs.Reserve(context.Background(), quota.ReserveRequest{
		Kind:      quota.KindNixCache,
		Key:       key,
		Identity:  id,
		Resources: quota.Resources{quota.ResourceCacheStorageBytes: amount},
	})
	if err != nil {
		t.Fatal(err)
	}
	return res
}

func commitReservation(t *testing.T, qs *QuotaStore, res quota.Reservation) {
	t.Helper()
	if err := qs.BeginCommit(context.Background(), res.ID); err != nil {
		t.Fatal(err)
	}
	if err := qs.Commit(context.Background(), res.ID); err != nil {
		t.Fatal(err)
	}
}

func committedAmount(t *testing.T, database *DB, id quota.Identity) int64 {
	t.Helper()
	var amount int64
	if err := database.QueryRow(`
		SELECT amount FROM quota_allocations
		WHERE repo_did = ? AND resource = ? AND kind = ? AND key = ?
	`, id.RepoDID, quota.ResourceCacheStorageBytes, quota.KindNixCache, "nar/object.nar").Scan(&amount); err != nil {
		t.Fatal(err)
	}
	return amount
}

func TestRepoDedup(t *testing.T) {
	qs := NewQuotaStore(newTestDB(t), cacheDefaults(0, 100))

	res := reserveBytes(t, qs, alice1, "hash123", 60)
	if !res.Allowed || res.ID == "" {
		t.Fatalf("expected allowed reservation, got %v", res)
	}
	commitReservation(t, qs, res)

	res2 := reserveBytes(t, qs, alice1, "hash123", 60)
	if !res2.Allowed {
		t.Fatal("expected reservation to be allowed due to repo dedup")
	}
	if res2.ID != "" {
		t.Fatalf("expected empty reservation ID for already committed object, got %q", res2.ID)
	}

	if reserveBytes(t, qs, alice1, "hash456", 60).Allowed {
		t.Fatal("expected reservation to be rejected (exceeds limit)")
	}
}

func TestOwnerUnionDedup(t *testing.T) {
	qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 0))

	res1 := reserveBytes(t, qs, alice1, "hash123", 60)
	if !res1.Allowed {
		t.Fatal("expected res1 to be allowed")
	}
	commitReservation(t, qs, res1)

	if !reserveBytes(t, qs, alice2, "hash123", 60).Allowed {
		t.Fatal("expected res2 to be allowed due to owner union dedup")
	}
	if reserveBytes(t, qs, alice2, "hash456", 50).Allowed {
		t.Fatal("expected res3 to be rejected")
	}
}

func TestDifferentOwnerLogicalCharging(t *testing.T) {
	qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 0))
	bob := quota.Identity{OwnerDID: "did:web:bob", RepoDID: "did:web:bob/repo1"}

	if !reserveBytes(t, qs, alice1, "hash123", 60).Allowed {
		t.Fatal("expected res1 to be allowed")
	}
	if !reserveBytes(t, qs, bob, "hash123", 60).Allowed {
		t.Fatal("expected res2 to be allowed")
	}
}

func TestBothLimits(t *testing.T) {
	qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 50))

	res := reserveBytes(t, qs, alice1, "hash123", 60)
	if res.Allowed || res.Reason != string(quota.ReasonRepoLimit) {
		t.Fatalf("expected rejection due to repo limit, got allowed=%t reason=%s", res.Allowed, res.Reason)
	}
}

func TestDefaultOverrideUnlimitedPrecedence(t *testing.T) {
	ctx := context.Background()
	qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 50))

	if reserveBytes(t, qs, alice1, "hash123", 60).Allowed {
		t.Fatal("expected rejection")
	}

	if err := qs.SetLimit(ctx, alice1.RepoDID, quota.ResourceCacheStorageBytes, 200); err != nil {
		t.Fatal(err)
	}
	if !reserveBytes(t, qs, alice1, "hash123", 60).Allowed {
		t.Fatal("expected allowed after repo override")
	}
	if reserveBytes(t, qs, alice1, "hash456", 120).Allowed {
		t.Fatal("expected rejection due to user limit")
	}

	if err := qs.SetLimit(ctx, alice1.OwnerDID, quota.ResourceCacheStorageBytes, 0); err != nil {
		t.Fatal(err)
	}
	if err := qs.SetLimit(ctx, alice1.RepoDID, quota.ResourceCacheStorageBytes, 0); err != nil {
		t.Fatal(err)
	}

	res := reserveBytes(t, qs, alice1, "hash456", 120)
	if !res.Allowed {
		t.Fatal("expected zero override to make user quota unlimited")
	}
	if res.Reason != quota.ReasonUnlimited {
		t.Fatalf("reason = %q, want %q", res.Reason, quota.ReasonUnlimited)
	}
}

func TestIdempotence(t *testing.T) {
	ctx := context.Background()
	qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 100))

	res1 := reserveBytes(t, qs, alice1, "hash123", 40)
	res2 := reserveBytes(t, qs, alice1, "hash123", 40)
	if res1.ID != res2.ID {
		t.Fatalf("expected identical reservation IDs, got %q and %q", res1.ID, res2.ID)
	}

	steps := []struct {
		name string
		do   func(context.Context, string) error
	}{{"begin commit", qs.BeginCommit}, {"commit", qs.Commit}, {"release", qs.Release}}
	for _, step := range steps {
		for _, id := range []string{res1.ID, res1.ID, ""} {
			if err := step.do(ctx, id); err != nil {
				t.Fatalf("%s %q: %v", step.name, id, err)
			}
		}
	}
}

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

	txCtx, cancel := context.WithCancel(ctx)
	err := qs.withTx(txCtx, func(*sql.Conn) error {
		cancel()
		return txCtx.Err()
	})
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("transaction error = %v, want context canceled", err)
	}

	if err := qs.SetLimit(ctx, "did:web:alice", quota.ResourceWorkflows, 1); err != nil {
		t.Fatalf("write after canceled transaction failed: %v", err)
	}
}

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

	for _, amount := range []int64{60, 80} {
		res := reserveBytes(t, qs, alice1, "nar/object.nar", amount)
		if !res.Allowed || res.ID == "" {
			t.Fatalf("reservation = %+v, want a new allowed reservation", res)
		}
		commitReservation(t, qs, res)
	}

	if amount := committedAmount(t, database, alice1); amount != 80 {
		t.Fatalf("committed amount = %d, want 80", amount)
	}
}

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

	commitReservation(t, qs, reserveBytes(t, qs, alice1, "nar/object.nar", 60))
	interrupted := reserveBytes(t, qs, alice1, "nar/object.nar", 80)
	if !interrupted.Allowed || interrupted.ID == "" {
		t.Fatalf("reservation = %+v, want a new allowed reservation", interrupted)
	}
	if err := qs.BeginCommit(context.Background(), interrupted.ID); err != nil {
		t.Fatal(err)
	}
	if err := qs.Recover(context.Background(), nil); err != nil {
		t.Fatal(err)
	}

	if amount := committedAmount(t, database, alice1); amount != 60 {
		t.Fatalf("recovered amount = %d, want the committed 60", amount)
	}
}

func TestRecoveryPhases(t *testing.T) {
	ctx := context.Background()
	database := newTestDB(t)

	const resource = "gpu_slices"
	qs := NewQuotaStore(database, quota.Defaults{
		quota.ScopeUser: quota.Resources{quota.ResourceCacheStorageBytes: 100, resource: 10},
		quota.ScopeRepo: quota.Resources{quota.ResourceCacheStorageBytes: 100, resource: 10},
	})

	res1 := reserveBytes(t, qs, alice1, "hash_reserved", 10)
	res2 := reserveBytes(t, qs, alice1, "hash_publishing", 20)
	if err := qs.BeginCommit(ctx, res2.ID); err != nil {
		t.Fatal(err)
	}
	res3, err := qs.Reserve(ctx, quota.ReserveRequest{
		Kind:      quota.KindWorkflow,
		Key:       "workflow_preserved",
		Resources: quota.Resources{resource: 1},
		Identity:  alice1,
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := qs.Recover(ctx, []string{res3.ID}); err != nil {
		t.Fatal(err)
	}

	count := func(query string, args ...any) int {
		t.Helper()
		var n int
		if err := database.QueryRowContext(ctx, query, args...).Scan(&n); err != nil {
			t.Fatal(err)
		}
		return n
	}
	if count(`SELECT COUNT(*) FROM quota_allocations WHERE key = 'hash_reserved'`) != 0 {
		t.Fatal("expected reserved object to NOT be committed")
	}
	if count(`SELECT COUNT(*) FROM quota_allocations WHERE key = 'hash_publishing'`) != 0 {
		t.Fatal("expected the publishing reservation to be released, not counted as an object")
	}
	if count(`SELECT COUNT(*) FROM quota_reservations WHERE id = ?`, res3.ID) != 1 {
		t.Fatal("expected preserved reservation to remain in reservations table")
	}
	if count(`SELECT COUNT(*) FROM quota_reservations WHERE id IN (?, ?)`, res1.ID, res2.ID) != 0 {
		t.Fatal("expected unpreserved reservations to be cleaned up")
	}
}

func TestConcurrency(t *testing.T) {
	qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 100))

	var wg sync.WaitGroup
	const numGoroutines = 10
	results := make([]quota.Reservation, numGoroutines)
	errs := make([]error, numGoroutines)

	for i := range numGoroutines {
		wg.Add(1)
		go func() {
			defer wg.Done()
			results[i], errs[i] = qs.Reserve(context.Background(), quota.ReserveRequest{
				Kind:      quota.KindNixCache,
				Key:       "hash_" + string(rune('a'+i)),
				Identity:  alice1,
				Resources: quota.Resources{quota.ResourceCacheStorageBytes: 15},
			})
		}()
	}
	wg.Wait()

	allowedCount := 0
	for i, res := range results {
		if errs[i] != nil {
			t.Fatalf("goroutine %d failed with error: %v", i, errs[i])
		}
		if res.Allowed {
			allowedCount++
		}
	}
	if allowedCount > 6 {
		t.Fatalf("expected at most 6 allowed reservations, got %d", allowedCount)
	}
	if allowedCount == 0 {
		t.Fatal("expected at least one reservation to be allowed")
	}
}

func TestMetricsSnapshot(t *testing.T) {
	ctx := context.Background()
	database := newTestDB(t)
	qs := NewQuotaStore(database, cacheDefaults(0, 100))

	repos := []struct {
		repo     string
		obj      string
		bytes    int64
		override int64
	}{
		{"did:web:alice/repo1", "obj1", 50, -2},
		{"did:web:alice/repo2", "obj2", 80, -2},
		{"did:web:alice/repo3", "obj3", 90, -2},
		{"did:web:alice/repo4", "obj4", 100, -2},
		{"did:web:alice/repo5", "obj5", 120, -2},
		{"did:web:alice/repo6", "obj6", 200, -1},
	}

	exec := func(query string, args ...any) {
		t.Helper()
		if _, err := database.ExecContext(ctx, query, args...); err != nil {
			t.Fatal(err)
		}
	}
	for _, r := range repos {
		exec(`INSERT INTO quota_repo_owners (repo_did, owner_did) VALUES (?, ?)`, r.repo, "did:web:alice")
		exec(`INSERT INTO quota_allocations (repo_did, resource, kind, key, amount) VALUES (?, 'cache_storage_bytes', 'nix_cache', ?, ?)`, r.repo, r.obj, r.bytes)
		if r.override != -2 {
			var val any
			if r.override >= 0 {
				val = r.override
			}
			exec(`INSERT INTO quota_limits (did, resource, max_amount) VALUES (?, 'cache_storage_bytes', ?)`, r.repo, val)
		}
	}
	exec(`INSERT INTO quota_repo_owners (repo_did, owner_did) VALUES (?, ?)`, "did:web:bob/idle", "did:web:bob")

	snapshot, err := qs.MetricsSnapshot(ctx)
	if err != nil {
		t.Fatal(err)
	}

	repoCounts := snapshot.Subjects[quota.ScopeRepo][quota.ResourceCacheStorageBytes]
	for bucket, want := range map[string]int64{"under_limit": 2, "near_limit": 2, "at_limit": 1, "over_limit": 1, "unlimited": 1} {
		if repoCounts[bucket] != want {
			t.Errorf("%s repos = %d, want %d (under_limit includes the idle default-limited repo)", bucket, repoCounts[bucket], want)
		}
	}
}

func TestUserLimitAllowZeroAdditional(t *testing.T) {
	qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 0))

	if !reserveBytes(t, qs, alice1, "hash123", 80).Allowed {
		t.Fatal("expected reservation to be allowed")
	}
	if err := qs.SetLimit(context.Background(), alice1.OwnerDID, quota.ResourceCacheStorageBytes, 50); err != nil {
		t.Fatal(err)
	}
	if !reserveBytes(t, qs, alice1, "hash123", 80).Allowed {
		t.Fatal("expected deduplicated reservation to be allowed even when over limit")
	}
}

func TestUnequalOwnerClaimsRejected(t *testing.T) {
	qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 0))

	res1 := reserveBytes(t, qs, alice1, "hash123", 60)
	if !res1.Allowed || res1.ID == "" {
		t.Fatalf("expected allowed reservation, got %v", res1)
	}
	commitReservation(t, qs, res1)

	if reserveBytes(t, qs, alice2, "hash123", 80).Allowed {
		t.Fatal("expected unequal larger claim to be rejected at user limit")
	}
	if !reserveBytes(t, qs, alice2, "hash123", 60).Allowed {
		t.Fatal("expected equal-amount cross-repo dedup claim to be allowed")
	}
	if !reserveBytes(t, qs, alice2, "hash123", 30).Allowed {
		t.Fatal("expected smaller unequal claim to be allowed since total is within limit")
	}
}

func TestUnequalOwnerClaimsCharged(t *testing.T) {
	qs := NewQuotaStore(newTestDB(t), cacheDefaults(200, 0))

	res1 := reserveBytes(t, qs, alice1, "hash123", 60)
	if !res1.Allowed || res1.ID == "" {
		t.Fatalf("expected allowed reservation, got %v", res1)
	}
	commitReservation(t, qs, res1)

	res2 := reserveBytes(t, qs, alice2, "hash123", 80)
	if !res2.Allowed || res2.ID == "" {
		t.Fatalf("expected allowed unequal claim, got %v", res2)
	}
	commitReservation(t, qs, res2)

	if reserveBytes(t, qs, alice1, "hash456", 70).Allowed {
		t.Fatal("expected new claim to be rejected since user limit has been exceeded by charged unequal claims")
	}
}

func TestLimitArithmeticDoesNotOverflow(t *testing.T) {
	ctx := context.Background()
	database := newTestDB(t)
	qs := NewQuotaStore(database, quota.Defaults{})
	id := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo"}
	const limit = int64(1 << 62)
	if err := qs.SetLimit(ctx, id.RepoDID, "capacity", limit); err != nil {
		t.Fatal(err)
	}
	reserve := func(key string, amount int64) bool {
		t.Helper()
		res, err := qs.Reserve(ctx, quota.ReserveRequest{
			Kind:      quota.KindWorkflow,
			Key:       key,
			Identity:  id,
			Resources: quota.Resources{"capacity": amount},
		})
		if err != nil {
			t.Fatal(err)
		}
		return res.Allowed
	}
	if !reserve("large", limit-1) || reserve("overflow", 2) {
		t.Fatal("overflow-safe limit admission produced the wrong decisions")
	}

	first, err := qs.Reserve(ctx, quota.ReserveRequest{
		Kind:      quota.KindWorkflow,
		Key:       "unlimited-1",
		Identity:  id,
		Resources: quota.Resources{"unlimited": quota.MaxResourceAmount},
	})
	if err != nil || !first.Allowed {
		t.Fatalf("first unlimited reservation = %+v, %v", first, err)
	}
	second, err := qs.Reserve(ctx, quota.ReserveRequest{
		Kind:      quota.KindWorkflow,
		Key:       "unlimited-2",
		Identity:  id,
		Resources: quota.Resources{"unlimited": quota.MaxResourceAmount},
	})
	if err != nil || second.Allowed || !second.Temporary {
		t.Fatalf("overflowing unlimited reservation = %+v, %v", second, err)
	}
	if _, err := qs.ListUsage(ctx); err != nil {
		t.Fatalf("safe maximum usage failed to aggregate: %v", err)
	}
	if _, err := database.ExecContext(ctx, `
		INSERT INTO quota_allocations (repo_did, resource, kind, key, amount)
		VALUES (?, 'unlimited', 'generic_cache', 'forced-overflow', ?)
	`, id.RepoDID, quota.MaxResourceAmount); err != nil {
		t.Fatal(err)
	}
	if _, err := qs.ListUsage(ctx); err == nil {
		t.Fatal("expected pre-existing overflowing usage to return an error")
	}
}