package ssh

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
)

type KeyPair struct {
	PrivateKey []byte
	PublicKey  []byte
	Passphrase string
}

func GenerateKeyPair(passphrase string) (*KeyPair, error) {
	if passphrase == "" {
		var err error
		passphrase, err = GenerateSecurePassphrase(32)
		if err != nil {
			return nil, fmt.Errorf("failed to generate passphrase: %w", err)
		}
	}

	// Create temporary directory for key generation
	tmpDir, err := os.MkdirTemp("", "sslh-lab-keygen-*")
	if err != nil {
		return nil, fmt.Errorf("failed to create temp directory: %w", err)
	}
	defer os.RemoveAll(tmpDir)

	tmpKeyPath := filepath.Join(tmpDir, "id_ed25519")

	// Use ssh-keygen to generate OpenSSH format key with passphrase
	cmd := exec.Command("ssh-keygen", "-t", "ed25519", "-f", tmpKeyPath, "-N", passphrase, "-C", "sslh-lab-generated")
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("failed to generate SSH key with ssh-keygen: %w", err)
	}

	// Read the generated private key
	privateKeyBytes, err := os.ReadFile(tmpKeyPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read generated private key: %w", err)
	}

	// Read the generated public key
	publicKeyBytes, err := os.ReadFile(tmpKeyPath + ".pub")
	if err != nil {
		return nil, fmt.Errorf("failed to read generated public key: %w", err)
	}

	return &KeyPair{
		PrivateKey: privateKeyBytes,
		PublicKey:  publicKeyBytes,
		Passphrase: passphrase,
	}, nil
}

func SaveKeyPair(keyPair *KeyPair, outputDir string) (string, string, error) {
	if err := os.MkdirAll(outputDir, 0700); err != nil {
		return "", "", fmt.Errorf("failed to create output directory: %w", err)
	}

	privateKeyPath := filepath.Join(outputDir, "id_ed25519")
	publicKeyPath := filepath.Join(outputDir, "id_ed25519.pub")

	if err := os.WriteFile(privateKeyPath, keyPair.PrivateKey, 0600); err != nil {
		return "", "", fmt.Errorf("failed to write private key: %w", err)
	}

	if err := os.WriteFile(publicKeyPath, keyPair.PublicKey, 0644); err != nil {
		return "", "", fmt.Errorf("failed to write public key: %w", err)
	}

	return privateKeyPath, publicKeyPath, nil
}

func LoadPublicKey(publicKeyPath string) (string, error) {
	publicKeyBytes, err := os.ReadFile(publicKeyPath)
	if err != nil {
		return "", fmt.Errorf("failed to read public key file: %w", err)
	}

	return string(publicKeyBytes), nil
}