package token

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hkdf"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"regexp"
	"strings"
	"time"
)

const (
	tokenPrefix       = "mailtok.v1"
	tokenSchema       = 1
	maxTokenBytes     = 2048
	maxTokenTTL       = 10 * time.Minute
	maxCursorBytes    = 512
	maxReferenceBytes = 512
)

var keyIDPattern = regexp.MustCompile(`^[a-z0-9_-]{1,32}$`)

type Key struct {
	ID       string
	Material []byte
}

type Keyring struct {
	Current Key
	Legacy  []Key
}

type Codec struct {
	keys       Keyring
	deployment string
	now        func() time.Time
}

func New(
	keys Keyring,
	deployment string,
	now func() time.Time,
) (*Codec, error) {
	if err := validateKey(keys.Current); err != nil {
		return nil, err
	}
	seen := map[string]struct{}{keys.Current.ID: {}}
	for _, key := range keys.Legacy {
		if err := validateKey(key); err != nil {
			return nil, err
		}
		if _, duplicate := seen[key.ID]; duplicate {
			return nil, fmt.Errorf("Mail token key ID is duplicated")
		}
		seen[key.ID] = struct{}{}
	}
	deployment = strings.TrimSpace(deployment)
	if deployment == "" || len(deployment) > 128 {
		return nil, fmt.Errorf("Mail token deployment audience is invalid")
	}
	if now == nil {
		now = time.Now
	}
	return &Codec{
		keys: Keyring{
			Current: cloneKey(keys.Current),
			Legacy:  cloneKeys(keys.Legacy),
		},
		deployment: deployment,
		now:        now,
	}, nil
}

func (codec *Codec) Seal(
	claims Claims,
	ttl time.Duration,
) (string, error) {
	if codec == nil {
		return "", fmt.Errorf("Mail token codec is not configured")
	}
	if ttl <= 0 || ttl > maxTokenTTL {
		return "", fmt.Errorf("Mail token lifetime is invalid")
	}
	now := codec.now().UTC()
	claims.SchemaVersion = tokenSchema
	claims.IssuedAtUnix = now.Unix()
	claims.ExpiresAtUnix = now.Add(ttl).Unix()
	if err := validateClaims(claims); err != nil {
		return "", err
	}
	plaintext, err := json.Marshal(claims)
	if err != nil {
		return "", fmt.Errorf("encode Mail token claims: %w", err)
	}
	aead, err := tokenAEAD(codec.keys.Current, codec.deployment)
	if err != nil {
		return "", err
	}
	nonce := make([]byte, aead.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", fmt.Errorf("generate Mail token nonce: %w", err)
	}
	aad := codec.aad(
		codec.keys.Current.ID,
		claims.Action,
		claims.Kind,
	)
	sealed := aead.Seal(nonce, nonce, plaintext, aad)
	encoded := strings.Join([]string{
		tokenPrefix,
		codec.keys.Current.ID,
		base64.RawURLEncoding.EncodeToString(sealed),
	}, ".")
	if len(encoded) > maxTokenBytes {
		return "", fmt.Errorf("Mail token exceeded size limit")
	}
	return encoded, nil
}

func (codec *Codec) Open(
	encoded string,
	binding Binding,
) (Claims, error) {
	if codec == nil || len(encoded) == 0 ||
		len(encoded) > maxTokenBytes {
		return Claims{}, errInvalidToken
	}
	parts := strings.Split(encoded, ".")
	if len(parts) != 4 ||
		parts[0]+"."+parts[1] != tokenPrefix {
		return Claims{}, errInvalidToken
	}
	key, ok := codec.key(parts[2])
	if !ok {
		return Claims{}, errInvalidToken
	}
	sealed, err := base64.RawURLEncoding.DecodeString(parts[3])
	if err != nil {
		return Claims{}, errInvalidToken
	}
	aead, err := tokenAEAD(key, codec.deployment)
	if err != nil || len(sealed) < aead.NonceSize() {
		return Claims{}, errInvalidToken
	}
	nonce := sealed[:aead.NonceSize()]
	plaintext, err := aead.Open(
		nil,
		nonce,
		sealed[aead.NonceSize():],
		codec.aad(key.ID, binding.Action, binding.Kind),
	)
	if err != nil {
		return Claims{}, errInvalidToken
	}
	var claims Claims
	decoder := json.NewDecoder(bytes.NewReader(plaintext))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&claims); err != nil {
		return Claims{}, errInvalidToken
	}
	if err := decoder.Decode(&struct{}{}); err != io.EOF {
		return Claims{}, errInvalidToken
	}
	if err := validateClaims(claims); err != nil ||
		!sameBinding(claims, binding) {
		return Claims{}, errInvalidToken
	}
	now := codec.now().UTC().Unix()
	if claims.IssuedAtUnix > now+30 ||
		claims.ExpiresAtUnix <= now ||
		claims.ExpiresAtUnix-claims.IssuedAtUnix >
			int64(maxTokenTTL/time.Second) {
		return Claims{}, errInvalidToken
	}
	return claims, nil
}

