package yum

// 95% of AI generated code
// human designed & reviewed

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"errors"
	"io"
	"os"

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

const (
	saltSize  = 16
	keySize   = 32
	nonceSize = 12

	// Argon2 params (tune if needed)
	argonTime    = 1
	argonMemory  = 64 * 1024 // 64 MB
	argonThreads = 4
)

// YumSave encrypts and saves data securely to disk.
func YumSave(path string, data []byte, password []byte) error {
	// Generate salt
	salt := make([]byte, saltSize)
	if _, err := rand.Read(salt); err != nil {
		return err
	}

	// Derive key
	key := argon2.IDKey(password, salt, argonTime, argonMemory, argonThreads, keySize)

	// Create AES cipher
	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return err
	}

	// Generate nonce
	nonce := make([]byte, nonceSize)
	if _, err := rand.Read(nonce); err != nil {
		return err
	}

	// Encrypt
	ciphertext := gcm.Seal(nil, nonce, data, nil)

	// File format: [salt | nonce | ciphertext]
	fileData := append(salt, nonce...)
	fileData = append(fileData, ciphertext...)

	return os.WriteFile(path, fileData, 0600)
}

// YumLoad loads and decrypts data from disk.
func YumLoad(path string, password []byte) ([]byte, error) {
	fileData, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}

	if len(fileData) < saltSize+nonceSize {
		return nil, errors.New("file too short")
	}

	// Extract components
	salt := fileData[:saltSize]
	nonce := fileData[saltSize : saltSize+nonceSize]
	ciphertext := fileData[saltSize+nonceSize:]

	// Derive key
	key := argon2.IDKey(password, salt, argonTime, argonMemory, argonThreads, keySize)

	// Recreate cipher
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	// Decrypt
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, errors.New("decryption failed (wrong password or corrupted data)")
	}

	return plaintext, nil
}

func YumSeed(buf []byte) error {
	_, err := io.ReadFull(rand.Reader, buf)
	if err != nil {
		for i := range buf {
			buf[i] = 0
		}
		return err
	}
	return nil
}