package ageio

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"filippo.io/age"
)

// GenerateTestKeyPair generates a test age key pair and returns paths
func GenerateTestKeyPair(dir string) (identityPath, recipientPath string, err error) {
	// Generate identity
	identity, err := age.GenerateX25519Identity()
	if err != nil {
		return "", "", fmt.Errorf("generate identity: %w", err)
	}

	// Create identity file
	identityPath = filepath.Join(dir, "age.key")
	if err := os.WriteFile(identityPath, []byte(identity.String()+"\n"), 0600); err != nil {
		return "", "", fmt.Errorf("write identity file: %w", err)
	}

	// Create recipient file
	recipientPath = filepath.Join(dir, "age.pub")
	recipient := identity.Recipient().String()
	if err := os.WriteFile(recipientPath, []byte(recipient+"\n"), 0644); err != nil {
		return "", "", fmt.Errorf("write recipient file: %w", err)
	}

	return identityPath, recipientPath, nil
}

// GenerateTestKeys generates test keys in memory
func GenerateTestKeys() (identity age.Identity, recipient age.Recipient, err error) {
	id, err := age.GenerateX25519Identity()
	if err != nil {
		return nil, nil, fmt.Errorf("generate identity: %w", err)
	}

	return id, id.Recipient(), nil
}

// MustGenerateTestKeyPair generates a test key pair or panics
func MustGenerateTestKeyPair(dir string) (identityPath, recipientPath string) {
	identityPath, recipientPath, err := GenerateTestKeyPair(dir)
	if err != nil {
		panic(fmt.Sprintf("failed to generate test key pair: %v", err))
	}
	return identityPath, recipientPath
}

// EncryptBytes is a convenience function for one-shot encryption
func EncryptBytes(plaintext []byte, recipients ...age.Recipient) ([]byte, error) {
	var buf bytes.Buffer
	w, err := age.Encrypt(&buf, recipients...)
	if err != nil {
		return nil, fmt.Errorf("create encryptor: %w", err)
	}

	if _, err := w.Write(plaintext); err != nil {
		return nil, fmt.Errorf("write plaintext: %w", err)
	}

	if err := w.Close(); err != nil {
		return nil, fmt.Errorf("close encryptor: %w", err)
	}

	return buf.Bytes(), nil
}

// DecryptBytes is a convenience function for one-shot decryption
func DecryptBytes(ciphertext []byte, identities ...age.Identity) ([]byte, error) {
	r, err := age.Decrypt(bytes.NewReader(ciphertext), identities...)
	if err != nil {
		return nil, fmt.Errorf("create decryptor: %w", err)
	}

	plaintext, err := io.ReadAll(r)
	if err != nil {
		return nil, fmt.Errorf("read plaintext: %w", err)
	}

	return plaintext, nil
}