package identity

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"errors"
	"fmt"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/novox/mesh-controller/internal/store"
)

func fresh(t *testing.T) *Identity {
	t.Helper()
	admin := os.Getenv("MESH_TEST_POSTGRES")
	if admin == "" {
		t.Skip("no MESH_TEST_POSTGRES; run `make check` to raise one")
	}
	name := fmt.Sprintf("ident_%d", time.Now().UnixNano()%10_000_000)

	conn, err := pgx.Connect(t.Context(), admin)
	if err != nil {
		t.Fatalf("cannot reach the test PostgreSQL: %v", err)
	}
	if _, err := conn.Exec(t.Context(), "create database "+name); err != nil {
		t.Fatalf("cannot create %s: %v", name, err)
	}
	conn.Close(t.Context())

	cut := strings.LastIndex(admin, "/")
	t.Setenv(store.Variable(Name), admin[:cut]+"/"+name+"?sslmode=disable")

	ident, err := Open(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		ident.Close()
		c, err := pgx.Connect(context.Background(), admin)
		if err != nil {
			return
		}
		defer c.Close(context.Background())
		_, _ = c.Exec(context.Background(), "drop database if exists "+name+" with (force)")
	})
	if err := ident.Ready(t.Context(), 20*time.Second); err != nil {
		t.Fatal(err)
	}
	migrations, err := Migrations()
	if err != nil {
		t.Fatal(err)
	}
	if _, err := ident.store.Migrate(t.Context(), migrations); err != nil {
		t.Fatal(err)
	}
	return ident
}

func TestNoKeyIsAnErrorRatherThanAnEmptyKey(t *testing.T) {
	// Signing with nothing, or with a key invented on the spot, produces declarations every
	// existing node correctly refuses — and that refusal looks like a compromise rather than a
	// control plane that lost its key.
	ident := fresh(t)
	if _, err := ident.Active(t.Context()); !errors.Is(err, ErrNoSigningKey) {
		t.Fatalf("expected ErrNoSigningKey, got %v", err)
	}
	if _, err := ident.Sign(t.Context(), []byte("anything")); !errors.Is(err, ErrNoSigningKey) {
		t.Fatalf("signing without a key gave %v", err)
	}
}

