package mill

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"io"
	"log/slog"
	"testing"

	"github.com/bluesky-social/indigo/atproto/syntax"
	gossh "golang.org/x/crypto/ssh"

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

type fakeDebugChecker struct {
	allow bool
	repos []syntax.DID
}

func (f *fakeDebugChecker) DebugKeyAllowed(_ context.Context, repo syntax.DID, _ gossh.PublicKey) bool {
	f.repos = append(f.repos, repo)
	return f.allow
}

func wireKey(t *testing.T) []byte {
	t.Helper()
	pub, _, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	key, err := gossh.NewPublicKey(pub)
	if err != nil {
		t.Fatal(err)
	}
	return key.Marshal()
}

type debugKeyEnv struct {
	m       *Mill
	sess    *millSession
	checker *fakeDebugChecker
	replies []*millv1.DebugKeyResponse
}

func newDebugKeyEnv(t *testing.T, allow bool) *debugKeyEnv {
	t.Helper()
	e := &debugKeyEnv{checker: &fakeDebugChecker{allow: allow}}
	e.m = New(slog.New(slog.NewTextHandler(io.Discard, nil)), Config{})
	e.m.SetDebugKeyChecker(e.checker)
	e.sess = newSession("node-a", "epoch-a", nil, scriptedEncoder(func(msg *millproto.Message) error {
		e.replies = append(e.replies, msg.GetDebugKeyResp())
		return nil
	}), slog.New(slog.NewTextHandler(io.Discard, nil)))
	lease := newLease("lease-1", "node-a", "epoch-a", "microvm")
	lease.repoDID = "did:plc:leaserepo"
	e.m.leases["lease-1"] = lease
	return e
}

func (e *debugKeyEnv) ask(t *testing.T, leaseID string, key []byte) *millv1.DebugKeyResponse {
	t.Helper()
	e.replies = nil
	e.m.handleDebugKeyRequest(e.sess, &millv1.DebugKeyRequest{RequestId: "r1", LeaseId: leaseID, PublicKey: key})
	if len(e.replies) != 1 || e.replies[0].GetRequestId() != "r1" {
		t.Fatalf("replies = %v, want exactly one for r1", e.replies)
	}
	return e.replies[0]
}

func TestMillAnswersADebugKeyFromTheLeasesRepo(t *testing.T) {
	for _, allow := range []bool{true, false} {
		e := newDebugKeyEnv(t, allow)
		resp := e.ask(t, "lease-1", wireKey(t))
		if resp.GetAllowed() != allow || resp.GetError() != "" {
			t.Errorf("checker said %v, mill answered %+v", allow, resp)
		}
		if len(e.checker.repos) != 1 || e.checker.repos[0] != "did:plc:leaserepo" {
			t.Errorf("checker asked about %v, want the lease's repo", e.checker.repos)
		}
	}
}

func TestMillRefusesADebugKeyForALeaseItDoesNotKnow(t *testing.T) {
	e := newDebugKeyEnv(t, true)
	if resp := e.ask(t, "no-such-lease", wireKey(t)); resp.GetAllowed() || resp.GetError() == "" {
		t.Errorf("unknown lease answered %+v, want a refusal with an error", resp)
	}

	// another node's lease is not this executor's to ask about
	other := newLease("lease-2", "node-b", "epoch-b", "microvm")
	other.repoDID = "did:plc:otherrepo"
	e.m.leases["lease-2"] = other
	if resp := e.ask(t, "lease-2", wireKey(t)); resp.GetAllowed() {
		t.Errorf("another node's lease answered %+v, want a refusal", resp)
	}

	// same node after a reconnect under a new epoch
	stale := newLease("lease-3", "node-a", "old-epoch", "microvm")
	stale.repoDID = "did:plc:otherrepo"
	e.m.leases["lease-3"] = stale
	if resp := e.ask(t, "lease-3", wireKey(t)); resp.GetAllowed() {
		t.Errorf("a stale epoch answered %+v, want a refusal", resp)
	}
	if len(e.checker.repos) != 0 {
		t.Errorf("checker consulted for %v, want it left alone", e.checker.repos)
	}
}

func TestMillRefusesADebugKeyItCannotCheck(t *testing.T) {
	e := newDebugKeyEnv(t, true)
	if resp := e.ask(t, "lease-1", []byte("not a key")); resp.GetAllowed() || resp.GetError() == "" {
		t.Errorf("garbage key answered %+v, want a refusal with an error", resp)
	}

	e.m.SetDebugKeyChecker(nil)
	if resp := e.ask(t, "lease-1", wireKey(t)); resp.GetAllowed() || resp.GetError() == "" {
		t.Errorf("no checker answered %+v, want a refusal with an error", resp)
	}
}