package mill

import (
	"context"
	"io"
	"log/slog"
	"runtime"
	"slices"
	"sync"
	"testing"
	"time"

	"tangled.org/core/notifier"

	"tangled.org/core/spindle/models"
	"tangled.org/core/spindle/nixcache"
	"tangled.org/core/spindle/quota"

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

type nixUse struct {
	cache  string
	hashes []string
}

type nixUseStore struct {
	*scriptedStore
	umu     sync.Mutex
	uses    []nixUse
	gate    chan struct{}
	entered chan struct{}
}

func (s *nixUseStore) RecordNixUse(_ context.Context, cache string, hashes []string) error {
	if s.gate != nil {
		s.entered <- struct{}{}
		<-s.gate
	}
	s.umu.Lock()
	defer s.umu.Unlock()
	s.uses = append(s.uses, nixUse{cache, hashes})
	return nil
}

func (s *nixUseStore) recorded() []nixUse {
	s.umu.Lock()
	defer s.umu.Unlock()
	return slices.Clone(s.uses)
}

func testRepoCache(t *testing.T) string {
	t.Helper()
	cache, err := nixcache.CacheName("did:plc:testrepo")
	if err != nil {
		t.Fatal(err)
	}
	return cache
}

func placeWithNixCache(t *testing.T, store quota.Store) (*Mill, *quotaCandidate, *RemoteLease) {
	t.Helper()
	m, _ := quotaMillWith(t, store, Config{NixCaches: true})
	cand, lease := placeWorkflow(t, m)
	return m, cand, lease
}

func placeWorkflow(t *testing.T, m *Mill) (*quotaCandidate, *RemoteLease) {
	t.Helper()
	cand := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048))
	slot, err := m.place(context.Background(), "dummy", models.WorkflowId{PipelineId: models.PipelineId("rkey1"), Name: "build"}, testWorkflow("build"))
	if err != nil {
		t.Fatalf("place: %v", err)
	}
	t.Cleanup(slot.Release)

	m.mu.Lock()
	var lease *RemoteLease
	for _, l := range m.leases {
		lease = l
	}
	m.mu.Unlock()
	return cand, lease
}

func nixObjectRequest(leaseID, cache string) *millv1.QuotaRequest {
	return &millv1.QuotaRequest{
		RequestId: "req-1",
		Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE,
		Kind:      string(quota.KindNixCache),
		Key:       quota.NixCacheKey(cache, "abc"),
		LeaseId:   leaseID,
		Resources: map[string]int64{"cache_storage_bytes": 4096},
		NixObject: &millv1.NixObject{
			Cache:      cache,
			StoreHash:  "abc",
			StorePath:  "/nix/store/abc-pkg",
			NarSize:    4096,
			References: []string{"def-dep"},
		},
	}
}

func captureQuotaReplies(cand *quotaCandidate) *[]*millv1.QuotaResponse {
	var replies []*millv1.QuotaResponse
	cand.sess.enc = scriptedEncoder(func(msg *millproto.Message) error {
		if resp := msg.GetQuotaResp(); resp != nil {
			replies = append(replies, resp)
		}
		return nil
	})
	return &replies
}

func TestCacheReserveOnAnAtticMillNeedsItsObject(t *testing.T) {
	for _, tt := range []struct {
		name    string
		mutate  func(*millv1.QuotaRequest)
		allowed bool
	}{
		{"no object", func(req *millv1.QuotaRequest) { req.NixObject = nil }, false},
		{"another repos cache", func(req *millv1.QuotaRequest) {
			req.NixObject.Cache = "someone-elses"
			req.Key = quota.NixCacheKey("someone-elses", "abc")
		}, false},
		{"another objects key", func(req *millv1.QuotaRequest) { req.Key = quota.NixCacheKey(req.NixObject.Cache, "other") }, false},
		{"its own object", func(*millv1.QuotaRequest) {}, true},
	} {
		t.Run(tt.name, func(t *testing.T) {
			store := &scriptedStore{verdicts: []quota.Reservation{
				{ID: "res-workflow", Allowed: true, Reason: quota.ReasonWithinLimit},
				{ID: "res-cache", Allowed: true, Reason: quota.ReasonWithinLimit},
			}}
			m, cand, lease := placeWithNixCache(t, store)
			replies := captureQuotaReplies(cand)
			before, _ := store.snapshot()
			req := nixObjectRequest(lease.id, testRepoCache(t))
			tt.mutate(req)

			m.handleQuotaRequest(cand.sess, req)

			if len(*replies) != 1 {
				t.Fatalf("replies = %+v, want one", *replies)
			}
			got := (*replies)[0]
			if got.GetAllowed() != tt.allowed || (got.GetError() == "") != tt.allowed {
				t.Fatalf("reply = %+v, want allowed=%v", got, tt.allowed)
			}
			after, _ := store.snapshot()
			if !tt.allowed {
				if len(after) != len(before) {
					t.Fatalf("store reserved %v for a refused request", after[len(before):])
				}
				return
			}
			want := &quota.NixObject{Cache: testRepoCache(t), StoreHash: "abc", StorePath: "/nix/store/abc-pkg", NarSize: 4096, References: []string{"def-dep"}}
			if obj := after[1].NixObject; obj == nil || obj.Cache != want.Cache || obj.StoreHash != want.StoreHash || obj.StorePath != want.StorePath || obj.NarSize != want.NarSize || !slices.Equal(obj.References, want.References) {
				t.Fatalf("object = %+v, want %+v", obj, want)
			}
		})
	}
}

