// Package keygen creates SSH key pairs natively in Go (no ssh-keygen
// dependency) and, optionally, appends a matching Host block to ~/.ssh/config.
//
// It is a rebuild of an older create_ssh_keypair.sh, fixing that script's
// cross-platform and safety problems:
//   - no macOS-only `sed -i ''` (public-key comments are handled in Go);
//   - the .pub file is kept, not deleted after display;
//   - ed25519 by default, RSA-4096 as an explicit opt-out;
//   - config appends are quoted/validated, back up the file first, and refuse
//     to create a duplicate Host alias.
package keygen

import (
	"crypto"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/rsa"
	"encoding/pem"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"

	"gitea.apointless.space/bsncubed/dial/internal/sshconf"
)

// Algorithm selects the key type.
type Algorithm string

const (
	AlgoED25519 Algorithm = "ed25519"
	AlgoRSA     Algorithm = "rsa"

	rsaBits = 4096
)

// Result holds the encoded key material ready to write to disk.
type Result struct {
	PrivatePEM []byte // OpenSSH-format private key (optionally passphrase-encrypted)
	PublicLine []byte // authorized_keys line, comment appended, trailing newline
}

// Generate creates a key pair of the given algorithm. comment is embedded in
// the private key and appended to the public line. A non-empty passphrase
// encrypts the private key at rest.
func Generate(algo Algorithm, comment string, passphrase []byte) (Result, error) {
	var privKey crypto.PrivateKey
	var pubKey crypto.PublicKey

	switch algo {
	case AlgoRSA:
		k, err := rsa.GenerateKey(rand.Reader, rsaBits)
		if err != nil {
			return Result{}, err
		}
		privKey, pubKey = k, &k.PublicKey
	case AlgoED25519, "":
		pub, priv, err := ed25519.GenerateKey(rand.Reader)
		if err != nil {
			return Result{}, err
		}
		privKey, pubKey = priv, pub
	default:
		return Result{}, fmt.Errorf("unknown algorithm %q", algo)
	}

	var block *pem.Block
	var err error
	if len(passphrase) > 0 {
		block, err = ssh.MarshalPrivateKeyWithPassphrase(privKey, comment, passphrase)
	} else {
		block, err = ssh.MarshalPrivateKey(privKey, comment)
	}
	if err != nil {
		return Result{}, fmt.Errorf("encoding private key: %w", err)
	}

	sshPub, err := ssh.NewPublicKey(pubKey)
	if err != nil {
		return Result{}, err
	}
	pubLine := strings.TrimRight(string(ssh.MarshalAuthorizedKey(sshPub)), "\n")
	if comment != "" {
		pubLine += " " + comment
	}
	pubLine += "\n"

	return Result{PrivatePEM: pem.EncodeToMemory(block), PublicLine: []byte(pubLine)}, nil
}

// WriteKeyPair writes the private key (0600) and public key (0644) into sshDir
// as <name> and <name>.pub. It refuses to overwrite existing files so a stray
// re-run can't clobber a key still in use.
func WriteKeyPair(sshDir, name string, r Result) (privPath, pubPath string, err error) {
	if err = os.MkdirAll(sshDir, 0o700); err != nil {
		return "", "", err
	}
	privPath = filepath.Join(sshDir, name)
	pubPath = privPath + ".pub"
	if exists(privPath) || exists(pubPath) {
		return privPath, pubPath, fmt.Errorf("key files for %q already exist; refusing to overwrite", name)
	}
	if err = os.WriteFile(privPath, r.PrivatePEM, 0o600); err != nil {
		return privPath, pubPath, err
	}
	if err = os.WriteFile(pubPath, r.PublicLine, 0o644); err != nil {
		return privPath, pubPath, err
	}
	return privPath, pubPath, nil
}

// HostEntry describes a Host block to append to ~/.ssh/config.
type HostEntry struct {
	Alias        string
	HostName     string
	User         string
	Port         string
	IdentityFile string
	Tags         []string
}

// AliasExists reports whether the config already declares the given Host alias.
func AliasExists(configPath, alias string) (bool, error) {
	hosts, err := sshconf.Parse(configPath)
	if err != nil {
		return false, err
	}
	for _, h := range hosts {
		if h.Alias == alias {
			return true, nil
		}
	}
	return false, nil
}

// AppendHostEntry appends a Host block to configPath. It validates fields
// (rejecting whitespace/newlines that could break or inject config), backs up
// an existing config first, and refuses to duplicate an existing alias.
func AppendHostEntry(configPath string, e HostEntry) (backup string, err error) {
	if err := validateEntry(e); err != nil {
		return "", err
	}
	if exists(configPath) {
		dup, err := AliasExists(configPath, e.Alias)
		if err != nil {
			return "", err
		}
		if dup {
			return "", fmt.Errorf("Host %q already exists in %s", e.Alias, configPath)
		}
		backup = fmt.Sprintf("%s.%s.bak", configPath, time.Now().Format("20060102-150405"))
		if err := copyFile(configPath, backup); err != nil {
			return "", fmt.Errorf("backing up config: %w", err)
		}
	} else if err := os.MkdirAll(filepath.Dir(configPath), 0o700); err != nil {
		return "", err
	}

	var b strings.Builder
	b.WriteString("\n")
	if len(e.Tags) > 0 {
		b.WriteString("# tag: " + strings.Join(e.Tags, ", ") + "\n")
	}
	b.WriteString("Host " + e.Alias + "\n")
	if e.HostName != "" {
		b.WriteString("    HostName " + e.HostName + "\n")
	}
	if e.User != "" {
		b.WriteString("    User " + e.User + "\n")
	}
	if e.Port != "" && e.Port != "22" {
		b.WriteString("    Port " + e.Port + "\n")
	}
	if e.IdentityFile != "" {
		b.WriteString("    IdentityFile " + e.IdentityFile + "\n")
		b.WriteString("    IdentitiesOnly yes\n")
	}

	f, err := os.OpenFile(configPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
	if err != nil {
		return backup, err
	}
	defer f.Close()
	if _, err := f.WriteString(b.String()); err != nil {
		return backup, err
	}
	return backup, nil
}

func validateEntry(e HostEntry) error {
	if strings.TrimSpace(e.Alias) == "" {
		return fmt.Errorf("host alias is required")
	}
	fields := map[string]string{
		"alias": e.Alias, "hostname": e.HostName, "user": e.User,
		"port": e.Port, "identityfile": e.IdentityFile,
	}
	for name, v := range fields {
		if strings.ContainsAny(v, "\n\r") {
			return fmt.Errorf("%s contains a newline", name)
		}
	}
	if strings.ContainsAny(e.Alias, " \t") {
		return fmt.Errorf("host alias %q must not contain spaces", e.Alias)
	}
	for _, t := range e.Tags {
		if strings.ContainsAny(t, "\n\r,") {
			return fmt.Errorf("tag %q must not contain newlines or commas", t)
		}
	}
	return nil
}

func exists(p string) bool {
	_, err := os.Stat(p)
	return err == nil
}

func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.OpenFile(dst, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
	if err != nil {
		return err
	}
	defer out.Close()
	if _, err := io.Copy(out, in); err != nil {
		return err
	}
	return out.Close()
}