func validateClaims(claims Claims) error {
	if claims.SchemaVersion != tokenSchema ||
		!validKind(claims.Kind) ||
		!validIdentity(claims.WorkspaceID) ||
		!validIdentity(claims.InstallationID) ||
		!validIdentity(claims.ConnectionID) ||
		!validIdentity(claims.Generation) ||
		!validIdentity(claims.AgentID) ||
		!validIdentity(claims.Action) ||
		claims.IssuedAtUnix <= 0 ||
		claims.ExpiresAtUnix <= claims.IssuedAtUnix {
		return errInvalidToken
	}
	if len(claims.Mailbox) > maxReferenceBytes ||
		len(claims.Part) > maxReferenceBytes ||
		len(claims.Cursor) > maxCursorBytes ||
		len(claims.FilterDigest) > 80 {
		return errInvalidToken
	}
	switch claims.Kind {
	case KindCursor:
		if claims.Cursor == "" ||
			claims.FilterDigest == "" ||
			claims.Limit <= 0 ||
			claims.Limit > 50 ||
			claims.UID != 0 ||
			claims.UIDValidity != 0 {
			return errInvalidToken
		}
	default:
		if claims.Mailbox == "" ||
			claims.UID == 0 ||
			claims.UIDValidity == 0 ||
			claims.Cursor != "" ||
			claims.FilterDigest != "" ||
			claims.Limit != 0 {
			return errInvalidToken
		}
	}
	return nil
}

func sameBinding(claims Claims, binding Binding) bool {
	return claims.Kind == binding.Kind &&
		claims.WorkspaceID == binding.WorkspaceID &&
		claims.InstallationID == binding.InstallationID &&
		claims.ConnectionID == binding.ConnectionID &&
		claims.Generation == binding.Generation &&
		claims.AgentID == binding.AgentID &&
		claims.Action == binding.Action
}

func (codec *Codec) aad(
	keyID string,
	action string,
	kind string,
) []byte {
	return []byte(strings.Join([]string{
		tokenPrefix,
		codec.deployment,
		"mail",
		keyID,
		action,
		kind,
	}, "\x00"))
}

func (codec *Codec) key(id string) (Key, bool) {
	if codec.keys.Current.ID == id {
		return codec.keys.Current, true
	}
	for _, key := range codec.keys.Legacy {
		if key.ID == id {
			return key, true
		}
	}
	return Key{}, false
}

func tokenAEAD(key Key, deployment string) (cipher.AEAD, error) {
	derived, err := hkdf.Key(
		sha256.New,
		key.Material,
		[]byte("agentteam/email/token/salt/v1"),
		"agentteam/email/token/aes-gcm/v1\x00"+deployment+"\x00"+key.ID,
		32,
	)
	if err != nil {
		return nil, fmt.Errorf("derive Mail token key: %w", err)
	}
	block, err := aes.NewCipher(derived)
	if err != nil {
		return nil, fmt.Errorf("create Mail token cipher: %w", err)
	}
	return cipher.NewGCM(block)
}

func validateKey(key Key) error {
	if !keyIDPattern.MatchString(key.ID) ||
		len(key.Material) != 32 {
		return fmt.Errorf("Mail token key is invalid")
	}
	return nil
}

func cloneKey(key Key) Key {
	return Key{
		ID:       key.ID,
		Material: append([]byte(nil), key.Material...),
	}
}

func cloneKeys(keys []Key) []Key {
	result := make([]Key, 0, len(keys))
	for _, key := range keys {
		result = append(result, cloneKey(key))
	}
	return result
}

func validIdentity(value string) bool {
	return strings.TrimSpace(value) == value &&
		value != "" &&
		len(value) <= 255
}

func validKind(value string) bool {
	switch value {
	case KindMessage, KindThread, KindDraft, KindAttachment, KindCursor:
		return true
	default:
		return false
	}
}

var errInvalidToken = fmt.Errorf("Mail token is invalid or expired")