package executor

import (
	"context"
	"errors"
	"fmt"
	"math/rand/v2"
	"sync"
	"testing"
	"time"

	"tangled.org/core/spindle/quota"

	millproto "tangled.org/core/spindle/mill/proto"
	millv1 "tangled.org/core/spindle/mill/proto/gen"
)

var testIdentity = quota.Identity{
	OwnerDID: "did:web:owner",
	RepoDID:  "did:web:repo",
}

func validReserveRequest(kind quota.Kind, key string, resource string, amount int64) quota.ReserveRequest {
	return quota.ReserveRequest{
		Kind:     kind,
		Key:      key,
		Identity: testIdentity,
		Resources: quota.Resources{
			resource: amount,
		},
	}
}

func TestQuotaClientReserveAndTransitions(t *testing.T) {
	client := NewQuotaClient()
	leaseID := "lease-abc-123"
	store := client.ForLease(leaseID)

	var (
		mu   sync.Mutex
		sent []*millproto.Message
	)

	client.SetSendFn(func(msg *millproto.Message) error {
		mu.Lock()
		sent = append(sent, msg)
		mu.Unlock()

		req := msg.GetQuotaReq()
		if req == nil {
			t.Error("expected QuotaReq message")
			return nil
		}

		go func() {
			switch req.GetOperation() {
			case millv1.QuotaOperation_QUOTA_OPERATION_RESERVE:
				client.HandleResponse(req.GetRequestId(), &millv1.QuotaResponse{
					RequestId:     req.GetRequestId(),
					ReservationId: "res-123",
					Allowed:       true,
					Reason:        "within_limits",
					Resource:      quota.ResourceCacheStorageBytes,
				})
			case millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT,
				millv1.QuotaOperation_QUOTA_OPERATION_COMMIT,
				millv1.QuotaOperation_QUOTA_OPERATION_RELEASE:
				client.HandleResponse(req.GetRequestId(), &millv1.QuotaResponse{
					RequestId: req.GetRequestId(),
					Allowed:   true,
				})
			}
		}()
		return nil
	})

	ctx := context.Background()

	res, err := store.Reserve(ctx, validReserveRequest(quota.KindNixCache, "nar-hash-1", quota.ResourceCacheStorageBytes, 2048))
	if err != nil {
		t.Fatalf("Reserve() unexpected error: %v", err)
	}
	if res.ID != "res-123" {
		t.Errorf("res.ID = %q, want %q", res.ID, "res-123")
	}
	if !res.Allowed {
		t.Error("res.Allowed = false, want true")
	}
	if res.Temporary {
		t.Error("res.Temporary = true, want false")
	}
	if res.Reason != "within_limits" {
		t.Errorf("res.Reason = %q, want %q", res.Reason, "within_limits")
	}
	if res.Resource != quota.ResourceCacheStorageBytes {
		t.Errorf("res.Resource = %q, want %q", res.Resource, quota.ResourceCacheStorageBytes)
	}

	mu.Lock()
	if len(sent) != 1 {
		t.Fatalf("sent %d messages, want 1", len(sent))
	}
	reqMsg := sent[0].GetQuotaReq()
	if reqMsg.GetOperation() != millv1.QuotaOperation_QUOTA_OPERATION_RESERVE {
		t.Errorf("operation = %v, want RESERVE", reqMsg.GetOperation())
	}
	if reqMsg.GetLeaseId() != leaseID {
		t.Errorf("lease_id = %q, want %q", reqMsg.GetLeaseId(), leaseID)
	}
	if reqMsg.GetKind() != string(quota.KindNixCache) {
		t.Errorf("kind = %q, want %q", reqMsg.GetKind(), quota.KindNixCache)
	}
	if reqMsg.GetKey() != "nar-hash-1" {
		t.Errorf("key = %q, want %q", reqMsg.GetKey(), "nar-hash-1")
	}
	if reqMsg.GetResources()["cache_storage_bytes"] != 2048 {
		t.Errorf("resources.cache_storage_bytes = %d, want 2048", reqMsg.GetResources()["cache_storage_bytes"])
	}
	mu.Unlock()

	if err := store.BeginCommit(ctx, "res-123"); err != nil {
		t.Fatalf("BeginCommit() unexpected error: %v", err)
	}

	mu.Lock()
	if len(sent) != 2 {
		t.Fatalf("sent %d messages, want 2", len(sent))
	}
	bcMsg := sent[1].GetQuotaReq()
	if bcMsg.GetOperation() != millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT {
		t.Errorf("operation = %v, want BEGIN_COMMIT", bcMsg.GetOperation())
	}
	if bcMsg.GetLeaseId() != leaseID {
		t.Errorf("lease_id = %q, want %q", bcMsg.GetLeaseId(), leaseID)
	}
	if bcMsg.GetReservationId() != "res-123" {
		t.Errorf("reservation_id = %q, want %q", bcMsg.GetReservationId(), "res-123")
	}
	mu.Unlock()

	if err := store.Commit(ctx, "res-123"); err != nil {
		t.Fatalf("Commit() unexpected error: %v", err)
	}

	mu.Lock()
	if len(sent) != 3 {
		t.Fatalf("sent %d messages, want 3", len(sent))
	}
	cMsg := sent[2].GetQuotaReq()
	if cMsg.GetOperation() != millv1.QuotaOperation_QUOTA_OPERATION_COMMIT {
		t.Errorf("operation = %v, want COMMIT", cMsg.GetOperation())
	}
	if cMsg.GetLeaseId() != leaseID {
		t.Errorf("lease_id = %q, want %q", cMsg.GetLeaseId(), leaseID)
	}
	if cMsg.GetReservationId() != "res-123" {
		t.Errorf("reservation_id = %q, want %q", cMsg.GetReservationId(), "res-123")
	}
	mu.Unlock()

	if err := store.Release(ctx, "res-123"); err != nil {
		t.Fatalf("Release() unexpected error: %v", err)
	}

	mu.Lock()
	if len(sent) != 4 {
		t.Fatalf("sent %d messages, want 4", len(sent))
	}
	rMsg := sent[3].GetQuotaReq()
	if rMsg.GetOperation() != millv1.QuotaOperation_QUOTA_OPERATION_RELEASE {
		t.Errorf("operation = %v, want RELEASE", rMsg.GetOperation())
	}
	if rMsg.GetLeaseId() != leaseID {
		t.Errorf("lease_id = %q, want %q", rMsg.GetLeaseId(), leaseID)
	}
	if rMsg.GetReservationId() != "res-123" {
		t.Errorf("reservation_id = %q, want %q", rMsg.GetReservationId(), "res-123")
	}
	mu.Unlock()

	if err := store.BeginCommit(ctx, ""); err != nil {
		t.Errorf("BeginCommit(\"\") error = %v", err)
	}
	if err := store.Commit(ctx, ""); err != nil {
		t.Errorf("Commit(\"\") error = %v", err)
	}
	if err := store.Release(ctx, ""); err != nil {
		t.Errorf("Release(\"\") error = %v", err)
	}

	mu.Lock()
	if len(sent) != 4 {
		t.Errorf("sent %d messages after empty transitions, want 4", len(sent))
	}
	mu.Unlock()
}

