package lagoon

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/hkdf"
	"crypto/rand"
	"crypto/sha256"
	"database/sql/driver"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"strings"
	"sync"

	"git.golem15.com/golem15/summercms/modules/backpack"
	"git.golem15.com/golem15/summercms/modules/compass"
)

const (
	encryptedFormatV1  = byte(0x10)
	encryptedNonceSize = 12
	encryptedKeySize   = 32
	columnKeyInfo      = "summercms.lagoon.encrypted.v1"
	redactedLiteral    = "[redacted]"
	appKeyErr          = "lagoon: app.key is empty or invalid (set SUMMER_APP__KEY to a 32-byte base64 value)"
)

// Encrypted is an AES-256-GCM at-rest cast. Plaintext is unexported; the only
// greppable accessor is Reveal. MarshalJSON, String, and GoString always redact.
type Encrypted struct {
	plaintext []byte
	set       bool
}

// NewEncrypted holds plaintext for a subsequent Value() write.
func NewEncrypted(plaintext string) Encrypted {
	return Encrypted{plaintext: []byte(plaintext), set: true}
}

// Reveal returns the plaintext, or "" if the value is unset.
func (e Encrypted) Reveal() string {
	if !e.set {
		return ""
	}
	return string(e.plaintext)
}

// MarshalJSON always emits a redaction, never plaintext.
func (e Encrypted) MarshalJSON() ([]byte, error) {
	return json.Marshal(redactedLiteral)
}

// String returns a fixed redaction literal.
func (e Encrypted) String() string { return redactedLiteral }

// GoString returns a fixed redaction so %#v cannot leak plaintext.
func (e Encrypted) GoString() string { return "lagoon.Encrypted{[redacted]}" }

// Scan decrypts versioned ciphertext using the published primary key, then
// each previous key. src may be string, []byte, or nil (SQL NULL).
func (e *Encrypted) Scan(src any) error {
	if e == nil {
		return fmt.Errorf("lagoon: encrypted scan on nil receiver")
	}
	if src == nil {
		e.plaintext = nil
		e.set = false
		return nil
	}
	var raw string
	switch v := src.(type) {
	case string:
		raw = v
	case []byte:
		raw = string(v)
	default:
		return fmt.Errorf("lagoon: encrypted scan unsupported type %T", src)
	}
	raw = strings.TrimSpace(raw)
	if raw == "" {
		e.plaintext = nil
		e.set = false
		return nil
	}
	plain, err := decryptCiphertext(raw)
	if err != nil {
		return err
	}
	e.plaintext = plain
	e.set = true
	return nil
}

// Value re-encrypts the current plaintext under the published primary key.
// Unset values store SQL NULL.
func (e Encrypted) Value() (driver.Value, error) {
	if !e.set {
		return nil, nil
	}
	return encryptPlaintext(e.plaintext)
}

type encryptionKeys struct {
	primary   []byte
	previous  [][]byte
	primaryID byte
}

var (
	keysMu      sync.RWMutex
	currentKeys *encryptionKeys
)

// PublishEncryptionKeys installs column-encryption keys for Encrypted Scan/Value.
// OpenFromApp calls this once per boot so row-level encrypt/decrypt does not
// re-read config. Pass a nil app from tests that only need the package-level
// key set. The backpack publish is best-effort once-per-boot; package-level
// keys are the live source of truth and may be rotated in tests.
func PublishEncryptionKeys(app *backpack.App, primaryKey []byte, previousKeys [][]byte) error {
	k, err := newEncryptionKeys(primaryKey, previousKeys)
	if err != nil {
		return err
	}
	setCurrentKeys(k)
	if app == nil {
		return nil
	}
	if _, ok := app.Lookup[*encryptionKeys](); ok {
		return nil
	}
	return app.Publish(k)
}

func newEncryptionKeys(primaryKey []byte, previousKeys [][]byte) (*encryptionKeys, error) {
	if len(primaryKey) != encryptedKeySize {
		return nil, fmt.Errorf(appKeyErr)
	}
	prev := make([][]byte, 0, len(previousKeys))
	for i, k := range previousKeys {
		if len(k) != encryptedKeySize {
			return nil, fmt.Errorf("lagoon: app.previous_keys[%d] is empty or invalid (set SUMMER_APP__KEY to a 32-byte base64 value)", i)
		}
		prev = append(prev, cloneBytes(k))
	}
	id := byte(len(prev) + 1)
	if id > 0x0F {
		id = 0x0F
	}
	return &encryptionKeys{
		primary:   cloneBytes(primaryKey),
		previous:  prev,
		primaryID: id,
	}, nil
}

func setCurrentKeys(k *encryptionKeys) {
	keysMu.Lock()
	currentKeys = k
	keysMu.Unlock()
}

func clearEncryptionKeys() {
	setCurrentKeys(nil)
}

func liveKeys() (*encryptionKeys, error) {
	keysMu.RLock()
	k := currentKeys
	keysMu.RUnlock()
	if k == nil || len(k.primary) != encryptedKeySize {
		return nil, fmt.Errorf(appKeyErr)
	}
	return k, nil
}

