package sshsupport

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"golang.org/x/crypto/ssh"
	"log"
	"os"
	"strings"
)

func generateHostKey(filename string, bitSize int) error {
	if _, err := os.Stat(filename); !os.IsNotExist(err) {
		log.Printf("Reusing key file '%s'", filename)
		return nil
	}
	privateKey, err := rsa.GenerateKey(rand.Reader, bitSize)
	if err != nil {
		return err
	}

	privateKeyPEM := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
	}

	privateKeyFile, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer privateKeyFile.Close()

	log.Printf("Generating key '%s'", filename)
	return pem.Encode(privateKeyFile, privateKeyPEM)
}

func GeneratePrivatePublicKey(bitsize int) (privateKeyData string, publicKeyData string, err error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, bitsize)
	if err != nil {
		return "", "", err
	}

	// Encode private key to PEM format
	privateKeyPEM := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
	}
	buf := bytes.Buffer{}
	if err := pem.Encode(&buf, privateKeyPEM); err != nil {
		return "", "", fmt.Errorf("failed to write private key: %v", err)
	}
	privateKeyData = string(buf.String())

	publicKey, err := ssh.NewPublicKey(&privateKey.PublicKey)
	if err != nil {
		return "", "", fmt.Errorf("failed to generate public key: %v", err)
	}

	publicKeyData = strings.TrimSpace(string(ssh.MarshalAuthorizedKey(publicKey)))
	return privateKeyData, publicKeyData, nil
}