func TestQuotaClientGenericValidation(t *testing.T) {
	client := NewQuotaClient()
	store := client.ForLease("lease-1")

	client.SetSendFn(func(msg *millproto.Message) error {
		t.Error("sendFn must not be called for invalid requests")
		return nil
	})

	ctx := context.Background()

	invalidRequests := []struct {
		name string
		req  quota.ReserveRequest
	}{
		{
			name: "invalid kind",
			req: quota.ReserveRequest{
				Kind:      quota.Kind("invalid-kind"),
				Key:       "key-1",
				Identity:  testIdentity,
				Resources: quota.Resources{quota.ResourceCacheStorageBytes: 100},
			},
		},
		{
			name: "empty key",
			req: quota.ReserveRequest{
				Kind:      quota.KindNixCache,
				Key:       "",
				Identity:  testIdentity,
				Resources: quota.Resources{quota.ResourceCacheStorageBytes: 100},
			},
		},
		{
			name: "empty resources",
			req: quota.ReserveRequest{
				Kind:      quota.KindNixCache,
				Key:       "key-1",
				Identity:  testIdentity,
				Resources: nil,
			},
		},
		{
			name: "negative resource amount",
			req: quota.ReserveRequest{
				Kind:      quota.KindNixCache,
				Key:       "key-1",
				Identity:  testIdentity,
				Resources: quota.Resources{quota.ResourceCacheStorageBytes: -50},
			},
		},
		{
			name: "empty owner DID",
			req: quota.ReserveRequest{
				Kind:      quota.KindNixCache,
				Key:       "key-1",
				Identity:  quota.Identity{OwnerDID: "", RepoDID: "did:web:repo"},
				Resources: quota.Resources{quota.ResourceCacheStorageBytes: 100},
			},
		},
		{
			name: "empty repo DID",
			req: quota.ReserveRequest{
				Kind:      quota.KindNixCache,
				Key:       "key-1",
				Identity:  quota.Identity{OwnerDID: "did:web:owner", RepoDID: ""},
				Resources: quota.Resources{quota.ResourceCacheStorageBytes: 100},
			},
		},
		{
			name: "invalid owner DID prefix",
			req: quota.ReserveRequest{
				Kind:      quota.KindNixCache,
				Key:       "key-1",
				Identity:  quota.Identity{OwnerDID: "not-a-did", RepoDID: "did:web:repo"},
				Resources: quota.Resources{quota.ResourceCacheStorageBytes: 100},
			},
		},
		{
			name: "invalid repo DID prefix",
			req: quota.ReserveRequest{
				Kind:      quota.KindNixCache,
				Key:       "key-1",
				Identity:  quota.Identity{OwnerDID: "did:web:owner", RepoDID: "not-a-did"},
				Resources: quota.Resources{quota.ResourceCacheStorageBytes: 100},
			},
		},
	}

	for _, tc := range invalidRequests {
		t.Run(tc.name, func(t *testing.T) {
			_, err := store.Reserve(ctx, tc.req)
			if err == nil {
				t.Fatalf("Reserve() expected validation error for %s", tc.name)
			}
		})
	}

	validKinds := []quota.Kind{quota.KindWorkflow, quota.KindNixCache, quota.KindGenericCache}
	for _, k := range validKinds {
		t.Run("valid kind "+string(k), func(t *testing.T) {
			called := false
			client.SetSendFn(func(msg *millproto.Message) error {
				called = true
				go client.HandleResponse(msg.GetQuotaReq().GetRequestId(), &millv1.QuotaResponse{
					RequestId:     msg.GetQuotaReq().GetRequestId(),
					ReservationId: "res-ok",
					Allowed:       true,
				})
				return nil
			})
			req := validReserveRequest(k, "key-valid", "custom_units", 100)
			res, err := store.Reserve(ctx, req)
			if err != nil {
				t.Fatalf("unexpected validation failure for kind %s: %v", k, err)
			}
			if !called || !res.Allowed {
				t.Fatalf("expected request to be sent and allowed for kind %s", k)
			}
		})
	}
}