// LoadAppKey reads app.key and app.previous_keys from cfg. Missing, short, or
// undecodable values fail loudly with no default (D-11).
func LoadAppKey(cfg *compass.Config) ([]byte, [][]byte, error) {
	if cfg == nil {
		return nil, nil, fmt.Errorf(appKeyErr)
	}
	primary, err := decodeAppKey(cfg.String("app.key"))
	if err != nil {
		return nil, nil, err
	}
	var previous [][]byte
	if v, ok := cfg.Lookup("app.previous_keys"); ok && v != nil {
		previous, err = decodePreviousKeys(v)
		if err != nil {
			return nil, nil, err
		}
	}
	return primary, previous, nil
}

func decodeAppKey(s string) ([]byte, error) {
	s = strings.TrimSpace(s)
	if s == "" {
		return nil, fmt.Errorf(appKeyErr)
	}
	raw, err := base64.StdEncoding.DecodeString(s)
	if err != nil || len(raw) != encryptedKeySize {
		return nil, fmt.Errorf(appKeyErr)
	}
	return raw, nil
}

func decodePreviousKeys(v any) ([][]byte, error) {
	items, ok := asStringList(v)
	if !ok {
		return nil, fmt.Errorf("lagoon: app.previous_keys is empty or invalid (set SUMMER_APP__KEY to a 32-byte base64 value)")
	}
	out := make([][]byte, 0, len(items))
	for _, item := range items {
		raw, err := decodeAppKey(item)
		if err != nil {
			return nil, fmt.Errorf("lagoon: app.previous_keys is empty or invalid (set SUMMER_APP__KEY to a 32-byte base64 value)")
		}
		out = append(out, raw)
	}
	return out, nil
}

func asStringList(v any) ([]string, bool) {
	switch list := v.(type) {
	case []string:
		return list, true
	case []any:
		out := make([]string, 0, len(list))
		for _, item := range list {
			s, ok := item.(string)
			if !ok {
				return nil, false
			}
			out = append(out, s)
		}
		return out, true
	default:
		return nil, false
	}
}

func deriveColumnKey(appKey []byte) ([]byte, error) {
	if len(appKey) != encryptedKeySize {
		return nil, fmt.Errorf(appKeyErr)
	}
	return hkdf.Key(sha256.New, appKey, nil, columnKeyInfo, encryptedKeySize)
}

func encryptPlaintext(plain []byte) (string, error) {
	keys, err := liveKeys()
	if err != nil {
		return "", err
	}
	colKey, err := deriveColumnKey(keys.primary)
	if err != nil {
		return "", err
	}
	block, err := aes.NewCipher(colKey)
	if err != nil {
		return "", fmt.Errorf("lagoon: encrypted aes: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", fmt.Errorf("lagoon: encrypted gcm: %w", err)
	}
	nonce := make([]byte, encryptedNonceSize)
	if _, err := rand.Read(nonce); err != nil {
		return "", fmt.Errorf("lagoon: encrypted nonce: %w", err)
	}
	sealed := gcm.Seal(nil, nonce, plain, nil)
	out := make([]byte, 1+len(nonce)+len(sealed))
	out[0] = encryptedFormatV1 | (keys.primaryID & 0x0F)
	copy(out[1:], nonce)
	copy(out[1+len(nonce):], sealed)
	return base64.StdEncoding.EncodeToString(out), nil
}

func decryptCiphertext(stored string) ([]byte, error) {
	keys, err := liveKeys()
	if err != nil {
		return nil, err
	}
	raw, err := base64.StdEncoding.DecodeString(stored)
	if err != nil {
		return nil, fmt.Errorf("lagoon: encrypted ciphertext is not base64")
	}
	if len(raw) < 1+encryptedNonceSize+16 {
		return nil, fmt.Errorf("lagoon: encrypted ciphertext is truncated")
	}
	if raw[0]&0xF0 != encryptedFormatV1 {
		return nil, fmt.Errorf("lagoon: encrypted ciphertext has unknown format")
	}
	nonce := raw[1 : 1+encryptedNonceSize]
	sealed := raw[1+encryptedNonceSize:]

	candidates := make([][]byte, 0, 1+len(keys.previous))
	candidates = append(candidates, keys.primary)
	candidates = append(candidates, keys.previous...)
	var last error
	for _, appKey := range candidates {
		plain, err := gcmOpen(appKey, nonce, sealed)
		if err == nil {
			return plain, nil
		}
		last = err
	}
	if last == nil {
		last = fmt.Errorf("lagoon: encrypted decrypt failed")
	}
	return nil, last
}

func gcmOpen(appKey, nonce, sealed []byte) ([]byte, error) {
	colKey, err := deriveColumnKey(appKey)
	if err != nil {
		return nil, err
	}
	block, err := aes.NewCipher(colKey)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	return gcm.Open(nil, nonce, sealed, nil)
}

func cloneBytes(b []byte) []byte {
	if b == nil {
		return nil
	}
	out := make([]byte, len(b))
	copy(out, b)
	return out
}

func loadEncryptionKeysFromApp(app *backpack.App) error {
	if app == nil || app.Config == nil {
		return fmt.Errorf(appKeyErr)
	}
	primary, previous, err := LoadAppKey(app.Config)
	if err != nil {
		return err
	}
	return PublishEncryptionKeys(app, primary, previous)
}