func TestNixCacheUsedRecordsOnlyTheLeasesOwnCache(t *testing.T) {
	store := &nixUseStore{scriptedStore: &scriptedStore{}}
	m, cand, lease := placeWithNixCache(t, store)

	m.handleNixCacheUsed(cand.sess, &millv1.NixCacheUsed{LeaseId: lease.id, Cache: testRepoCache(t), StoreHashes: []string{"a", "b"}})
	m.handleNixCacheUsed(cand.sess, &millv1.NixCacheUsed{LeaseId: lease.id, Cache: "someone-elses", StoreHashes: []string{"c"}})
	m.handleNixCacheUsed(cand.sess, &millv1.NixCacheUsed{LeaseId: "no-such-lease", Cache: testRepoCache(t), StoreHashes: []string{"d"}})

	m.cfg.NixCaches = false
	m.handleNixCacheUsed(cand.sess, &millv1.NixCacheUsed{LeaseId: lease.id, Cache: testRepoCache(t), StoreHashes: []string{"e"}})

	got := store.recorded()
	if len(got) != 1 || got[0].cache != testRepoCache(t) || !slices.Equal(got[0].hashes, []string{"a", "b"}) {
		t.Fatalf("recorded = %+v, want only the lease's own cache", got)
	}
}

func TestRestoredLeaseKeepsItsReposNixCache(t *testing.T) {
	store := &nixUseStore{scriptedStore: &scriptedStore{verdicts: []quota.Reservation{
		{ID: "res-workflow", Allowed: true, Reason: quota.ReasonWithinLimit},
		{ID: "res-cache", Allowed: true, Reason: quota.ReasonWithinLimit},
	}}}
	m, cand, lease := placeWithNixCache(t, store)
	bdb := m.db

	n := notifier.New()
	restarted := New(slog.New(slog.NewTextHandler(io.Discard, nil)), Config{ReconnectGrace: time.Minute, NixCaches: true})
	qm := quota.NewManager(store, 10*time.Millisecond, nil)
	t.Cleanup(qm.Close)
	restarted.Attach(bdb, &n, qm)
	if err := restarted.RestoreState(); err != nil {
		t.Fatalf("RestoreState: %v", err)
	}
	replies := captureQuotaReplies(cand)

	restarted.handleQuotaRequest(cand.sess, nixObjectRequest(lease.id, testRepoCache(t)))
	restarted.handleQuotaRequest(cand.sess, nixObjectRequest(lease.id, "someone-elses"))
	restarted.handleNixCacheUsed(cand.sess, &millv1.NixCacheUsed{LeaseId: lease.id, Cache: testRepoCache(t), StoreHashes: []string{"a"}})
	restarted.handleNixCacheUsed(cand.sess, &millv1.NixCacheUsed{LeaseId: lease.id, Cache: "someone-elses", StoreHashes: []string{"b"}})

	if len(*replies) != 2 || !(*replies)[0].GetAllowed() || (*replies)[0].GetError() != "" {
		t.Fatalf("own cache reply = %+v, want it allowed", *replies)
	}
	if (*replies)[1].GetAllowed() || (*replies)[1].GetError() == "" {
		t.Fatalf("other cache reply = %+v, want a refusal", (*replies)[1])
	}
	if got := store.recorded(); len(got) != 1 || got[0].cache != testRepoCache(t) || !slices.Equal(got[0].hashes, []string{"a"}) {
		t.Fatalf("recorded = %+v, want only the repo's own cache", got)
	}
}

func TestNixCacheUsedReportsBeyondTheConcurrencyBoundAreDropped(t *testing.T) {
	store := &nixUseStore{scriptedStore: &scriptedStore{}, gate: make(chan struct{}), entered: make(chan struct{}, nixUseConcurrency+1)}
	m, cand, lease := placeWithNixCache(t, store)
	report := &millv1.NixCacheUsed{LeaseId: lease.id, Cache: testRepoCache(t), StoreHashes: []string{"a"}}

	for range nixUseConcurrency {
		m.startNixCacheUsed(cand.sess, report)
		<-store.entered
	}
	m.startNixCacheUsed(cand.sess, report)
	close(store.gate)

	deadline := time.Now().Add(5 * time.Second)
	for len(m.nixUseSlots) != 0 {
		if time.Now().After(deadline) {
			t.Fatal("recording goroutines never finished")
		}
		runtime.Gosched()
	}
	if got := store.recorded(); len(got) != nixUseConcurrency {
		t.Fatalf("recorded %d reports, want %d: the one over the bound is dropped", len(got), nixUseConcurrency)
	}
}