func TestQuotaClientResourceResponsePropagation(t *testing.T) {
	ctx := context.Background()

	t.Run("resource roundtrip", func(t *testing.T) {
		client := NewQuotaClient()
		store := client.ForLease("lease-1")

		client.SetSendFn(func(msg *millproto.Message) error {
			go client.HandleResponse(msg.GetQuotaReq().GetRequestId(), &millv1.QuotaResponse{
				RequestId:     msg.GetQuotaReq().GetRequestId(),
				ReservationId: "res-mem",
				Allowed:       true,
				Reason:        "within_limits",
				Resource:      quota.ResourceMemoryMiB,
			})
			return nil
		})

		res, err := store.Reserve(ctx, validReserveRequest(quota.KindWorkflow, "wf-1", quota.ResourceMemoryMiB, 4096))
		if err != nil {
			t.Fatalf("Reserve() error = %v", err)
		}
		if res.Resource != quota.ResourceMemoryMiB {
			t.Errorf("res.Resource = %q, want %q", res.Resource, quota.ResourceMemoryMiB)
		}
		if res.ID != "res-mem" {
			t.Errorf("res.ID = %q, want %q", res.ID, "res-mem")
		}
	})

	t.Run("temporary denial", func(t *testing.T) {
		client := NewQuotaClient()
		store := client.ForLease("lease-1")

		client.SetSendFn(func(msg *millproto.Message) error {
			go client.HandleResponse(msg.GetQuotaReq().GetRequestId(), &millv1.QuotaResponse{
				RequestId: msg.GetQuotaReq().GetRequestId(),
				Allowed:   false,
				Temporary: true,
				Reason:    "repo_queue_full",
			})
			return nil
		})

		res, err := store.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 500))
		if err != nil {
			t.Fatalf("Reserve() error = %v", err)
		}
		if res.Allowed {
			t.Error("res.Allowed = true, want false")
		}
		if !res.Temporary {
			t.Error("res.Temporary = false, want true")
		}
		if res.Reason != "repo_queue_full" {
			t.Errorf("res.Reason = %q, want %q", res.Reason, "repo_queue_full")
		}
	})

	t.Run("permanent denial", func(t *testing.T) {
		client := NewQuotaClient()
		store := client.ForLease("lease-1")

		client.SetSendFn(func(msg *millproto.Message) error {
			go client.HandleResponse(msg.GetQuotaReq().GetRequestId(), &millv1.QuotaResponse{
				RequestId: msg.GetQuotaReq().GetRequestId(),
				Allowed:   false,
				Temporary: false,
				Reason:    "user_limit_exceeded",
			})
			return nil
		})

		res, err := store.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 500))
		if err != nil {
			t.Fatalf("Reserve() error = %v", err)
		}
		if res.Allowed {
			t.Error("res.Allowed = true, want false")
		}
		if res.Temporary {
			t.Error("res.Temporary = true, want false")
		}
		if res.Reason != "user_limit_exceeded" {
			t.Errorf("res.Reason = %q, want %q", res.Reason, "user_limit_exceeded")
		}
	})

	t.Run("allowed and temporary contradictory response", func(t *testing.T) {
		client := NewQuotaClient()
		store := client.ForLease("lease-1")

		client.SetSendFn(func(msg *millproto.Message) error {
			go client.HandleResponse(msg.GetQuotaReq().GetRequestId(), &millv1.QuotaResponse{
				RequestId:     msg.GetQuotaReq().GetRequestId(),
				ReservationId: "res-bad",
				Allowed:       true,
				Temporary:     true,
			})
			return nil
		})

		_, err := store.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 500))
		if err == nil {
			t.Fatal("expected error for allowed+temporary response")
		}
	})

	t.Run("mill error response", func(t *testing.T) {
		client := NewQuotaClient()
		store := client.ForLease("lease-1")

		client.SetSendFn(func(msg *millproto.Message) error {
			go client.HandleResponse(msg.GetQuotaReq().GetRequestId(), &millv1.QuotaResponse{
				RequestId: msg.GetQuotaReq().GetRequestId(),
				Error:     "quota store database failure",
			})
			return nil
		})

		_, err := store.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 500))
		if err == nil || err.Error() != "quota store database failure" {
			t.Fatalf("expected error %q, got %v", "quota store database failure", err)
		}
	})
}

