package orchestrator

import (
	"context"
	"fmt"
	"log"
	"time"

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

type ClusterJoiner struct {
	ExistingNode    string
	ClusterName     string
	JoinFingerprint string
}

func (c *ClusterJoiner) Join(ctx context.Context, hostIP string, privateKey string, publicKey string) error {
	client, err := c.connect(hostIP, privateKey)
	if err != nil {
		return fmt.Errorf("ssh connect to %s: %w", hostIP, err)
	}
	defer client.Close()

	// Join the cluster
	cmd := fmt.Sprintf("pvecm add %s --force", c.ExistingNode)
	log.Printf("cluster: running on %s: %s", hostIP, cmd)
	if err := c.runCmd(client, cmd); err != nil {
		return fmt.Errorf("pvecm add failed: %w", err)
	}
	log.Printf("cluster: %s joined successfully", hostIP)

	// Remove the ephemeral key from authorized_keys
	escaped := escapeForSed(publicKey)
	removeCmd := fmt.Sprintf(`sed -i '\|%s|d' /root/.ssh/authorized_keys`, escaped)
	if err := c.runCmd(client, removeCmd); err != nil {
		log.Printf("cluster: warning: failed to remove ephemeral key from %s: %v", hostIP, err)
	} else {
		log.Printf("cluster: ephemeral key removed from %s", hostIP)
	}

	return nil
}

func (c *ClusterJoiner) connect(hostIP string, privateKeyPEM string) (*ssh.Client, error) {
	signer, err := ssh.ParsePrivateKey([]byte(privateKeyPEM))
	if err != nil {
		return nil, fmt.Errorf("parse ephemeral key: %w", err)
	}
	config := &ssh.ClientConfig{
		User:            "root",
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(signer)},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         30 * time.Second,
	}
	return ssh.Dial("tcp", hostIP+":22", config)
}

func (c *ClusterJoiner) runCmd(client *ssh.Client, cmd string) error {
	session, err := client.NewSession()
	if err != nil {
		return fmt.Errorf("ssh session: %w", err)
	}
	defer session.Close()
	output, err := session.CombinedOutput(cmd)
	if err != nil {
		return fmt.Errorf("%w\noutput: %s", err, string(output))
	}
	return nil
}

func escapeForSed(s string) string {
	// Trim trailing newline and escape sed delimiter
	result := ""
	for _, c := range s {
		if c == '|' {
			result += `\|`
		} else if c == '\n' {
			continue
		} else {
			result += string(c)
		}
	}
	return result
}