package infra

import (
	"bytes"
	"fmt"
	"os/exec"
	"strings"
)

// WGKeys holds a WireGuard Curve25519 key pair and an optional preshared key,
// all encoded as base64 strings.
type WGKeys struct {
	PrivateKey   string // base64 Curve25519 private key
	PublicKey    string // base64 Curve25519 public key
	PresharedKey string // base64 preshared key, empty if not requested
}

// WGKeygen generates a WireGuard key pair using `wg genkey` / `wg pubkey`.
// If withPSK is true it also runs `wg genpsk` to produce a preshared key.
// Requires the `wg` binary in PATH (install wireguard-tools).
// NEVER log PrivateKey or PresharedKey in plain text.
func WGKeygen(withPSK bool) (WGKeys, error) {
	// Generate private key
	privOut, err := runWG(nil, "genkey")
	if err != nil {
		return WGKeys{}, fmt.Errorf("wg genkey: %w", err)
	}
	privateKey := strings.TrimSpace(privOut)

	// Derive public key from private key
	pubOut, err := runWG(strings.NewReader(privateKey), "pubkey")
	if err != nil {
		return WGKeys{}, fmt.Errorf("wg pubkey: %w", err)
	}
	publicKey := strings.TrimSpace(pubOut)

	keys := WGKeys{
		PrivateKey: privateKey,
		PublicKey:  publicKey,
	}

	if withPSK {
		pskOut, err := runWG(nil, "genpsk")
		if err != nil {
			return WGKeys{}, fmt.Errorf("wg genpsk: %w", err)
		}
		keys.PresharedKey = strings.TrimSpace(pskOut)
	}

	return keys, nil
}

// runWG executes `wg <subcommand>`, optionally piping stdin, and returns stdout.
// stderr is captured and included in the error when exit code != 0.
func runWG(stdin interface{ Read([]byte) (int, error) }, subcommand string) (string, error) {
	cmd := exec.Command("wg", subcommand)
	if stdin != nil {
		cmd.Stdin = stdin
	}
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		return "", fmt.Errorf("%w: %s", err, strings.TrimSpace(stderr.String()))
	}
	return stdout.String(), nil
}