func TestEstablishingTwiceKeepsTheFirstKey(t *testing.T) {
	// The control plane runs this at every start. A second key generated by a restart is a mesh
	// whose nodes all hold the wrong public half — every declaration refused, by every node,
	// with nothing visibly having gone wrong.
	ident := fresh(t)
	first, err := ident.Establish(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	second, err := ident.Establish(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	if first.ID != second.ID || string(first.Public) != string(second.Public) {
		t.Error("a second Establish replaced the signing key; every node would hold the wrong one")
	}
}

func TestTwoProcessesStartingTogetherAgreeOnOneKey(t *testing.T) {
	// A restart while another copy is coming up. Both find nothing and both generate; only one
	// insert may survive, and the loser must read back the winner rather than return the key it
	// generated and did not store.
	ident := fresh(t)

	var wg sync.WaitGroup
	keys := make([]SigningKey, 6)
	errs := make([]error, 6)
	for i := range keys {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			keys[i], errs[i] = ident.Establish(context.Background())
		}(i)
	}
	wg.Wait()

	for i, err := range errs {
		if err != nil {
			t.Fatalf("establish %d failed: %v", i, err)
		}
	}
	for i, k := range keys {
		if k.ID != keys[0].ID {
			t.Errorf("establish %d got key %s, establish 0 got %s — they disagree", i, k.ID, keys[0].ID)
		}
	}

	var count int
	if err := ident.store.Pool().QueryRow(t.Context(),
		`select count(*) from signing_key`).Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != 1 {
		t.Errorf("%d signing keys exist; exactly one may be active", count)
	}
}

func TestASignatureVerifiesAgainstThePublicHalfThatTravels(t *testing.T) {
	// The whole point: a node holds only the public half, from a token it may have received
	// months ago, and must be able to tell a real declaration from a forged one.
	ident := fresh(t)
	key, err := ident.Establish(t.Context())
	if err != nil {
		t.Fatal(err)
	}

	declaration := []byte(`{"declaration":1,"resources":[]}`)
	signature, err := ident.Sign(t.Context(), declaration)
	if err != nil {
		t.Fatal(err)
	}
	if !Verify(key.Public, declaration, signature) {
		t.Fatal("a declaration this control plane signed did not verify against the key it hands out")
	}
}

func TestATamperedDeclarationDoesNotVerify(t *testing.T) {
	// Since the host applies whatever the link delivers, a forged declaration is the whole
	// machine. This is the check that stands between those two facts.
	ident := fresh(t)
	key, err := ident.Establish(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	signature, err := ident.Sign(t.Context(), []byte(`{"resources":["harmless"]}`))
	if err != nil {
		t.Fatal(err)
	}
	if Verify(key.Public, []byte(`{"resources":["something else entirely"]}`), signature) {
		t.Fatal("a signature made over one declaration verified against a different one")
	}
}

func TestAnotherControlPlanesSignatureIsRefused(t *testing.T) {
	// "This is not from the mesh I joined" — the case ADR 0004 requires a host to tell apart
	// from "this is malformed".
	mine := fresh(t)
	theirs := fresh(t)
	myKey, err := mine.Establish(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	if _, err := theirs.Establish(t.Context()); err != nil {
		t.Fatal(err)
	}

	declaration := []byte(`{"declaration":1}`)
	theirSignature, err := theirs.Sign(t.Context(), declaration)
	if err != nil {
		t.Fatal(err)
	}
	if Verify(myKey.Public, declaration, theirSignature) {
		t.Fatal("a signature from a different control plane verified against this one's key")
	}
}

func TestTheFingerprintIsOfThePublicHalf(t *testing.T) {
	ident := fresh(t)
	key, err := ident.Establish(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	if len(key.Fingerprint()) != 64 {
		t.Errorf("fingerprint is %q", key.Fingerprint())
	}
	// And it must not be derivable from something that is not the key.
	if key.Fingerprint() == (SigningKey{Public: make([]byte, 32)}).Fingerprint() {
		t.Error("the fingerprint does not depend on the key")
	}
}

func TestANodeIsVerifiedByAKeyItGenerated(t *testing.T) {
	// The whole of proving a node is that node. The mesh holds only the public half, so this
	// verification is the same operation the node performs on every declaration, in reverse.
	ident := fresh(t)
	public, private, err := ed25519.GenerateKey(nil)
	if err != nil {
		t.Fatal(err)
	}
	node := uuid(t)

	if _, err := ident.RecordNodeKey(t.Context(), node, public); err != nil {
		t.Fatal(err)
	}
	challenge := []byte("prove you are that node")
	if err := ident.VerifyNode(t.Context(), node, challenge, ed25519.Sign(private, challenge)); err != nil {
		t.Fatalf("a node signing with its own key was refused: %v", err)
	}
}

func TestAnotherMachinesKeyDoesNotIdentifyThisNode(t *testing.T) {
	ident := fresh(t)
	mine, _, err := ed25519.GenerateKey(nil)
	if err != nil {
		t.Fatal(err)
	}
	_, theirPrivate, err := ed25519.GenerateKey(nil)
	if err != nil {
		t.Fatal(err)
	}
	node := uuid(t)
	if _, err := ident.RecordNodeKey(t.Context(), node, mine); err != nil {
		t.Fatal(err)
	}

	challenge := []byte("prove you are that node")
	err = ident.VerifyNode(t.Context(), node, challenge, ed25519.Sign(theirPrivate, challenge))
	if !errors.Is(err, ErrNotThisNode) {
		t.Fatalf("a different machine's signature was accepted: %v", err)
	}
}

func TestReEnrolmentRevokesTheMachineItReplaced(t *testing.T) {
	// novox/hq ADR 0004's stolen-laptop case. Two live identities for one node record means the
	// machine that was replaced goes on being believed, which is the thing revocation exists for.
	ident := fresh(t)
	node := uuid(t)
	oldPublic, oldPrivate, err := ed25519.GenerateKey(nil)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := ident.RecordNodeKey(t.Context(), node, oldPublic); err != nil {
		t.Fatal(err)
	}

	newPublic, newPrivate, err := ed25519.GenerateKey(nil)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := ident.RecordNodeKey(t.Context(), node, newPublic); err != nil {
		t.Fatal(err)
	}

	challenge := []byte("still me?")
	if err := ident.VerifyNode(t.Context(), node, challenge, ed25519.Sign(oldPrivate, challenge)); !errors.Is(err, ErrNotThisNode) {
		t.Error("the replaced machine is still believed")
	}
	if err := ident.VerifyNode(t.Context(), node, challenge, ed25519.Sign(newPrivate, challenge)); err != nil {
		t.Errorf("the new machine was refused: %v", err)
	}
}

func TestOnlyOneKeyIsEverLiveForANode(t *testing.T) {
	// Enforced by the database rather than by the order of two statements, because the window
	// between them is exactly when two live identities would exist.
	ident := fresh(t)
	node := uuid(t)
	for i := 0; i < 4; i++ {
		public, _, err := ed25519.GenerateKey(nil)
		if err != nil {
			t.Fatal(err)
		}
		if _, err := ident.RecordNodeKey(t.Context(), node, public); err != nil {
			t.Fatal(err)
		}
	}

	var live int
	if err := ident.store.Pool().QueryRow(t.Context(),
		`select count(*) from node_key where node = $1 and revoked is null`, node).Scan(&live); err != nil {
		t.Fatal(err)
	}
	if live != 1 {
		t.Errorf("%d live keys for one node; exactly one may be", live)
	}
}

func TestOnlyThePublicHalfIsEverStored(t *testing.T) {
	// The property that makes a copy of this database worthless to whoever takes it: it is a list
	// of who to believe, not a set of credentials.
	ident := fresh(t)
	public, private, err := ed25519.GenerateKey(nil)
	if err != nil {
		t.Fatal(err)
	}
	node := uuid(t)
	if _, err := ident.RecordNodeKey(t.Context(), node, public); err != nil {
		t.Fatal(err)
	}

	var stored []byte
	if err := ident.store.Pool().QueryRow(t.Context(),
		`select public from node_key where node = $1`, node).Scan(&stored); err != nil {
		t.Fatal(err)
	}
	if len(stored) != ed25519.PublicKeySize {
		t.Errorf("stored %d bytes for a node key; a public key is %d and a private key is %d",
			len(stored), ed25519.PublicKeySize, ed25519.PrivateKeySize)
	}
	if bytes.Contains(private, stored) && bytes.Contains(stored, private[:32]) {
		t.Error("what is stored looks like part of the private key")
	}
}

func TestAnUnknownNodeAndARevokedKeyAreRefusedAlike(t *testing.T) {
	ident := fresh(t)
	_, private, err := ed25519.GenerateKey(nil)
	if err != nil {
		t.Fatal(err)
	}
	challenge := []byte("hello")
	err = ident.VerifyNode(t.Context(), uuid(t), challenge, ed25519.Sign(private, challenge))
	if !errors.Is(err, ErrNotThisNode) {
		t.Fatalf("an unknown node gave %v", err)
	}
}

func TestSomethingThatIsNotAKeyIsRefused(t *testing.T) {
	ident := fresh(t)
	if _, err := ident.RecordNodeKey(t.Context(), uuid(t), []byte("far too short")); err == nil {
		t.Fatal("a truncated key was recorded as a node's identity")
	}
}

// uuid gives each test its own node id. The node records live in another context's database
// (novox/hq ADR 0008), so there is nothing here to reference and nothing to create.
func uuid(t *testing.T) string {
	t.Helper()
	var id string
	if err := freshConn(t).QueryRow(t.Context(), `select gen_random_uuid()`).Scan(&id); err != nil {
		t.Fatal(err)
	}
	return id
}

func freshConn(t *testing.T) *pgx.Conn {
	t.Helper()
	conn, err := pgx.Connect(t.Context(), os.Getenv("MESH_TEST_POSTGRES"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { conn.Close(context.Background()) })
	return conn
}

func TestTwoEnrolmentsAtOnceStillLeaveOneLiveKey(t *testing.T) {
	// The case the database constraint is for, and the only one: sequential calls are already
	// safe because RecordNodeKey revokes before it inserts. Two at once both revoke what they
	// found and both insert, and without the partial unique index the node ends with two live
	// identities — the replaced machine still believed, which is the whole thing revocation
	// exists to prevent.
	//
	// Some of these are expected to fail. What must not happen is two of them succeeding.
	ident := fresh(t)
	node := uuid(t)

	var wg sync.WaitGroup
	errs := make([]error, 6)
	for i := range errs {
		public, _, err := ed25519.GenerateKey(nil)
		if err != nil {
			t.Fatal(err)
		}
		wg.Add(1)
		go func(i int, public ed25519.PublicKey) {
			defer wg.Done()
			_, errs[i] = ident.RecordNodeKey(context.Background(), node, public)
		}(i, public)
	}
	wg.Wait()

	var live int
	if err := ident.store.Pool().QueryRow(t.Context(),
		`select count(*) from node_key where node = $1 and revoked is null`, node).Scan(&live); err != nil {
		t.Fatal(err)
	}
	if live != 1 {
		t.Errorf("%d live keys after six concurrent enrolments; exactly one may be live", live)
	}

	succeeded := 0
	for _, err := range errs {
		if err == nil {
			succeeded++
		}
	}
	if succeeded == 0 {
		t.Error("every concurrent enrolment failed; at least one must win")
	}
}