package signing

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

const maximumTTL = 5 * time.Minute

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

type SigningKey struct {
	ID      string
	Version int
}

type IssueRequest struct {
	Subject       string
	PrincipalName string
	Audience      string
	JWTID         string
	Scope         string
}

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 {
	issuer string
	ttl    time.Duration
	signer RS256Signer
}

func NewIssuer(issuer string, ttl time.Duration, signer RS256Signer) (*Issuer, error) {
	if issuer == "" || signer == nil {
		return nil, errors.New("issuer and signer are required")
	}
	if ttl <= 0 || ttl > maximumTTL {
		return nil, errors.New("TTL must be between zero and five minutes")
	}
	return &Issuer{issuer: issuer, ttl: ttl, signer: signer}, nil
}

func (i *Issuer) Issue(ctx context.Context, request IssueRequest) (string, error) {
	return i.issueAt(ctx, request, time.Now())
}

func (i *Issuer) issueAt(ctx context.Context, request IssueRequest, now time.Time) (string, error) {
	if err := validateIssueRequest(request); 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")
	}

	claims := Claims{
		Issuer:        i.issuer,
		Subject:       request.Subject,
		PrincipalName: request.PrincipalName,
		Audience:      []string{request.Audience},
		IssuedAt:      now.Unix(),
		NotBefore:     now.Unix(),
		ExpiresAt:     now.Add(i.ttl).Unix(),
		JWTID:         request.JWTID,
		Scope:         request.Scope,
	}
	signingInput, err := encodeSigningInput(key.ID, claims)
	if err != nil {
		return "", err
	}
	signature, err := i.signer.SignRS256(ctx, key, []byte(signingInput))
	if err != nil {
		return "", fmt.Errorf("sign JWT: %w", err)
	}
	return serializeCompact(signingInput, signature), nil
}

func validateIssueRequest(request IssueRequest) error {
	if request.Subject == "" || request.Audience == "" || request.JWTID == "" || request.Scope == "" {
		return errors.New("required JWT input is missing")
	}
	if strings.ContainsAny(request.Subject, "\r\n") || strings.ContainsAny(request.JWTID, "\r\n") {
		return errors.New("invalid JWT input")
	}
	return nil
}

func encodeSigningInput(keyID string, claims Claims) (string, error) {
	header, err := json.Marshal(map[string]string{"alg": "RS256", "kid": keyID, "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)
	}
	return base64.RawURLEncoding.EncodeToString(header) + "." + base64.RawURLEncoding.EncodeToString(payload), nil
}

func serializeCompact(signingInput string, signature []byte) string {
	return signingInput + "." + base64.RawURLEncoding.EncodeToString(signature)
}