package signing

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

	"github.com/go-jose/go-jose/v4"
)

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")
	}

	payload, err := json.Marshal(claims)
	if err != nil {
		return "", fmt.Errorf("encode claims: %w", err)
	}

	opaque := &contextSigner{ctx: ctx, signer: i.signer, key: key}
	options := (&jose.SignerOptions{}).
		WithType(jose.ContentType("at+jwt")).
		WithHeader(jose.HeaderKey("kid"), key.ID)
	signer, err := jose.NewSigner(jose.SigningKey{Algorithm: jose.RS256, Key: opaque}, options)
	if err != nil {
		return "", fmt.Errorf("create JWT signer: %w", err)
	}
	jws, err := signer.Sign(payload)
	if err != nil {
		return "", fmt.Errorf("sign JWT: %w", err)
	}
	compact, err := jws.CompactSerialize()
	if err != nil {
		return "", fmt.Errorf("serialize JWT: %w", err)
	}
	return compact, nil
}

type contextSigner struct {
	ctx    context.Context
	signer RS256Signer
	key    SigningKey
}

func (*contextSigner) Public() *jose.JSONWebKey {
	return nil
}

func (*contextSigner) Algs() []jose.SignatureAlgorithm {
	return []jose.SignatureAlgorithm{jose.RS256}
}

func (s *contextSigner) SignPayload(payload []byte, algorithm jose.SignatureAlgorithm) ([]byte, error) {
	if algorithm != jose.RS256 {
		return nil, errors.New("unsupported signing algorithm")
	}
	return s.signer.SignRS256(s.ctx, s.key, payload)
}

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
}