package edgeencrypt

import (
	"crypto/rsa"
	"fmt"
	"slices"
	"time"

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

type Token struct {
	BearerToken string `json:"token"`   // actual bearer token
	Version     string `json:"version"` // 1,2
}

const (
	Week            = 8760 * time.Hour
	Year            = 52 * Week
	DefaultDuration = Year
)

type Role string

const (
	Encryption = Role("encryption")
	Decryption = Role("decryption")
)

var AllRoles = []Role{Encryption, Decryption}

func (r Role) Valid() bool {
	if len(r) == 0 {
		return false
	}
	return slices.Contains(AllRoles, r)
}

type EncryptionClaims struct {
	jwt.StandardClaims
	BannerEdgeID string `json:"bannerEdgeID,omitempty"`
	ChannelID    string `json:"channelId"`
	Channel      string `json:"channel"`
	Role         Role   `json:"role"`
}

func (c *EncryptionClaims) Valid() error {
	if len(c.ChannelID) == 0 {
		return fmt.Errorf("channelID is required")
	}
	if len(c.Channel) == 0 {
		return fmt.Errorf("channel is required")
	}
	if !c.Role.Valid() {
		return fmt.Errorf("invalid role: %v", c.Role)
	}
	if c.Role == Encryption && len(c.BannerEdgeID) == 0 {
		return fmt.Errorf("invalid bannerEdgeID is required for encryption role")
	}
	return c.StandardClaims.Valid()
}

func (c *EncryptionClaims) ValidChannel(channelID string) bool {
	return c.ChannelID == channelID
}

func (c *EncryptionClaims) ValidChannelName(channelName string) bool {
	return c.Channel == channelName
}

func (c *EncryptionClaims) HasRole(role Role) bool {
	return c.Role == role
}

// CreateToken use a private key to generate a bearer token to authenticate with edge encryption/decryption services
func CreateToken(method jwt.SigningMethod, key interface{}, duration time.Duration, channelID, channelName string, role Role, banner ...string) (string, error) {
	if method == nil {
		return "", fmt.Errorf("signing method is required to create token")
	}
	if !validSigningMethod(method) {
		return "", fmt.Errorf("unsupported signing method: %T, must be RSA256", method)
	}
	if key == nil {
		return "", fmt.Errorf("signing key is required to create token")
	}
	if duration == 0 {
		duration = DefaultDuration
	}
	if len(channelID) == 0 {
		return "", fmt.Errorf("channelID is required")
	}
	if len(channelName) == 0 {
		return "", fmt.Errorf("channelName is required")
	}
	if !role.Valid() {
		return "", fmt.Errorf("invalid role: %v", role)
	}
	var bannerID string
	if role == Encryption {
		if len(banner) == 0 {
			return "", fmt.Errorf("bannerEdgeID is required for encryption role")
		}
		bannerID = banner[0]
	}
	c := &EncryptionClaims{
		BannerEdgeID: bannerID,
		ChannelID:    channelID,
		Channel:      channelName,
		Role:         role,
		StandardClaims: jwt.StandardClaims{
			ExpiresAt: time.Now().Add(duration).Unix(),
		},
	}

	claims := jwt.NewWithClaims(method, c)
	token, err := claims.SignedString(key)
	if err != nil {
		return "", fmt.Errorf("failed to sign token: %v", err)
	}
	return token, nil
}

// FromToken rsa public key, bearerToken as input, encryption claims as output
func FromToken(publicKey *rsa.PublicKey, bearerToken string) (*EncryptionClaims, error) {
	c := &EncryptionClaims{}
	_, err := jwt.ParseWithClaims(bearerToken, c,
		func(token *jwt.Token) (interface{}, error) {
			if !validSigningMethod(token.Method) {
				return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
			}
			return publicKey, nil
		})
	if err != nil {
		return nil, fmt.Errorf("failed to parse token: %v", err)
	}
	return c, nil
}

func validSigningMethod(method jwt.SigningMethod) bool {
	if method == nil {
		return false
	}
	_, isRSA := method.(*jwt.SigningMethodRSA)
	_, isKMS := method.(*SigningMethodKMS)
	return isRSA || isKMS
}