package keys

import (
	"crypto/md5"
	"encoding/base64"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

const authorizedKeysPath = "/root/.ssh/authorized_keys"

func ReadAuthorizedKeys() ([]string, error) {
	data, err := os.ReadFile(authorizedKeysPath)
	if os.IsNotExist(err) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}

	var lines []string
	for _, line := range strings.Split(string(data), "\n") {
		line = strings.TrimSpace(line)
		if line != "" && !strings.HasPrefix(line, "#") {
			lines = append(lines, line)
		}
	}
	return lines, nil
}

func WriteAuthorizedKeys(keys []string) error {
	dir := filepath.Dir(authorizedKeysPath)
	if err := os.MkdirAll(dir, 0700); err != nil {
		return fmt.Errorf("mkdir %s: %w", dir, err)
	}

	content := strings.Join(keys, "\n")
	if len(keys) > 0 {
		content += "\n"
	}

	tmpPath := authorizedKeysPath + ".tmp"
	if err := os.WriteFile(tmpPath, []byte(content), 0600); err != nil {
		return fmt.Errorf("write tmp: %w", err)
	}

	if err := os.Rename(tmpPath, authorizedKeysPath); err != nil {
		os.Remove(tmpPath)
		return fmt.Errorf("rename: %w", err)
	}

	return os.Chmod(authorizedKeysPath, 0600)
}

func FingerprintLines(lines []string) map[string]bool {
	fp := make(map[string]bool, len(lines))
	for _, line := range lines {
		fp[fingerprint(line)] = true
	}
	return fp
}

func StateChanged(current, desired []string) bool {
	if len(current) != len(desired) {
		return true
	}
	cur := FingerprintLines(current)
	for _, line := range desired {
		if !cur[fingerprint(line)] {
			return true
		}
	}
	return false
}

func fingerprint(pubKey string) string {
	parts := strings.Fields(pubKey)
	if len(parts) < 2 {
		return pubKey
	}
	raw, err := base64.StdEncoding.DecodeString(parts[1])
	if err != nil {
		return pubKey
	}
	sum := md5.Sum(raw)
	var pairs []string
	for _, b := range sum {
		pairs = append(pairs, fmt.Sprintf("%02x", b))
	}
	return "MD5:" + strings.Join(pairs, ":")
}

// KeyGenOptions controls how ssh-keygen is invoked.
type KeyGenOptions struct {
	KeyType    string // ed25519 (default), rsa, ecdsa
	KeySize    int    // bits; used for rsa and ecdsa
	Passphrase string // empty = no passphrase
	Comment    string // embedded in the public key
}

// GenerateKeyPair generates an SSH keypair and returns the public key.
// The private key is written to keyPath; keyPath+".pub" holds the public key.
func GenerateKeyPair(keyPath string, opts KeyGenOptions) (string, error) {
	if err := os.MkdirAll(filepath.Dir(keyPath), 0700); err != nil {
		return "", err
	}

	keyType := opts.KeyType
	if keyType == "" {
		keyType = "ed25519"
	}

	args := []string{
		"-t", keyType,
		"-f", keyPath,
		"-N", opts.Passphrase,
		"-C", opts.Comment,
	}
	if opts.KeySize > 0 && keyType != "ed25519" {
		args = append(args, "-b", fmt.Sprintf("%d", opts.KeySize))
	}

	cmd := exec.Command("ssh-keygen", args...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("ssh-keygen: %w: %s", err, out)
	}

	pubData, err := os.ReadFile(keyPath + ".pub")
	if err != nil {
		return "", fmt.Errorf("read pubkey: %w", err)
	}
	return strings.TrimSpace(string(pubData)), nil
}