package token

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"strings"
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

func generateTestKey(t *testing.T) (ed25519.PublicKey, ed25519.PrivateKey) {
	t.Helper()
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("generate key: %v", err)
	}
	return pub, priv
}

// b64url encodes a string as base64url without padding.
func b64url(s string) string {
	return base64.RawURLEncoding.EncodeToString([]byte(s))
}

const testIssuer = "https://auth.example.com"

func TestIssueAndValidateToken(t *testing.T) {
	pub, priv := generateTestKey(t)
	roles := []string{"admin", "reader"}

	tokenStr, claims, err := IssueToken(priv, testIssuer, "user-uuid-1", roles, time.Hour)
	if err != nil {
		t.Fatalf("IssueToken: %v", err)
	}
	if tokenStr == "" {
		t.Fatal("IssueToken returned empty token string")
	}
	if claims.JTI == "" {
		t.Error("JTI must not be empty")
	}
	if claims.Subject != "user-uuid-1" {
		t.Errorf("Subject = %q, want %q", claims.Subject, "user-uuid-1")
	}

	// Validate the token.
	got, err := ValidateToken(pub, tokenStr, testIssuer)
	if err != nil {
		t.Fatalf("ValidateToken: %v", err)
	}
	if got.Subject != "user-uuid-1" {
		t.Errorf("validated Subject = %q, want %q", got.Subject, "user-uuid-1")
	}
	if got.JTI != claims.JTI {
		t.Errorf("validated JTI = %q, want %q", got.JTI, claims.JTI)
	}
	if len(got.Roles) != 2 {
		t.Errorf("validated Roles = %v, want 2 roles", got.Roles)
	}
}

// TestValidateTokenWrongAlgorithm verifies that tokens with non-EdDSA alg are
// rejected immediately, before any signature verification.
// Security: This tests the core defence against algorithm-confusion attacks.
func TestValidateTokenWrongAlgorithm(t *testing.T) {
	_, priv := generateTestKey(t)
	pub, _ := generateTestKey(t) // different key — but alg check should fail first

	// Forge a token signed with HMAC-SHA256 (alg: HS256).
	hmacToken := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
		"iss": testIssuer,
		"sub": "attacker",
		"iat": time.Now().Unix(),
		"exp": time.Now().Add(time.Hour).Unix(),
		"jti": "fake-jti",
	})
	// Use the Ed25519 public key bytes as the HMAC secret (classic alg confusion).
	hs256Signed, err := hmacToken.SignedString([]byte(priv.Public().(ed25519.PublicKey)))
	if err != nil {
		t.Fatalf("sign HS256 token: %v", err)
	}

	_, err = ValidateToken(pub, hs256Signed, testIssuer)
	if err == nil {
		t.Fatal("expected error for HS256 token, got nil")
	}
	if !errors.Is(err, ErrWrongAlgorithm) {
		t.Errorf("expected ErrWrongAlgorithm, got: %v", err)
	}
}

// TestValidateTokenAlgNone verifies that "none" algorithm is rejected.
// Security: "none" algorithm tokens have no signature and must always be
// rejected regardless of payload content.
func TestValidateTokenAlgNone(t *testing.T) {
	pub, _ := generateTestKey(t)

	// Construct a "none" algorithm token manually.
	// golang-jwt/v5 disallows signing with "none" directly, so we craft it
	// using raw base64url encoding.
	header := `{"alg":"none","typ":"JWT"}`
	payload := `{"iss":"https://auth.example.com","sub":"evil","iat":1000000,"exp":9999999999,"jti":"evil-jti"}`
	noneToken := b64url(header) + "." + b64url(payload) + "."

	_, err := ValidateToken(pub, noneToken, testIssuer)
	if err == nil {
		t.Fatal("expected error for 'none' algorithm token, got nil")
	}
}

