package maintenance

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"io"

	"golang.org/x/crypto/hkdf"
)

const (
	EnvelopeVersion uint32 = 1
	// MaxEnvelopePlaintext bounds the crypto allocation. Request/result parsers
	// impose their smaller operation-specific limits before execution.
	MaxEnvelopePlaintext = 1 << 20
	envelopePrefix       = "CME1"
	ephemeralBytes       = 65
	nonceBytes           = 12
	tagBytes             = 16
	envelopeInfo         = "carlos-maintenance-v1"
)

type Direction string

const (
	DirectionInput       Direction = "input"
	DirectionResult      Direction = "result"
	DirectionEnvironment Direction = "environment"
)

var ErrEnvelope = errors.New("maintenance: invalid private envelope")

// Binding authenticates identifiers already validated by the request handler.
// It does not establish ownership or authorize an operation.
type Binding struct {
	Version   uint32
	Direction Direction
	Account   string
	App       string
	Host      string
	RequestID string
	AttemptID string
	Recipient string
}

// RecipientFingerprint names the P-256 key that can open the envelope.
func RecipientFingerprint(pub *ecdh.PublicKey) (string, error) {
	if pub == nil || pub.Curve() != ecdh.P256() {
		return "", ErrEnvelope
	}
	sum := sha256.Sum256(pub.Bytes())
	return hex.EncodeToString(sum[:]), nil
}

func appendField(dst []byte, field []byte) []byte {
	dst = binary.BigEndian.AppendUint32(dst, uint32(len(field)))
	return append(dst, field...)
}
func bindingAAD(b Binding, pub *ecdh.PublicKey) ([]byte, error) {
	fingerprint, err := RecipientFingerprint(pub)
	if err != nil || b.Version != EnvelopeVersion || b.Recipient != fingerprint {
		return nil, ErrEnvelope
	}
	if b.Direction != DirectionInput && b.Direction != DirectionResult && b.Direction != DirectionEnvironment {
		return nil, ErrEnvelope
	}
	fields := []struct {
		s   string
		max int
	}{{b.Account, 64}, {b.App, 64}, {b.Host, 253}, {b.RequestID, 64}, {b.AttemptID, 64}}
	for _, f := range fields {
		if len(f.s) == 0 || len(f.s) > f.max {
			return nil, ErrEnvelope
		}
		for i := range len(f.s) {
			if f.s[i] < 0x21 || f.s[i] > 0x7e {
				return nil, ErrEnvelope
			}
		}
	}
	aad := appendField(nil, []byte(envelopePrefix))
	aad = appendField(aad, binary.BigEndian.AppendUint32(nil, b.Version))
	aad = appendField(aad, []byte(b.Direction))
	for _, f := range fields {
		aad = appendField(aad, []byte(f.s))
	}
	return appendField(aad, []byte(b.Recipient)), nil
}
func envelopeGCM(shared []byte) (cipher.AEAD, error) {
	key := make([]byte, 32)
	if _, err := io.ReadFull(hkdf.New(sha256.New, shared, nil, []byte(envelopeInfo)), key); err != nil {
		return nil, ErrEnvelope
	}
	block, err := aes.NewCipher(key)
	clear(key)
	if err != nil {
		return nil, ErrEnvelope
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, ErrEnvelope
	}
	return gcm, nil
}

// Seal encrypts private input, output or an environment snapshot for one
// recipient and one immutable operation binding.
func Seal(pub *ecdh.PublicKey, b Binding, plaintext []byte) (string, error) {
	if len(plaintext) > MaxEnvelopePlaintext {
		return "", ErrEnvelope
	}
	if _, err := bindingAAD(b, pub); err != nil {
		return "", ErrEnvelope
	}
	eph, err := ecdh.P256().GenerateKey(rand.Reader)
	if err != nil {
		return "", ErrEnvelope
	}
	nonce := make([]byte, nonceBytes)
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", ErrEnvelope
	}
	return sealWith(eph, nonce, pub, b, plaintext)
}

// sealWith permits deterministic vectors without replacing production entropy.
func sealWith(eph *ecdh.PrivateKey, nonce []byte, pub *ecdh.PublicKey, b Binding, plaintext []byte) (string, error) {
	if eph == nil || eph.Curve() != ecdh.P256() || len(nonce) != nonceBytes || len(plaintext) > MaxEnvelopePlaintext {
		return "", ErrEnvelope
	}
	aad, err := bindingAAD(b, pub)
	if err != nil {
		return "", ErrEnvelope
	}
	shared, err := eph.ECDH(pub)
	if err != nil {
		return "", ErrEnvelope
	}
	defer clear(shared)
	gcm, err := envelopeGCM(shared)
	if err != nil {
		return "", ErrEnvelope
	}
	raw := make([]byte, 0, len(envelopePrefix)+ephemeralBytes+nonceBytes+len(plaintext)+tagBytes)
	raw = append(raw, []byte(envelopePrefix)...)
	raw = append(raw, eph.PublicKey().Bytes()...)
	raw = append(raw, nonce...)
	raw = gcm.Seal(raw, nonce, plaintext, aad)
	return base64.StdEncoding.EncodeToString(raw), nil
}

// Open refuses changed bindings and noncanonical or oversized envelopes before
// returning plaintext. No error includes caller input or decrypted material.
func Open(priv *ecdh.PrivateKey, b Binding, envelope string) ([]byte, error) {
	const overhead = len(envelopePrefix) + ephemeralBytes + nonceBytes + tagBytes
	if priv == nil || priv.Curve() != ecdh.P256() || len(envelope) > base64.StdEncoding.EncodedLen(overhead+MaxEnvelopePlaintext) {
		return nil, ErrEnvelope
	}
	aad, err := bindingAAD(b, priv.PublicKey())
	if err != nil {
		return nil, ErrEnvelope
	}
	raw, err := base64.StdEncoding.Strict().DecodeString(envelope)
	if err != nil || len(raw) < overhead || len(raw) > overhead+MaxEnvelopePlaintext || base64.StdEncoding.EncodeToString(raw) != envelope {
		return nil, ErrEnvelope
	}
	if string(raw[:len(envelopePrefix)]) != envelopePrefix {
		return nil, ErrEnvelope
	}
	pointEnd := len(envelopePrefix) + ephemeralBytes
	eph, err := ecdh.P256().NewPublicKey(raw[len(envelopePrefix):pointEnd])
	if err != nil {
		return nil, ErrEnvelope
	}
	shared, err := priv.ECDH(eph)
	if err != nil {
		return nil, ErrEnvelope
	}
	defer clear(shared)
	gcm, err := envelopeGCM(shared)
	if err != nil {
		return nil, ErrEnvelope
	}
	plaintext, err := gcm.Open(nil, raw[pointEnd:pointEnd+nonceBytes], raw[pointEnd+nonceBytes:], aad)
	if err != nil {
		return nil, ErrEnvelope
	}
	return plaintext, nil
}