package signing

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"
)

var rawURLEncoding = base64.RawURLEncoding

type RS256Signer interface {
	ActiveKey(ctx context.Context) (SigningKey, error)
	SignRS256(ctx context.Context, key SigningKey, signingInput []byte) ([]byte, error)
}

type SigningKey struct {
	ID      string
	Version int
}

type Claims struct {
	Issuer        string   `json:"iss"`
	Subject       string   `json:"sub"`
	PrincipalName string   `json:"principal_name,omitempty"`
	Audience      []string `json:"aud"`
	IssuedAt      int64    `json:"iat"`
	NotBefore     int64    `json:"nbf"`
	ExpiresAt     int64    `json:"exp"`
	JWTID         string   `json:"jti"`
	Scope         string   `json:"scope"`
}

type Issuer struct {
	signer RS256Signer
	now    func() time.Time
}

func NewIssuer(signer RS256Signer) (*Issuer, error) {
	if signer == nil {
		return nil, errors.New("signer is required")
	}
	return &Issuer{signer: signer, now: time.Now}, nil
}

func (i *Issuer) Sign(ctx context.Context, claims Claims) (string, error) {
	if err := validateClaims(claims, i.now()); err != nil {
		return "", err
	}

	key, err := i.signer.ActiveKey(ctx)
	if err != nil {
		return "", fmt.Errorf("select signing key: %w", err)
	}
	if key.ID == "" || key.Version < 1 {
		return "", errors.New("signer returned an invalid signing key")
	}

	header, err := json.Marshal(map[string]string{
		"alg": "RS256",
		"kid": key.ID,
		"typ": "at+jwt",
	})
	if err != nil {
		return "", fmt.Errorf("encode protected header: %w", err)
	}
	payload, err := json.Marshal(claims)
	if err != nil {
		return "", fmt.Errorf("encode claims: %w", err)
	}

	signingInput := rawURLEncoding.EncodeToString(header) + "." + rawURLEncoding.EncodeToString(payload)
	signature, err := i.signer.SignRS256(ctx, key, []byte(signingInput))
	if err != nil {
		return "", fmt.Errorf("sign JWT: %w", err)
	}
	return signingInput + "." + rawURLEncoding.EncodeToString(signature), nil
}

func validateClaims(claims Claims, now time.Time) error {
	if claims.Issuer == "" || claims.Subject == "" || len(claims.Audience) != 1 || claims.JWTID == "" || claims.Scope == "" {
		return errors.New("required JWT claim is missing")
	}
	if strings.ContainsAny(claims.Subject, "\r\n") || strings.ContainsAny(claims.JWTID, "\r\n") {
		return errors.New("invalid JWT claim")
	}
	if claims.ExpiresAt <= claims.NotBefore || claims.NotBefore < claims.IssuedAt {
		return errors.New("invalid JWT time window")
	}
	if time.Unix(claims.ExpiresAt, 0).After(now.Add(5 * time.Minute)) {
		return errors.New("JWT lifetime exceeds five minutes")
	}
	return nil
}