package auth

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"errors"
	"math/big"
	"net/http"
	"net/http/httptest"
	"time"

	"github.com/MicahParks/jwkset"
	itkent "github.com/tidbcloud/devops/entity/itoken"

	jwt "github.com/golang-jwt/jwt/v5"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
)

type mockKeyfunc struct {
	key any
	err error
	ctx context.Context
}

func (m *mockKeyfunc) Keyfunc(token *jwt.Token) (any, error) {
	return m.key, m.err
}

func (m *mockKeyfunc) KeyfuncCtx(ctx context.Context) jwt.Keyfunc {
	m.ctx = ctx
	return func(token *jwt.Token) (any, error) {
		return m.Keyfunc(token)
	}
}

func (m *mockKeyfunc) Storage() jwkset.Storage {
	return nil
}

func (m *mockKeyfunc) VerificationKeySet(ctx context.Context) (jwt.VerificationKeySet, error) {
	return jwt.VerificationKeySet{}, nil
}

var _ = Describe("[TokenValidator]", func() {
	var (
		secret []byte
		ctx    context.Context
	)

	BeforeEach(func() {
		secret = []byte("secret")
		ctx = context.Background()
	})

	signHS256Token := func(claims jwt.MapClaims) string {
		token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
		signed, err := token.SignedString(secret)
		Expect(err).To(BeNil())
		return signed
	}

	Describe("ValidateToken", func() {
		It("validates token and returns claims", func() {
			keyFunc := &mockKeyfunc{key: secret}
			validator := &tokenValidator{
				keyFunc:  keyFunc,
				audience: "devops",
			}

			claims := jwt.MapClaims{
				"aud":   "devops",
				"email": "owner@example.com",
				"exp":   time.Now().Add(time.Hour).Unix(),
			}
			tokenString := signHS256Token(claims)

			result, err := validator.ValidateToken(ctx, tokenString)
			Expect(err).To(BeNil())
			Expect(result["aud"]).To(Equal("devops"))
			Expect(result["email"]).To(Equal("owner@example.com"))

			email, err := GetEmailClaim(result, "email")
			Expect(err).To(BeNil())
			Expect(email).To(Equal("owner@example.com"))
		})

		It("returns error on audience mismatch", func() {
			keyFunc := &mockKeyfunc{key: secret}
			validator := &tokenValidator{
				keyFunc:  keyFunc,
				audience: "devops",
			}

			claims := jwt.MapClaims{
				"aud": "other",
				"exp": time.Now().Add(time.Hour).Unix(),
			}
			tokenString := signHS256Token(claims)

			_, err := validator.ValidateToken(ctx, tokenString)
			Expect(err).To(MatchError("audience mismatch"))
		})

		It("returns error when keyfunc fails", func() {
			keyFunc := &mockKeyfunc{err: errors.New("boom")}
			validator := &tokenValidator{
				keyFunc: keyFunc,
			}

			claims := jwt.MapClaims{
				"exp": time.Now().Add(time.Hour).Unix(),
			}
			tokenString := signHS256Token(claims)

			_, err := validator.ValidateToken(ctx, tokenString)
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("failed to parse token"))
		})
	})

	Describe("NewTokenValidator with JWKSURL", func() {
		signRS256Token := func(key *rsa.PrivateKey, kid string, claims jwt.MapClaims) string {
			token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
			token.Header["kid"] = kid
			signed, err := token.SignedString(key)
			Expect(err).ToNot(HaveOccurred())
			return signed
		}

		buildJWK := func(pub *rsa.PublicKey, kid string) itkent.JWK {
			return itkent.JWK{
				Kty: "RSA",
				Use: "sig",
				Alg: jwt.SigningMethodRS256.Alg(),
				Kid: kid,
				N:   base64.RawURLEncoding.EncodeToString(pub.N.Bytes()),
				E:   base64.RawURLEncoding.EncodeToString(big.NewInt(int64(pub.E)).Bytes()),
			}
		}

		It("validates token against standard jwks response", func() {
			key, err := rsa.GenerateKey(rand.Reader, 2048)
			Expect(err).ToNot(HaveOccurred())

			jwksSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				_ = json.NewEncoder(w).Encode(map[string]any{
					"keys": []any{buildJWK(&key.PublicKey, "kid-1")},
				})
			}))
			defer jwksSrv.Close()

			validator, err := NewTokenValidator(TokenValidatorOptions{
				JWKSURL:  jwksSrv.URL,
				Audience: "jump-gateway",
			})
			Expect(err).ToNot(HaveOccurred())

			token := signRS256Token(key, "kid-1", jwt.MapClaims{
				"aud":   "jump-gateway",
				"exp":   time.Now().Add(time.Hour).Unix(),
				"email": "owner@example.com",
			})

			claims, err := validator.ValidateToken(context.Background(), token)
			Expect(err).ToNot(HaveOccurred())
			Expect(claims["aud"]).To(Equal("jump-gateway"))
		})

		It("returns error when audience mismatches in jwksurl mode", func() {
			key, err := rsa.GenerateKey(rand.Reader, 2048)
			Expect(err).ToNot(HaveOccurred())

			jwksSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				_ = json.NewEncoder(w).Encode(map[string]any{
					"keys": []any{buildJWK(&key.PublicKey, "kid-1")},
				})
			}))
			defer jwksSrv.Close()

			validator, err := NewTokenValidator(TokenValidatorOptions{
				JWKSURL:  jwksSrv.URL,
				Audience: "jump-gateway",
			})
			Expect(err).ToNot(HaveOccurred())

			token := signRS256Token(key, "kid-1", jwt.MapClaims{
				"aud": "another-aud",
				"exp": time.Now().Add(time.Hour).Unix(),
			})

			_, err = validator.ValidateToken(context.Background(), token)
			Expect(err).To(MatchError("audience mismatch"))
		})
	})

	Describe("audienceMatches", func() {
		It("matches string audience", func() {
			Expect(audienceMatches("devops", "devops")).To(BeTrue())
			Expect(audienceMatches("other", "devops")).To(BeFalse())
		})

		It("matches slice of strings audience", func() {
			Expect(audienceMatches([]string{"devops", "other"}, "devops")).To(BeTrue())
			Expect(audienceMatches([]string{"other"}, "devops")).To(BeFalse())
		})

		It("matches slice of any audience", func() {
			Expect(audienceMatches([]any{"devops", 1}, "devops")).To(BeTrue())
			Expect(audienceMatches([]any{"other"}, "devops")).To(BeFalse())
		})

		It("returns false for unsupported audience types", func() {
			Expect(audienceMatches(42, "devops")).To(BeFalse())
		})
	})

	Describe("GetEmailClaim", func() {
		It("returns error when email is missing", func() {
			_, err := GetEmailClaim(jwt.MapClaims{"sub": "user"}, "email")
			Expect(err).To(MatchError("no email found"))
		})
	})
})