package itoken

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"math/big"
	"testing"
	"time"

	jwt "github.com/golang-jwt/jwt/v5"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
	itkent "github.com/tidbcloud/devops/entity/itoken"
	"github.com/tidbcloud/devops/pkg/infra/config"
	"github.com/tidbcloud/devops/pkg/infra/utils/pointer"
	"github.com/tidbcloud/devops/pkg/infra/utils/test"
	itksvcapi "github.com/tidbcloud/devops/pkg/service/itoken/api"
)

func TestInternalTokenService(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Internal Token Service Suite", test.ReporterConfig())
}

var _ = Describe("[InternalTokenService]", func() {
	var (
		ctx           context.Context
		_             = test.Framework()
		rsaKey        *rsa.PrivateKey
		rsaPEM        string
		ecdsaPEM      string
		expIssuer     string
		audience      string
		serviceConfig config.InternalTokenServiceConfig
	)

	BeforeEach(func() {
		ctx = context.Background()
		rsaKey = mustGenerateRSAKey()
		rsaPEM = mustMarshalRSAPrivateKeyPEM(rsaKey)
		ecdsaPEM = mustMarshalECDSAPrivateKeyPEM(mustGenerateECDSAKey())
		expIssuer = "dev.devops.pingcap.com"
		audience = "devops-internal"
		serviceConfig = config.InternalTokenServiceConfig{
			DefaultTTL:    pointer.Pointer(15 * time.Minute),
			SigningKeyPEM: rsaPEM,
			KeyID:         "devops-key-1",
		}
	})

	newService := func(conf config.InternalTokenServiceConfig) (itksvcapi.InternalTokenServiceAPI, error) {
		return NewService(ServiceParam{
			Conf: &config.RootConfig{
				Services: config.Services{
					InternalToken: conf,
				},
			},
		})
	}

	Describe("NewService", func() {
		It("initializes successfully with a valid rsa pem", func() {
			svc, err := newService(serviceConfig)
			Expect(err).NotTo(HaveOccurred())
			Expect(svc).NotTo(BeNil())
		})

		It("fails fast for an invalid pem", func() {
			serviceConfig.SigningKeyPEM = "not-a-pem"

			_, err := newService(serviceConfig)
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("parse jwt token signing key"))
		})

		It("fails fast for a non-rsa key", func() {
			serviceConfig.SigningKeyPEM = ecdsaPEM

			_, err := newService(serviceConfig)
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("parse jwt token signing key"))
		})
	})

	Describe("IssueJWTToken", func() {
		It("issues a token with standard and custom claims", func() {
			svc, err := newService(serviceConfig)
			Expect(err).NotTo(HaveOccurred())

			expiresAt := time.Now().Add(30 * time.Minute).Round(time.Second)
			jti := "provided-jti"
			issued, err := svc.IssueJWTToken(ctx, itksvcapi.IssueJWTTokenParam{
				Subject:   "user@example.com",
				Audience:  audience,
				ExpiresAt: &expiresAt,
				JWTID:     &jti,
				CustomClaims: map[string]any{
					"role": "admin",
				},
			})
			Expect(err).NotTo(HaveOccurred())
			Expect(issued.Issuer).To(Equal(expIssuer))
			Expect(issued.Audience).To(Equal(audience))
			Expect(issued.KeyID).To(Equal(serviceConfig.KeyID))
			Expect(issued.ExpireAt.Unix()).To(Equal(expiresAt.Unix()))

			token, err := jwt.Parse(issued.Token, func(token *jwt.Token) (any, error) {
				Expect(token.Method.Alg()).To(Equal(jwt.SigningMethodRS256.Alg()))
				Expect(token.Header["kid"]).To(Equal(serviceConfig.KeyID))
				return &rsaKey.PublicKey, nil
			})
			Expect(err).NotTo(HaveOccurred())
			Expect(token.Valid).To(BeTrue())

			claims, ok := token.Claims.(jwt.MapClaims)
			Expect(ok).To(BeTrue())
			Expect(claims["iss"]).To(Equal(expIssuer))
			Expect(claims["aud"]).To(Equal(audience))
			Expect(claims["sub"]).To(Equal("user@example.com"))
			Expect(claims["jti"]).To(Equal(jti))
			Expect(claims["role"]).To(Equal("admin"))
			Expect(claims["iat"]).NotTo(BeNil())
			Expect(claims["exp"]).To(BeEquivalentTo(expiresAt.Unix()))
		})

		It("uses the provided audience in the issued token", func() {
			svc, err := newService(serviceConfig)
			Expect(err).NotTo(HaveOccurred())

			overrideAudience := "custom-audience"
			issued, err := svc.IssueJWTToken(ctx, itksvcapi.IssueJWTTokenParam{
				Subject:  "user@example.com",
				Audience: overrideAudience,
			})
			Expect(err).NotTo(HaveOccurred())
			Expect(issued.Audience).To(Equal(overrideAudience))

			token, err := jwt.Parse(issued.Token, func(token *jwt.Token) (any, error) {
				return &rsaKey.PublicKey, nil
			})
			Expect(err).NotTo(HaveOccurred())

			claims, ok := token.Claims.(jwt.MapClaims)
			Expect(ok).To(BeTrue())
			Expect(claims["aud"]).To(Equal(overrideAudience))
		})

		It("rejects empty audience", func() {
			svc, err := newService(serviceConfig)
			Expect(err).NotTo(HaveOccurred())

			_, err = svc.IssueJWTToken(ctx, itksvcapi.IssueJWTTokenParam{
				Subject:  "user@example.com",
				Audience: "",
			})
			Expect(err).To(MatchError("jwt token audience is required"))
		})

		It("applies the default ttl when expiration is not provided", func() {
			svc, err := newService(serviceConfig)
			Expect(err).NotTo(HaveOccurred())

			before := time.Now()
			issued, err := svc.IssueJWTToken(ctx, itksvcapi.IssueJWTTokenParam{
				Subject:  "user@example.com",
				Audience: audience,
			})
			after := time.Now()

			Expect(err).NotTo(HaveOccurred())
			Expect(issued.ExpireAt).To(BeTemporally(">=", before.Add(*serviceConfig.DefaultTTL)))
			Expect(issued.ExpireAt).To(BeTemporally("<=", after.Add(*serviceConfig.DefaultTTL).Add(time.Second)))
		})

		It("generates a jti when not provided", func() {
			svc, err := newService(serviceConfig)
			Expect(err).NotTo(HaveOccurred())

			issued, err := svc.IssueJWTToken(ctx, itksvcapi.IssueJWTTokenParam{
				Subject:  "user@example.com",
				Audience: audience,
			})
			Expect(err).NotTo(HaveOccurred())

			token, err := jwt.Parse(issued.Token, func(token *jwt.Token) (any, error) {
				return &rsaKey.PublicKey, nil
			})
			Expect(err).NotTo(HaveOccurred())

			claims, ok := token.Claims.(jwt.MapClaims)
			Expect(ok).To(BeTrue())
			jti, ok := claims["jti"].(string)
			Expect(ok).To(BeTrue())
			Expect(jti).NotTo(BeEmpty())
		})

		It("rejects custom claims that overwrite reserved claims", func() {
			svc, err := newService(serviceConfig)
			Expect(err).NotTo(HaveOccurred())

			_, err = svc.IssueJWTToken(ctx, itksvcapi.IssueJWTTokenParam{
				Subject:  "user@example.com",
				Audience: audience,
				CustomClaims: map[string]any{
					"iss": "override",
				},
			})
			Expect(err).To(MatchError(`custom claim "iss" is reserved`))
		})
	})

	Describe("GetJWKS", func() {
		It("exports a public jwks that can verify issued tokens", func() {
			svc, err := newService(serviceConfig)
			Expect(err).NotTo(HaveOccurred())

			jwks, err := svc.GetJWKS(ctx)
			Expect(err).NotTo(HaveOccurred())
			Expect(jwks.Keys).To(HaveLen(1))
			Expect(jwks.Keys[0].Kty).To(Equal("RSA"))
			Expect(jwks.Keys[0].Use).To(Equal("sig"))
			Expect(jwks.Keys[0].Alg).To(Equal(jwt.SigningMethodRS256.Alg()))
			Expect(jwks.Keys[0].Kid).To(Equal(serviceConfig.KeyID))

			publicKey, err := jwkToRSAPublicKey(jwks.Keys[0])
			Expect(err).NotTo(HaveOccurred())
			Expect(publicKey.N.Cmp(rsaKey.PublicKey.N)).To(Equal(0))
			Expect(publicKey.E).To(Equal(rsaKey.PublicKey.E))

			issued, err := svc.IssueJWTToken(ctx, itksvcapi.IssueJWTTokenParam{
				Subject:  "user@example.com",
				Audience: audience,
			})
			Expect(err).NotTo(HaveOccurred())

			token, err := jwt.Parse(issued.Token, func(token *jwt.Token) (any, error) {
				return publicKey, nil
			})
			Expect(err).NotTo(HaveOccurred())
			Expect(token.Valid).To(BeTrue())
		})
	})
})

