package connector_oauth

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"strings"
)

const cipherEnvelopeVersion = "v1"

type AESGCMCipher struct {
	activeKeyID string
	keys        map[string][]byte
	random      io.Reader
}

func NewAESGCMCipher(activeKeyID string, keys map[string][]byte) (*AESGCMCipher, error) {
	activeKeyID = strings.TrimSpace(activeKeyID)
	if !validCipherKeyID(activeKeyID) || len(keys) == 0 {
		return nil, ErrEncryptionKeyMissing
	}
	cloned := make(map[string][]byte, len(keys))
	for keyID, key := range keys {
		if !validCipherKeyID(keyID) || (len(key) != 16 && len(key) != 24 && len(key) != 32) {
			return nil, ErrEncryptionKeyMissing
		}
		cloned[keyID] = append([]byte(nil), key...)
	}
	if _, ok := cloned[activeKeyID]; !ok {
		return nil, ErrEncryptionKeyMissing
	}
	return &AESGCMCipher{activeKeyID: activeKeyID, keys: cloned, random: rand.Reader}, nil
}

func (c *AESGCMCipher) Seal(plaintext []byte, secretContext SecretContext) (string, string, error) {
	if c == nil {
		return "", "", ErrEncryptionKeyMissing
	}
	key, ok := c.keys[c.activeKeyID]
	if !ok {
		return "", "", ErrEncryptionKeyMissing
	}
	gcm, err := newGCM(key)
	if err != nil {
		return "", "", err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(c.random, nonce); err != nil {
		return "", "", fmt.Errorf("generate secret nonce: %w", err)
	}
	aad, err := secretAAD(c.activeKeyID, secretContext)
	if err != nil {
		return "", "", err
	}
	sealed := gcm.Seal(nil, nonce, plaintext, aad)
	encode := base64.RawURLEncoding.EncodeToString
	return strings.Join([]string{cipherEnvelopeVersion, c.activeKeyID, encode(nonce), encode(sealed)}, "."), c.activeKeyID, nil
}

func (c *AESGCMCipher) Open(ciphertext string, secretContext SecretContext) ([]byte, error) {
	if c == nil {
		return nil, ErrEncryptionKeyMissing
	}
	parts := strings.Split(ciphertext, ".")
	if len(parts) != 4 || parts[0] != cipherEnvelopeVersion || strings.TrimSpace(parts[1]) == "" {
		return nil, ErrCiphertextInvalid
	}
	keyID := parts[1]
	if !validCipherKeyID(keyID) {
		return nil, ErrCiphertextInvalid
	}
	key, ok := c.keys[keyID]
	if !ok {
		return nil, ErrEncryptionKeyMissing
	}
	nonce, err := base64.RawURLEncoding.DecodeString(parts[2])
	if err != nil {
		return nil, ErrCiphertextInvalid
	}
	sealed, err := base64.RawURLEncoding.DecodeString(parts[3])
	if err != nil {
		return nil, ErrCiphertextInvalid
	}
	gcm, err := newGCM(key)
	if err != nil || len(nonce) != gcm.NonceSize() {
		return nil, ErrCiphertextInvalid
	}
	aad, err := secretAAD(keyID, secretContext)
	if err != nil {
		return nil, err
	}
	plaintext, err := gcm.Open(nil, nonce, sealed, aad)
	if err != nil {
		return nil, ErrCiphertextInvalid
	}
	return plaintext, nil
}

func validCipherKeyID(keyID string) bool {
	return keyID != "" && keyID == strings.TrimSpace(keyID) && !strings.Contains(keyID, ".")
}

func newGCM(key []byte) (cipher.AEAD, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, ErrEncryptionKeyMissing
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, ErrEncryptionKeyMissing
	}
	return gcm, nil
}

func secretAAD(keyID string, secretContext SecretContext) ([]byte, error) {
	secretContext.UserID = strings.TrimSpace(secretContext.UserID)
	secretContext.ConnectorKey = strings.TrimSpace(secretContext.ConnectorKey)
	secretContext.RecordID = strings.TrimSpace(secretContext.RecordID)
	if strings.TrimSpace(keyID) == "" || secretContext.UserID == "" || secretContext.ConnectorKey == "" || secretContext.RecordID == "" || secretContext.Generation < 1 {
		return nil, ErrInvalidInput
	}
	switch secretContext.Purpose {
	case SecretAuthorizationVerifier, SecretPrincipalAccessToken, SecretPrincipalRefreshToken,
		SecretAuthorizationResult, SecretRefreshResult:
	default:
		return nil, ErrInvalidInput
	}
	return json.Marshal(struct {
		EnvelopeVersion string        `json:"version"`
		KeyID           string        `json:"key_id"`
		Purpose         SecretPurpose `json:"purpose"`
		UserID          string        `json:"user_id"`
		ConnectorKey    string        `json:"connector_key"`
		RecordID        string        `json:"record_id"`
		Generation      int64         `json:"generation"`
		Version         int64         `json:"version_number,omitempty"`
	}{cipherEnvelopeVersion, keyID, secretContext.Purpose, secretContext.UserID, secretContext.ConnectorKey, secretContext.RecordID, secretContext.Generation, secretContext.Version})
}