func TestQuotaClientLeaseBinding(t *testing.T) {
	client := NewQuotaClient()
	ctx := context.Background()

	var (
		mu         sync.Mutex
		seenLeases []string
	)

	client.SetSendFn(func(msg *millproto.Message) error {
		req := msg.GetQuotaReq()
		mu.Lock()
		seenLeases = append(seenLeases, req.GetLeaseId())
		mu.Unlock()

		go client.HandleResponse(req.GetRequestId(), &millv1.QuotaResponse{
			RequestId:     req.GetRequestId(),
			ReservationId: "res-" + req.GetLeaseId(),
			Allowed:       true,
		})
		return nil
	})

	s1 := client.ForLease("lease-alpha")
	s2 := client.ForLease("lease-beta")

	_, err := s1.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-a", quota.ResourceCacheStorageBytes, 100))
	if err != nil {
		t.Fatalf("s1.Reserve() error = %v", err)
	}

	_, err = s2.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-b", quota.ResourceCacheStorageBytes, 200))
	if err != nil {
		t.Fatalf("s2.Reserve() error = %v", err)
	}

	mu.Lock()
	if len(seenLeases) != 2 || seenLeases[0] != "lease-alpha" || seenLeases[1] != "lease-beta" {
		t.Fatalf("seen leases = %v, want [lease-alpha, lease-beta]", seenLeases)
	}
	mu.Unlock()

	sEmpty := client.ForLease("")
	if _, err := sEmpty.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-c", quota.ResourceCacheStorageBytes, 100)); err == nil {
		t.Error("sEmpty.Reserve() expected error")
	}
	if err := sEmpty.BeginCommit(ctx, "res-1"); err == nil {
		t.Error("sEmpty.BeginCommit() expected error")
	}
	if err := sEmpty.Commit(ctx, "res-1"); err == nil {
		t.Error("sEmpty.Commit() expected error")
	}
	if err := sEmpty.Release(ctx, "res-1"); err == nil {
		t.Error("sEmpty.Release() expected error")
	}
}

