package cryptokit

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"strings"

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

const maxCryptoMessageBytes = 64 << 20

func parseAgeRecipient(value string) (age.Recipient, error) {
	value = strings.TrimSpace(value)
	if recipient, err := age.ParseX25519Recipient(value); err == nil {
		return recipient, nil
	}
	if recipient, err := agessh.ParseRecipient(value); err == nil {
		return recipient, nil
	}
	return nil, errors.New("unsupported age recipient: expected native X25519 or SSH Ed25519/RSA")
}

func parseAgeIdentity(value string) (age.Identity, error) {
	value = strings.TrimSpace(value)
	if identity, err := age.ParseX25519Identity(value); err == nil {
		return identity, nil
	}
	if identity, err := agessh.ParseIdentity([]byte(value)); err == nil {
		return identity, nil
	}
	return nil, errors.New("unsupported age identity: expected native X25519 or SSH Ed25519/RSA private key")
}

// EncryptAge encrypts to a user-provided native age X25519 recipient or SSH
// Ed25519/RSA recipient and returns ASCII-armored age text.
func EncryptAge(message []byte, recipientKey string) ([]byte, error) {
	recipient, err := parseAgeRecipient(recipientKey)
	if err != nil {
		return nil, fmt.Errorf("parse age recipient: %w", err)
	}
	var output bytes.Buffer
	armored := agearmor.NewWriter(&output)
	writer, err := age.Encrypt(armored, recipient)
	if err != nil {
		_ = armored.Close()
		return nil, fmt.Errorf("create age encryption: %w", err)
	}
	if _, err := writer.Write(message); err != nil {
		_ = writer.Close()
		_ = armored.Close()
		return nil, fmt.Errorf("write age message: %w", err)
	}
	if err := writer.Close(); err != nil {
		_ = armored.Close()
		return nil, fmt.Errorf("close age message: %w", err)
	}
	if err := armored.Close(); err != nil {
		return nil, fmt.Errorf("close age armor: %w", err)
	}
	return output.Bytes(), nil
}

// DecryptAge decrypts an armored or binary age message with a user-provided
// native X25519 identity or an unencrypted SSH Ed25519/RSA private key.
func DecryptAge(message []byte, identityKey string) ([]byte, error) {
	identity, err := parseAgeIdentity(identityKey)
	if err != nil {
		return nil, fmt.Errorf("parse age identity: %w", err)
	}
	var input io.Reader = bytes.NewReader(message)
	if bytes.HasPrefix(bytes.TrimSpace(message), []byte(agearmor.Header)) {
		input = agearmor.NewReader(input)
	}
	reader, err := age.Decrypt(input, identity)
	if err != nil {
		return nil, fmt.Errorf("create age decryption: %w", err)
	}
	plaintext, err := io.ReadAll(io.LimitReader(reader, maxCryptoMessageBytes+1))
	if err != nil {
		return nil, fmt.Errorf("read age message: %w", err)
	}
	if len(plaintext) > maxCryptoMessageBytes {
		return nil, errors.New("decrypted age message exceeds the 64 MiB limit")
	}
	return plaintext, nil
}