package authn

import (
	"context"
	"crypto/ed25519"
	"errors"
	"testing"
	"time"

	"sourcecraft.dev/bigbes/sr-ht-core/auth"
	"github.com/dolthub/dolt/go/libraries/doltcore/creds"
)

const testAud = "dolt.srht.bigb.es"

type stubKey struct {
	pubkey   []byte
	username string
}

type stubKeyStore struct {
	keys     map[string]stubKey
	touched  []string
	byKIDErr error
	touchErr error
}

func (s *stubKeyStore) ByKID(ctx context.Context, kid string) ([]byte, string, error) {
	if s.byKIDErr != nil {
		return nil, "", s.byKIDErr
	}
	k, ok := s.keys[kid]
	if !ok {
		return nil, "", errors.New("no such key")
	}
	return k.pubkey, k.username, nil
}

func (s *stubKeyStore) TouchLastUsed(ctx context.Context, kid string) error {
	if s.touchErr != nil {
		return s.touchErr
	}
	s.touched = append(s.touched, kid)
	return nil
}

// doltKeypair generates an Ed25519 keypair and its dolt key id, and returns a
// KeyStore that recognises it plus a MetaBackend that knows its owner.
func doltKeypair(t *testing.T, username string) (ed25519.PrivateKey, string, *stubKeyStore) {
	t.Helper()
	pub, priv, err := ed25519.GenerateKey(nil)
	if err != nil {
		t.Fatal(err)
	}
	kid := creds.PubKeyToKIDStr(pub)
	ks := &stubKeyStore{keys: map[string]stubKey{
		kid: {pubkey: pub, username: username},
	}}
	return priv, kid, ks
}

func validSub(kid string) string { return jwtSubjectPrefix + kid }

func TestResolveDoltJWT_Valid(t *testing.T) {
	withStubBackend(t, &stubBackend{users: map[string]auth.AuthContext{
		"bigbes": sampleUser(1, "bigbes", auth.USER_TYPE_USER),
	}})
	priv, kid, ks := doltKeypair(t, "bigbes")

	raw := mintDoltJWT(t, priv, kid, testAud, creds.ClientIssuer, validSub(kid), time.Now().Add(30*time.Second))
	ac, err := ResolveDoltJWT(testCtx(), raw, testAud, ks)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if ac.Username != "bigbes" || ac.UserID != 1 {
		t.Fatalf("wrong user resolved: %+v", ac)
	}
	if ac.AuthMethod != AuthMethodDoltKey {
		t.Fatalf("AuthMethod = %q, want %q", ac.AuthMethod, AuthMethodDoltKey)
	}
	if len(ks.touched) != 1 || ks.touched[0] != kid {
		t.Fatalf("expected TouchLastUsed(%q), got %v", kid, ks.touched)
	}
}

func TestResolveDoltJWT_IgnoresIssuer(t *testing.T) {
	withStubBackend(t, &stubBackend{users: map[string]auth.AuthContext{
		"bigbes": sampleUser(1, "bigbes", auth.USER_TYPE_USER),
	}})
	priv, kid, ks := doltKeypair(t, "bigbes")

	// Deliberately weird issuer; must still be accepted (iss is ignored).
	raw := mintDoltJWT(t, priv, kid, testAud, "whatever.example.com", validSub(kid), time.Now().Add(30*time.Second))
	if _, err := ResolveDoltJWT(testCtx(), raw, testAud, ks); err != nil {
		t.Fatalf("issuer must be ignored, got %v", err)
	}
}

func TestResolveDoltJWT_WrongKey(t *testing.T) {
	withStubBackend(t, &stubBackend{users: map[string]auth.AuthContext{
		"bigbes": sampleUser(1, "bigbes", auth.USER_TYPE_USER),
	}})
	_, kid, ks := doltKeypair(t, "bigbes")

	// Sign with a different private key than the one registered for kid.
	_, wrongPriv, err := ed25519.GenerateKey(nil)
	if err != nil {
		t.Fatal(err)
	}
	raw := mintDoltJWT(t, wrongPriv, kid, testAud, creds.ClientIssuer, validSub(kid), time.Now().Add(30*time.Second))
	_, err = ResolveDoltJWT(testCtx(), raw, testAud, ks)
	if !errors.Is(err, ErrInvalidToken) {
		t.Fatalf("wrong signing key must wrap ErrInvalidToken, got %v", err)
	}
}

func TestResolveDoltJWT_Expired(t *testing.T) {
	withStubBackend(t, &stubBackend{users: map[string]auth.AuthContext{
		"bigbes": sampleUser(1, "bigbes", auth.USER_TYPE_USER),
	}})
	priv, kid, ks := doltKeypair(t, "bigbes")

	// Expired well beyond the leeway window.
	raw := mintDoltJWT(t, priv, kid, testAud, creds.ClientIssuer, validSub(kid), time.Now().Add(-10*time.Minute))
	_, err := ResolveDoltJWT(testCtx(), raw, testAud, ks)
	if !errors.Is(err, ErrInvalidToken) {
		t.Fatalf("expired token must wrap ErrInvalidToken, got %v", err)
	}
}