func TestQuotaClientDisconnect(t *testing.T) {
	t.Run("fails closed when not connected", func(t *testing.T) {
		client := NewQuotaClient()
		store := client.ForLease("lease-1")
		ctx := context.Background()

		if _, err := store.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 100)); err == nil {
			t.Error("Reserve() expected error when not connected")
		}
		if err := store.BeginCommit(ctx, "res-1"); err == nil {
			t.Error("BeginCommit() expected error when not connected")
		}
		if err := store.Commit(ctx, "res-1"); err == nil {
			t.Error("Commit() expected error when not connected")
		}
		if err := store.Release(ctx, "res-1"); err == nil {
			t.Error("Release() expected error when not connected")
		}
	})

	t.Run("in-flight unblocked by OnDisconnect", func(t *testing.T) {
		client := NewQuotaClient()
		store := client.ForLease("lease-1")
		sent := make(chan struct{})

		client.SetSendFn(func(msg *millproto.Message) error {
			close(sent)
			return nil
		})

		done := make(chan error, 1)
		go func() {
			_, err := store.Reserve(context.Background(), validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 100))
			done <- err
		}()

		<-sent
		client.OnDisconnect()

		select {
		case err := <-done:
			if err == nil || err.Error() != "quota client disconnected" {
				t.Fatalf("expected %q error, got %v", "quota client disconnected", err)
			}
		case <-time.After(2 * time.Second):
			t.Fatal("timed out waiting for Reserve() to return after OnDisconnect")
		}
	})

	t.Run("late duplicate responses", func(t *testing.T) {
		client := NewQuotaClient()
		client.HandleResponse("unknown-id", &millv1.QuotaResponse{
			RequestId:     "unknown-id",
			ReservationId: "res-late",
			Allowed:       true,
		})
	})

	t.Run("reconnect works", func(t *testing.T) {
		client := NewQuotaClient()
		store := client.ForLease("lease-1")

		client.SetSendFn(func(msg *millproto.Message) error {
			return errors.New("temporary failure")
		})

		_, err := store.Reserve(context.Background(), validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 100))
		if err == nil {
			t.Fatal("expected error on failed send")
		}

		client.OnDisconnect()

		client.SetSendFn(func(msg *millproto.Message) error {
			go client.HandleResponse(msg.GetQuotaReq().GetRequestId(), &millv1.QuotaResponse{
				RequestId:     msg.GetQuotaReq().GetRequestId(),
				ReservationId: "res-reconnected",
				Allowed:       true,
			})
			return nil
		})

		res, err := store.Reserve(context.Background(), validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 100))
		if err != nil {
			t.Fatalf("Reserve() after reconnect error = %v", err)
		}
		if res.ID != "res-reconnected" {
			t.Errorf("res.ID = %q, want %q", res.ID, "res-reconnected")
		}
	})

	t.Run("response vs disconnect race", func(t *testing.T) {
		for range 20 {
			client := NewQuotaClient()
			store := client.ForLease("lease-1")

			client.SetSendFn(func(msg *millproto.Message) error {
				return nil
			})

			done := make(chan error, 1)
			go func() {
				_, err := store.Reserve(context.Background(), validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 100))
				done <- err
			}()

			time.Sleep(time.Duration(rand.IntN(5)) * time.Millisecond)
			client.OnDisconnect()

			select {
			case err := <-done:
				if err == nil {
					t.Fatal("expected an error on disconnect")
				}
			case <-time.After(2 * time.Second):
				t.Fatal("timed out waiting for Reserve() during disconnect race")
			}
		}
	})
}

