package auth

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"errors"
	"git.yuntuo.hk/k2safe/go-omni/cache"
	"github.com/golang-jwt/jwt/v5"
	"strings"
	"testing"
	"time"
)

func TestJWTValidationAndRotation(t *testing.T) {
	pub, private, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	j, err := NewJWT(JWTConfig{Issuer: "issuer", Audience: "api", Keys: map[string]JWTKey{"first": {Algorithm: "EdDSA", PublicKey: pub}}})
	if err != nil {
		t.Fatal(err)
	}
	claims := JWTClaims{RegisteredClaims: jwt.RegisteredClaims{Issuer: "issuer", Subject: "user1", Audience: jwt.ClaimStrings{"api"}, ExpiresAt: jwt.NewNumericDate(time.Now().Add(time.Hour))}, TenantID: "tenant1", Roles: []string{"reader"}}
	sign := func(c JWTClaims, kid string, method jwt.SigningMethod, key any) string {
		t.Helper()
		token := jwt.NewWithClaims(method, c)
		token.Header["kid"] = kid
		raw, err := token.SignedString(key)
		if err != nil {
			t.Fatal(err)
		}
		return raw
	}
	raw := sign(claims, "first", jwt.SigningMethodEdDSA, private)
	// Constructor clones caller-owned key bytes.
	pub[0] ^= 1
	p, err := j.Authenticate(context.Background(), raw)
	if err != nil || p.TenantID != "tenant1" || !p.HasRole("reader") {
		t.Fatalf("%+v %v", p, err)
	}
	for name, change := range map[string]func(*JWTClaims){"issuer": func(c *JWTClaims) { c.Issuer = "bad" }, "audience": func(c *JWTClaims) { c.Audience = jwt.ClaimStrings{"other"} }, "expired": func(c *JWTClaims) { c.ExpiresAt = jwt.NewNumericDate(time.Now().Add(-time.Hour)) }, "noExpiry": func(c *JWTClaims) { c.ExpiresAt = nil }, "future": func(c *JWTClaims) { c.NotBefore = jwt.NewNumericDate(time.Now().Add(time.Hour)) }, "subject": func(c *JWTClaims) { c.Subject = "" }} {
		t.Run(name, func(t *testing.T) {
			bad := claims
			change(&bad)
			if _, err = j.Authenticate(context.Background(), sign(bad, "first", jwt.SigningMethodEdDSA, private)); !errors.Is(err, ErrInvalidToken) {
				t.Fatal(err)
			}
		})
	}
	for _, invalid := range []string{sign(claims, "missing", jwt.SigningMethodEdDSA, private), sign(claims, "first", jwt.SigningMethodHS256, []byte("invalid secret")), strings.Repeat("x", 8193)} {
		if _, err = j.Authenticate(context.Background(), invalid); !errors.Is(err, ErrInvalidToken) {
			t.Fatal(err)
		}
	}
	if err = j.ReplaceKeys(map[string]JWTKey{"second": {Algorithm: "EdDSA", PublicKey: private.Public()}}); err != nil {
		t.Fatal(err)
	}
	if _, err = j.Authenticate(context.Background(), raw); !errors.Is(err, ErrInvalidToken) {
		t.Fatal(err)
	}
	if _, err = j.Authenticate(context.Background(), sign(claims, "second", jwt.SigningMethodEdDSA, private)); err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err = j.Authenticate(ctx, raw); !errors.Is(err, context.Canceled) {
		t.Fatal(err)
	}
}
func TestSessionsRevokeAndIsolation(t *testing.T) {
	mem, _ := cache.NewMemory(cache.MemoryConfig{})
	defer mem.Close()
	s, err := NewSessions(mem, "session:", time.Minute)
	if err != nil {
		t.Fatal(err)
	}
	ctx := context.Background()
	token, err := s.Issue(ctx, Principal{UserID: "u1", TenantID: "t1", Roles: []string{"reader"}})
	if err != nil {
		t.Fatal(err)
	}
	key, _ := s.key(token)
	if strings.Contains(key, token) {
		t.Fatal("plaintext session secret stored")
	}
	p, err := s.Authenticate(ctx, token)
	if err != nil || p.UserID != "u1" {
		t.Fatal(p, err)
	}
	p.Roles[0] = "admin"
	p, _ = s.Authenticate(ctx, token)
	if p.HasRole("admin") {
		t.Fatal("principal alias")
	}
	if err = s.Revoke(ctx, token); err != nil {
		t.Fatal(err)
	}
	if _, err = s.Authenticate(ctx, token); !errors.Is(err, ErrInvalidToken) {
		t.Fatal(err)
	}
	if _, err = s.Authenticate(ctx, "bad"); !errors.Is(err, ErrInvalidToken) {
		t.Fatal(err)
	}
}