func TestResolveDoltJWT_WrongAudience(t *testing.T) {
	withStubBackend(t, &stubBackend{users: map[string]auth.AuthContext{
		"bigbes": sampleUser(1, "bigbes", auth.USER_TYPE_USER),
	}})
	priv, kid, ks := doltKeypair(t, "bigbes")

	raw := mintDoltJWT(t, priv, kid, "some.other.host", creds.ClientIssuer, validSub(kid), time.Now().Add(30*time.Second))
	_, err := ResolveDoltJWT(testCtx(), raw, testAud, ks)
	if !errors.Is(err, ErrInvalidToken) {
		t.Fatalf("wrong audience must wrap ErrInvalidToken, got %v", err)
	}
}

func TestResolveDoltJWT_BadSubject(t *testing.T) {
	withStubBackend(t, &stubBackend{users: map[string]auth.AuthContext{
		"bigbes": sampleUser(1, "bigbes", auth.USER_TYPE_USER),
	}})
	priv, kid, ks := doltKeypair(t, "bigbes")

	raw := mintDoltJWT(t, priv, kid, testAud, creds.ClientIssuer, "doltClientCredentials/wrong", time.Now().Add(30*time.Second))
	_, err := ResolveDoltJWT(testCtx(), raw, testAud, ks)
	if !errors.Is(err, ErrInvalidToken) {
		t.Fatalf("bad subject must wrap ErrInvalidToken, got %v", err)
	}
}

func TestResolveDoltJWT_UnknownKID(t *testing.T) {
	withStubBackend(t, &stubBackend{users: map[string]auth.AuthContext{
		"bigbes": sampleUser(1, "bigbes", auth.USER_TYPE_USER),
	}})
	priv, kid, _ := doltKeypair(t, "bigbes")
	emptyKS := &stubKeyStore{keys: map[string]stubKey{}}

	raw := mintDoltJWT(t, priv, kid, testAud, creds.ClientIssuer, validSub(kid), time.Now().Add(30*time.Second))
	_, err := ResolveDoltJWT(testCtx(), raw, testAud, emptyKS)
	if !errors.Is(err, ErrInvalidToken) {
		t.Fatalf("unknown kid must wrap ErrInvalidToken, got %v", err)
	}
}

func TestResolveDoltJWT_KIDPubkeyMismatch(t *testing.T) {
	withStubBackend(t, &stubBackend{users: map[string]auth.AuthContext{
		"bigbes": sampleUser(1, "bigbes", auth.USER_TYPE_USER),
	}})
	priv, kid, _ := doltKeypair(t, "bigbes")

	// Register the kid against an unrelated public key: integrity check fails.
	otherPub, _, err := ed25519.GenerateKey(nil)
	if err != nil {
		t.Fatal(err)
	}
	badKS := &stubKeyStore{keys: map[string]stubKey{kid: {pubkey: otherPub, username: "bigbes"}}}

	raw := mintDoltJWT(t, priv, kid, testAud, creds.ClientIssuer, validSub(kid), time.Now().Add(30*time.Second))
	_, err = ResolveDoltJWT(testCtx(), raw, testAud, badKS)
	if !errors.Is(err, ErrInvalidToken) {
		t.Fatalf("kid/pubkey mismatch must wrap ErrInvalidToken, got %v", err)
	}
}

func TestResolveDoltJWT_TouchFailureIsNonFatal(t *testing.T) {
	withStubBackend(t, &stubBackend{users: map[string]auth.AuthContext{
		"bigbes": sampleUser(1, "bigbes", auth.USER_TYPE_USER),
	}})
	priv, kid, ks := doltKeypair(t, "bigbes")
	ks.touchErr = errors.New("db write failed")

	raw := mintDoltJWT(t, priv, kid, testAud, creds.ClientIssuer, validSub(kid), time.Now().Add(30*time.Second))
	if _, err := ResolveDoltJWT(testCtx(), raw, testAud, ks); err != nil {
		t.Fatalf("TouchLastUsed failure must not fail auth, got %v", err)
	}
}

func TestResolveDoltJWT_UserLookupFailureIsTransient(t *testing.T) {
	withStubBackend(t, &stubBackend{lookupErr: errBackendDown})
	priv, kid, ks := doltKeypair(t, "bigbes")

	raw := mintDoltJWT(t, priv, kid, testAud, creds.ClientIssuer, validSub(kid), time.Now().Add(30*time.Second))
	_, err := ResolveDoltJWT(testCtx(), raw, testAud, ks)
	if err == nil {
		t.Fatal("expected an error")
	}
	if errors.Is(err, ErrInvalidToken) {
		t.Fatalf("user-lookup failure must be transient, got %v", err)
	}
}

// PubKeyToKIDStr is dolt's own derivation; this asserts our generated keypair's
// kid is self-consistent (the value production and tests both rely on).
func TestKIDDerivationMatchesDolt(t *testing.T) {
	dc, err := creds.GenerateCredentials()
	if err != nil {
		t.Fatal(err)
	}
	if got := creds.PubKeyToKIDStr(dc.PubKey); got != dc.KeyIDBase32Str() {
		t.Fatalf("PubKeyToKIDStr(%x) = %q, want %q", dc.PubKey, got, dc.KeyIDBase32Str())
	}
}