// TestValidateTokenExpired verifies that expired tokens are rejected.
func TestValidateTokenExpired(t *testing.T) {
	pub, priv := generateTestKey(t)

	// Issue a token with a negative expiry (already expired).
	tokenStr, _, err := IssueToken(priv, testIssuer, "user", nil, -time.Minute)
	if err != nil {
		t.Fatalf("IssueToken: %v", err)
	}

	_, err = ValidateToken(pub, tokenStr, testIssuer)
	if err == nil {
		t.Fatal("expected error for expired token, got nil")
	}
	if !errors.Is(err, ErrExpiredToken) {
		t.Errorf("expected ErrExpiredToken, got: %v", err)
	}
}

// TestValidateTokenTamperedSignature verifies that signature tampering is caught.
func TestValidateTokenTamperedSignature(t *testing.T) {
	pub, priv := generateTestKey(t)

	tokenStr, _, err := IssueToken(priv, testIssuer, "user", nil, time.Hour)
	if err != nil {
		t.Fatalf("IssueToken: %v", err)
	}

	// Tamper: flip a byte in the signature (last segment).
	parts := strings.Split(tokenStr, ".")
	if len(parts) != 3 {
		t.Fatalf("unexpected token format: %d parts", len(parts))
	}
	sigBytes, err := base64.RawURLEncoding.DecodeString(parts[2])
	if err != nil {
		t.Fatalf("decode signature: %v", err)
	}
	sigBytes[0] ^= 0x01 // flip one bit
	parts[2] = base64.RawURLEncoding.EncodeToString(sigBytes)
	tampered := strings.Join(parts, ".")

	_, err = ValidateToken(pub, tampered, testIssuer)
	if err == nil {
		t.Fatal("expected error for tampered signature, got nil")
	}
}

// TestValidateTokenWrongKey verifies that a token signed with a different key
// is rejected.
func TestValidateTokenWrongKey(t *testing.T) {
	_, priv := generateTestKey(t)
	wrongPub, _ := generateTestKey(t)

	tokenStr, _, err := IssueToken(priv, testIssuer, "user", nil, time.Hour)
	if err != nil {
		t.Fatalf("IssueToken: %v", err)
	}

	_, err = ValidateToken(wrongPub, tokenStr, testIssuer)
	if err == nil {
		t.Fatal("expected error for wrong key, got nil")
	}
}

// TestValidateTokenWrongIssuer verifies that tokens from a different issuer
// are rejected.
func TestValidateTokenWrongIssuer(t *testing.T) {
	pub, priv := generateTestKey(t)

	tokenStr, _, err := IssueToken(priv, "https://evil.example.com", "user", nil, time.Hour)
	if err != nil {
		t.Fatalf("IssueToken: %v", err)
	}

	_, err = ValidateToken(pub, tokenStr, testIssuer)
	if err == nil {
		t.Fatal("expected error for wrong issuer, got nil")
	}
}

// TestJTIsAreUnique verifies that two issued tokens have different JTIs.
func TestJTIsAreUnique(t *testing.T) {
	_, priv := generateTestKey(t)

	_, c1, err := IssueToken(priv, testIssuer, "user", nil, time.Hour)
	if err != nil {
		t.Fatalf("IssueToken (1): %v", err)
	}
	_, c2, err := IssueToken(priv, testIssuer, "user", nil, time.Hour)
	if err != nil {
		t.Fatalf("IssueToken (2): %v", err)
	}
	if c1.JTI == c2.JTI {
		t.Error("two issued tokens have the same JTI")
	}
}

// TestClaimsHasRole verifies role checking.
func TestClaimsHasRole(t *testing.T) {
	c := &Claims{Roles: []string{"admin", "reader"}}
	if !c.HasRole("admin") {
		t.Error("expected HasRole(admin) = true")
	}
	if !c.HasRole("reader") {
		t.Error("expected HasRole(reader) = true")
	}
	if c.HasRole("writer") {
		t.Error("expected HasRole(writer) = false")
	}
}