func mustGenerateRSAKey() *rsa.PrivateKey {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		panic(fmt.Sprintf("generate rsa key: %v", err))
	}
	return key
}

func mustGenerateECDSAKey() *ecdsa.PrivateKey {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		panic(fmt.Sprintf("generate ecdsa key: %v", err))
	}
	return key
}

func mustMarshalRSAPrivateKeyPEM(key *rsa.PrivateKey) string {
	return string(pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	}))
}

func mustMarshalECDSAPrivateKeyPEM(key *ecdsa.PrivateKey) string {
	bytes, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		panic(fmt.Sprintf("marshal ecdsa key: %v", err))
	}
	return string(pem.EncodeToMemory(&pem.Block{
		Type:  "EC PRIVATE KEY",
		Bytes: bytes,
	}))
}

func jwkToRSAPublicKey(jwk itkent.JWK) (*rsa.PublicKey, error) {
	nBytes, err := base64.RawURLEncoding.DecodeString(jwk.N)
	if err != nil {
		return nil, fmt.Errorf("decode modulus: %w", err)
	}
	eBytes, err := base64.RawURLEncoding.DecodeString(jwk.E)
	if err != nil {
		return nil, fmt.Errorf("decode exponent: %w", err)
	}

	return &rsa.PublicKey{
		N: new(big.Int).SetBytes(nBytes),
		E: int(new(big.Int).SetBytes(eBytes).Int64()),
	}, nil
}