package auth

import (
	"context"
	"crypto"
	"crypto/hmac"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"io"
	"math/big"
	"net/http"
	"strings"
	"testing"
	"time"
)

const (
	testIssuer  = "https://auth.example.test"
	testSubject = "00000000-0000-0000-0000-000000000001"
	testKID     = "test-key"
)

func TestValidateAcceptsValidToken(t *testing.T) {
	env := newTestJWTEnv(t)
	validator, err := NewValidator(env.issuer, env.jwksURL, WithNow(func() time.Time {
		return env.now
	}), WithHTTPClient(env.client))
	if err != nil {
		t.Fatalf("NewValidator() error = %v", err)
	}

	principal, err := validator.Validate(context.Background(), env.token(t, tokenOptions{}))
	if err != nil {
		t.Fatalf("Validate() error = %v", err)
	}

	if principal.Subject != testSubject {
		t.Fatalf("Subject = %q, want %q", principal.Subject, testSubject)
	}
	if principal.Email != "user@example.test" {
		t.Fatalf("Email = %q", principal.Email)
	}
	if principal.Role != "authenticated" {
		t.Fatalf("Role = %q", principal.Role)
	}
}

func TestValidateRejectsExpiredToken(t *testing.T) {
	env := newTestJWTEnv(t)
	validator, err := NewValidator(env.issuer, env.jwksURL, WithNow(func() time.Time {
		return env.now
	}), WithHTTPClient(env.client))
	if err != nil {
		t.Fatalf("NewValidator() error = %v", err)
	}

	_, err = validator.Validate(context.Background(), env.token(t, tokenOptions{
		expiresAt: env.now.Add(-time.Minute),
	}))
	if err == nil {
		t.Fatal("Validate() error = nil, want error")
	}
}

func TestValidateRejectsManipulatedToken(t *testing.T) {
	env := newTestJWTEnv(t)
	validator, err := NewValidator(env.issuer, env.jwksURL, WithNow(func() time.Time {
		return env.now
	}), WithHTTPClient(env.client))
	if err != nil {
		t.Fatalf("NewValidator() error = %v", err)
	}

	token := env.token(t, tokenOptions{})
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		t.Fatalf("token has %d parts, want 3", len(parts))
	}

	var claims map[string]any
	mustDecodeJWTPart(t, parts[1], &claims)
	claims["email"] = "attacker@example.test"
	parts[1] = mustEncodeJSON(t, claims)

	_, err = validator.Validate(context.Background(), strings.Join(parts, "."))
	if err == nil {
		t.Fatal("Validate() error = nil, want error")
	}
}

func TestValidateAcceptsLegacyHS256Token(t *testing.T) {
	now := time.Unix(1_800_000_000, 0).UTC()
	validator, err := NewValidator(testIssuer, "", WithHMACSecret("local-test-secret"), WithNow(func() time.Time {
		return now
	}))
	if err != nil {
		t.Fatalf("NewValidator() error = %v", err)
	}

	token := hs256Token(t, "local-test-secret", now, now.Add(time.Hour))
	principal, err := validator.Validate(context.Background(), token)
	if err != nil {
		t.Fatalf("Validate() error = %v", err)
	}

	if principal.Subject != testSubject {
		t.Fatalf("Subject = %q, want %q", principal.Subject, testSubject)
	}
}

type testJWTEnv struct {
	issuer  string
	jwksURL string
	now     time.Time
	key     *rsa.PrivateKey
	client  *http.Client
}

type tokenOptions struct {
	expiresAt time.Time
}

func newTestJWTEnv(t *testing.T) testJWTEnv {
	t.Helper()

	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("GenerateKey() error = %v", err)
	}

	jwks := jwksResponse{Keys: []jwk{{
		KID: testKID,
		Kty: "RSA",
		Alg: "RS256",
		N:   base64.RawURLEncoding.EncodeToString(privateKey.PublicKey.N.Bytes()),
		E:   base64.RawURLEncoding.EncodeToString(big.NewInt(int64(privateKey.PublicKey.E)).Bytes()),
	}}}

	jwksBody, err := json.Marshal(jwks)
	if err != nil {
		t.Fatalf("Marshal() error = %v", err)
	}

	return testJWTEnv{
		issuer:  testIssuer,
		jwksURL: "https://jwks.example.test",
		now:     time.Unix(1_800_000_000, 0).UTC(),
		key:     privateKey,
		client: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
			return &http.Response{
				StatusCode: http.StatusOK,
				Header:     make(http.Header),
				Body:       io.NopCloser(strings.NewReader(string(jwksBody))),
				Request:    req,
			}, nil
		})},
	}
}

func (e testJWTEnv) token(t *testing.T, opts tokenOptions) string {
	t.Helper()

	expiresAt := opts.expiresAt
	if expiresAt.IsZero() {
		expiresAt = e.now.Add(time.Hour)
	}

	header := mustEncodeJSON(t, map[string]any{
		"alg": "RS256",
		"kid": testKID,
		"typ": "JWT",
	})
	claims := mustEncodeJSON(t, map[string]any{
		"iss":   e.issuer,
		"sub":   testSubject,
		"email": "user@example.test",
		"role":  "authenticated",
		"iat":   e.now.Unix(),
		"exp":   expiresAt.Unix(),
	})

	signingInput := header + "." + claims
	digest := sha256.Sum256([]byte(signingInput))
	signature, err := rsa.SignPKCS1v15(rand.Reader, e.key, crypto.SHA256, digest[:])
	if err != nil {
		t.Fatalf("SignPKCS1v15() error = %v", err)
	}

	return signingInput + "." + base64.RawURLEncoding.EncodeToString(signature)
}

func mustEncodeJSON(t *testing.T, value any) string {
	t.Helper()

	data, err := json.Marshal(value)
	if err != nil {
		t.Fatalf("Marshal() error = %v", err)
	}
	return base64.RawURLEncoding.EncodeToString(data)
}

func mustDecodeJWTPart(t *testing.T, part string, out any) {
	t.Helper()

	data, err := base64.RawURLEncoding.DecodeString(part)
	if err != nil {
		t.Fatalf("DecodeString() error = %v", err)
	}
	if err := json.Unmarshal(data, out); err != nil {
		t.Fatalf("Unmarshal() error = %v", err)
	}
}

func hs256Token(t *testing.T, secret string, now time.Time, expiresAt time.Time) string {
	t.Helper()

	header := mustEncodeJSON(t, map[string]any{
		"alg": "HS256",
		"typ": "JWT",
	})
	claims := mustEncodeJSON(t, map[string]any{
		"iss":   testIssuer,
		"sub":   testSubject,
		"email": "user@example.test",
		"role":  "authenticated",
		"iat":   now.Unix(),
		"exp":   expiresAt.Unix(),
	})
	signingInput := header + "." + claims

	mac := hmac.New(sha256.New, []byte(secret))
	_, _ = mac.Write([]byte(signingInput))
	return signingInput + "." + base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
}

type roundTripFunc func(*http.Request) (*http.Response, error)

func (fn roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
	return fn(req)
}