package restore

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/pem"
	"fmt"
	"os"
	"path/filepath"
	"strings"

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

// SSHKeyPair holds an ephemeral SSH key pair
type SSHKeyPair struct {
	PrivateKeyPath string
	PublicKey      string
}

// GenerateEphemeralKey creates a temporary ED25519 SSH key pair
func GenerateEphemeralKey() (*SSHKeyPair, error) {
	// Generate ED25519 key pair
	pubKey, privKey, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("failed to generate key: %w", err)
	}

	// Convert to SSH format
	sshPubKey, err := ssh.NewPublicKey(pubKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create SSH public key: %w", err)
	}

	// Marshal public key
	pubKeyStr := strings.TrimSpace(string(ssh.MarshalAuthorizedKey(sshPubKey)))

	// Create temp directory for key
	tmpDir, err := os.MkdirTemp("", "recover-ssh-")
	if err != nil {
		return nil, fmt.Errorf("failed to create temp dir: %w", err)
	}

	// Write private key in OpenSSH format
	privKeyPath := filepath.Join(tmpDir, "id_ed25519")

	// Marshal private key to OpenSSH format
	pemBlock, err := ssh.MarshalPrivateKey(privKey, "")
	if err != nil {
		return nil, fmt.Errorf("failed to marshal private key: %w", err)
	}

	privKeyPEM := pem.EncodeToMemory(pemBlock)
	if err := os.WriteFile(privKeyPath, privKeyPEM, 0600); err != nil {
		return nil, fmt.Errorf("failed to write private key: %w", err)
	}

	return &SSHKeyPair{
		PrivateKeyPath: privKeyPath,
		PublicKey:      pubKeyStr,
	}, nil
}

// Cleanup removes the ephemeral key files
func (k *SSHKeyPair) Cleanup() {
	if k.PrivateKeyPath != "" {
		os.RemoveAll(filepath.Dir(k.PrivateKeyPath))
	}
}

// runSSHKeyMerge merges original authorized_keys with ephemeral key
func (p *Pipeline) runSSHKeyMerge(ctx context.Context) error {
	// First, backup current authorized_keys
	backupCmd := "cp /root/.ssh/authorized_keys /root/.ssh/authorized_keys.ephemeral 2>/dev/null || true"
	p.remoteCmd(ctx, backupCmd)

	// Check if we have original keys in the restored /root
	checkCmd := "cat /root/.ssh/authorized_keys.original 2>/dev/null || cat /srv/restore/root/.ssh/authorized_keys 2>/dev/null || echo ''"
	originalKeys, _ := p.remoteCmdOutput(ctx, checkCmd)

	// Get current (ephemeral) keys
	currentKeys, _ := p.remoteCmdOutput(ctx, "cat /root/.ssh/authorized_keys 2>/dev/null || echo ''")

	// Merge keys (unique)
	allKeys := make(map[string]bool)
	for _, key := range strings.Split(currentKeys, "\n") {
		key = strings.TrimSpace(key)
		if key != "" && !strings.HasPrefix(key, "#") {
			allKeys[key] = true
		}
	}
	for _, key := range strings.Split(originalKeys, "\n") {
		key = strings.TrimSpace(key)
		if key != "" && !strings.HasPrefix(key, "#") {
			allKeys[key] = true
		}
	}

	// Write merged keys
	var mergedKeys []string
	for key := range allKeys {
		mergedKeys = append(mergedKeys, key)
	}

	mergeCmd := fmt.Sprintf("mkdir -p /root/.ssh && echo '%s' > /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys",
		strings.Join(mergedKeys, "\n"))

	return p.remoteCmd(ctx, mergeCmd)
}