package authtoken

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"maps"
	"strings"
	"time"
)

const (
	SessionLifetime = 24 * time.Hour

	algorithmHS256 = "HS256"
	tokenType      = "JWT"
	maxTokenBytes  = 4096
	segmentCount   = 3
	segmentSep     = "."
)

var (
	ErrNoToken     = errors.New("no token found")
	ErrInvalid     = errors.New("token is invalid")
	ErrExpired     = errors.New("token is expired")
	ErrNotYetValid = errors.New("token is not yet valid")
	ErrRevoked     = errors.New("token is revoked")
)

type RevocationStore interface {
	Revoke(id string, expiresAt time.Time) error
	IsRevoked(id string) (bool, error)
}

type Claims map[string]any

type header struct {
	Algorithm string `json:"alg"`
	Type      string `json:"typ,omitempty"`
}

type Auth struct {
	secret      []byte
	clock       func() time.Time
	revocations RevocationStore
}

type Option func(*Auth)

func WithClock(clock func() time.Time) Option {
	return func(a *Auth) { a.clock = clock }
}

func WithRevocationStore(store RevocationStore) Option {
	return func(a *Auth) { a.revocations = store }
}

func New(secret []byte, options ...Option) *Auth {
	auth := &Auth{secret: secret, clock: time.Now}
	for _, option := range options {
		option(auth)
	}
	return auth
}

func (a *Auth) Encode(claims map[string]any) (string, error) {
	payload := maps.Clone(claims)
	if payload == nil {
		payload = map[string]any{}
	}
	now := a.clock()
	if _, ok := payload["iat"]; !ok {
		payload["iat"] = now.Unix()
	}
	if _, ok := payload["exp"]; !ok {
		payload["exp"] = now.Add(SessionLifetime).Unix()
	}
	if _, ok := payload["jti"]; !ok {
		payload["jti"] = rand.Text()
	}

	headerSegment, err := encodeSegment(header{Algorithm: algorithmHS256, Type: tokenType})
	if err != nil {
		return "", err
	}
	payloadSegment, err := encodeSegment(payload)
	if err != nil {
		return "", err
	}

	signingInput := headerSegment + segmentSep + payloadSegment
	return signingInput + segmentSep + a.sign(signingInput), nil
}

func (a *Auth) Decode(token string) (Claims, error) {
	if token == "" {
		return nil, ErrNoToken
	}
	if len(token) > maxTokenBytes {
		return nil, ErrInvalid
	}

	segments := strings.Split(token, segmentSep)
	if len(segments) != segmentCount {
		return nil, ErrInvalid
	}

	var tokenHeader header
	if err := decodeSegment(segments[0], &tokenHeader); err != nil || tokenHeader.Algorithm != algorithmHS256 {
		return nil, ErrInvalid
	}

	signingInput := segments[0] + segmentSep + segments[1]
	if !hmac.Equal([]byte(a.sign(signingInput)), []byte(segments[2])) {
		return nil, ErrInvalid
	}

	var claims Claims
	if err := decodeSegment(segments[1], &claims); err != nil || claims == nil {
		return nil, ErrInvalid
	}

	if err := a.validateTimes(claims); err != nil {
		return nil, err
	}
	if err := a.checkRevocation(claims); err != nil {
		return nil, err
	}
	return claims, nil
}

func (a *Auth) Revoke(claims Claims) error {
	id, ok := claims["jti"].(string)
	if !ok || a.revocations == nil {
		return nil
	}
	expiry, present, err := numericDate(claims, "exp")
	if err != nil || !present {
		return ErrInvalid
	}
	return a.revocations.Revoke(id, expiry)
}

func (a *Auth) checkRevocation(claims Claims) error {
	raw, present := claims["jti"]
	if !present {
		return nil
	}
	id, ok := raw.(string)
	if !ok {
		return ErrInvalid
	}
	if a.revocations == nil {
		return nil
	}
	revoked, err := a.revocations.IsRevoked(id)
	if err != nil {
		return errors.Join(ErrInvalid, err)
	}
	if revoked {
		return ErrRevoked
	}
	return nil
}

func (a *Auth) validateTimes(claims Claims) error {
	now := a.clock()

	expiry, present, err := numericDate(claims, "exp")
	if err != nil {
		return err
	}
	if !present {
		return ErrInvalid
	}
	if !now.Before(expiry) {
		return ErrExpired
	}

	notBefore, present, err := numericDate(claims, "nbf")
	if err != nil {
		return err
	}
	if present && now.Before(notBefore) {
		return ErrNotYetValid
	}
	return nil
}

func numericDate(claims Claims, name string) (time.Time, bool, error) {
	raw, present := claims[name]
	if !present {
		return time.Time{}, false, nil
	}
	seconds, ok := raw.(float64)
	if !ok {
		return time.Time{}, false, ErrInvalid
	}
	return time.Unix(int64(seconds), 0), true, nil
}

func (a *Auth) sign(signingInput string) string {
	mac := hmac.New(sha256.New, a.secret)
	mac.Write([]byte(signingInput))
	return base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
}

func encodeSegment(value any) (string, error) {
	encoded, err := json.Marshal(value)
	if err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(encoded), nil
}

func decodeSegment(segment string, into any) error {
	decoded, err := base64.RawURLEncoding.DecodeString(segment)
	if err != nil {
		return err
	}
	return json.Unmarshal(decoded, into)
}