func TestQuotaClientContextCancellation(t *testing.T) {
	client := NewQuotaClient()
	store := client.ForLease("lease-1")
	sent := make(chan struct{})

	client.SetSendFn(func(msg *millproto.Message) error {
		close(sent)
		return nil
	})

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)

	go func() {
		_, err := store.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 100))
		done <- err
	}()

	<-sent
	cancel()

	select {
	case err := <-done:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("expected context.Canceled, got %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for Reserve() to return on context cancellation")
	}

	client.mu.Lock()
	pendingCount := len(client.pending)
	client.mu.Unlock()
	if pendingCount != 0 {
		t.Errorf("pending map has %d entries, want 0", pendingCount)
	}
}

func TestQuotaClientConcurrentCorrelations(t *testing.T) {
	client := NewQuotaClient()
	store := client.ForLease("lease-1")

	client.SetSendFn(func(msg *millproto.Message) error {
		req := msg.GetQuotaReq()
		go func(reqID, key string) {
			time.Sleep(time.Duration(rand.IntN(10)) * time.Millisecond)
			client.HandleResponse(reqID, &millv1.QuotaResponse{
				RequestId:     reqID,
				ReservationId: "res-" + key,
				Allowed:       true,
				Reason:        "within_limits",
				Resource:      quota.ResourceCacheStorageBytes,
			})
		}(req.GetRequestId(), req.GetKey())
		return nil
	})

	const count = 20
	var wg sync.WaitGroup
	wg.Add(count)

	for i := range count {
		go func(idx int) {
			defer wg.Done()
			key := fmt.Sprintf("key-%d", idx)
			res, err := store.Reserve(context.Background(), validReserveRequest(quota.KindNixCache, key, quota.ResourceCacheStorageBytes, int64(idx+1)*100))
			if err != nil {
				t.Errorf("goroutine %d Reserve() error: %v", idx, err)
				return
			}
			expectedID := "res-" + key
			if res.ID != expectedID {
				t.Errorf("goroutine %d res.ID = %q, want %q", idx, res.ID, expectedID)
			}
		}(i)
	}

	wg.Wait()
}