package security

import (
	"crypto/ed25519"
	"crypto/rand"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"os"
	"path/filepath"

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

// SSHKeyMode is the file mode for the SSH private key. Matches the
// CA key mode (REQ-033 spirit: 0600 for private keys).
const SSHKeyMode os.FileMode = 0o600

// SSHPubMode is the file mode for the SSH public key (authorized_keys
// line). Matches the CA cert mode (0644 for public material).
const SSHPubMode os.FileMode = 0o644

const (
	sshKeyFile = "orca_ssh_key"
	sshPubFile = "orca_ssh_key.pub"
)

// GenerateOrLoadSSHKey returns the orca SSH keypair, generating it
// lazily on first call (D-037). The key is Ed25519 (smaller, faster,
// more secure than RSA for SSH auth), persisted as PKCS8 PEM to
// dir/orca_ssh_key (0600) and dir/orca_ssh_key.pub (0644).
//
// Idempotent: if both files exist with valid content, they are loaded
// and returned without regeneration. This matches the CAInit fast-path
// pattern (D-036 idempotency).
//
// Returns:
//   - keyPEM: PKCS8 PEM private key (parses with ssh.ParsePrivateKey)
//   - pubLine: authorized_keys line (ssh-ed25519 AAAA... comment\n)
func GenerateOrLoadSSHKey(dir string) (keyPEM, pubLine []byte, err error) {
	if dir == "" {
		return nil, nil, errors.New("GenerateOrLoadSSHKey: dir is required")
	}
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return nil, nil, fmt.Errorf("GenerateOrLoadSSHKey: mkdir: %w", err)
	}
	keyPath := filepath.Join(dir, sshKeyFile)
	pubPath := filepath.Join(dir, sshPubFile)

	// Fast path: existing key — load and return.
	if ok, err := bothExist(keyPath, pubPath); err != nil {
		return nil, nil, err
	} else if ok {
		keyPEM, err := os.ReadFile(keyPath)
		if err != nil {
			return nil, nil, fmt.Errorf("read SSH key: %w", err)
		}
		pubLine, err := os.ReadFile(pubPath)
		if err != nil {
			return nil, nil, fmt.Errorf("read SSH pub: %w", err)
		}
		return keyPEM, pubLine, nil
	}

	// Generate Ed25519 keypair.
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, nil, fmt.Errorf("GenerateOrLoadSSHKey: ed25519 gen: %w", err)
	}

	// Serialize private key as PKCS8 PEM (consistent with ca.key/server.key).
	keyDER, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return nil, nil, fmt.Errorf("GenerateOrLoadSSHKey: marshal key: %w", err)
	}
	keyPEM = pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: keyDER})

	// Serialize public key as authorized_keys line.
	sshPub, err := ssh.NewPublicKey(pub)
	if err != nil {
		return nil, nil, fmt.Errorf("GenerateOrLoadSSHKey: new pubkey: %w", err)
	}
	pubLine = ssh.MarshalAuthorizedKey(sshPub)

	// Persist with correct modes (atomic write + chmod).
	if err := writeAtomic(keyPath, SSHKeyMode, keyPEM); err != nil {
		return nil, nil, fmt.Errorf("write SSH key: %w", err)
	}
	if err := writeAtomic(pubPath, SSHPubMode, pubLine); err != nil {
		return nil, nil, fmt.Errorf("write SSH pub: %w", err)
	}

	return keyPEM, pubLine, nil
}