import (
    "bytes"
    "io"
    "strings"

    "filippo.io/age"
    "filippo.io/age/armor"
)

func EncryptPayload(payload []byte, publicKey string) ([]byte, error) {
    recipient, err := age.ParseX25519Recipient(publicKey)
    if err != nil {
        return nil, err
    }

    var buf bytes.Buffer
    // Optional: wrap with armor for ASCII output (can omit if binary preferred)
    w := armor.NewWriter(&buf)
    encryptor := age.NewEncryptor(w, recipient)
    _, err = encryptor.Write(payload)
    if err != nil {
        return nil, err
    }
    if err := encryptor.Close(); err != nil {
        return nil, err
    }
    if err := w.Close(); err != nil {
        return nil, err
    }

    return buf.Bytes(), nil
}


func DecryptPayload(ciphertext []byte, privateKey string) ([]byte, error) {
    identity, err := age.ParseX25519Identity(privateKey)
    if err != nil {
        return nil, err
    }

    // Support armored input:
    r := bytes.NewReader(ciphertext)
    decoder := armor.NewReader(r)

    decryptor, err := age.Decrypt(decoder, identity)
    if err != nil {
        return nil, err
    }
    defer decryptor.Close()

    plaintext, err := io.ReadAll(decryptor)
    if err != nil {
        return nil, err
    }
    return plaintext, nil
}