// Package sshcert builds SSH certificates with Shellstream extensions,
// bridging SPIFFE identity and Guildhouse governance metadata.
package sshcert

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"math/big"
	"net/url"
	"strings"
	"time"

	"github.com/guildhouse-cooperative/guildhouse-spire-plugins/pkg/shellstream"
	"golang.org/x/crypto/ssh"
)

// Config holds SSH certificate builder configuration.
type Config struct {
	// TrustDomain is the SPIFFE trust domain.
	TrustDomain string

	// MaxValidSeconds is the server-enforced maximum certificate TTL.
	// SSH-SVID spec mandates 3600 (1 hour).
	MaxValidSeconds uint64

	// MinValidSeconds is the server-enforced minimum certificate TTL.
	// SSH-SVID spec mandates 30 seconds.
	MinValidSeconds uint64

	// DefaultValidSeconds is the TTL used when the request does not specify one.
	// SSH-SVID spec recommends 300 (5 minutes).
	DefaultValidSeconds uint64

	// AllowedPrincipals restricts which SSH principals may appear in certificates.
	// If empty, only the SPIFFE ID leaf component is permitted.
	AllowedPrincipals []string

	// CAKey is the SSH CA signing key. If nil, Build will generate an ephemeral one.
	CAKey ssh.Signer
}

// CertRequest describes an SSH certificate to build.
type CertRequest struct {
	// SpiffeID is the workload's SPIFFE ID (used as principal).
	SpiffeID string

	// Extensions are the Shellstream governance extensions to embed.
	Extensions *shellstream.ShellstreamExtensions

	// ValidSeconds is the requested certificate lifetime in seconds.
	// Clamped to [MinValidSeconds, MaxValidSeconds] by the Builder.
	ValidSeconds uint64

	// Principals are additional SSH principals beyond the SPIFFE ID.
	// Validated against Builder.AllowedPrincipals.
	Principals []string
}

// Builder creates SSH certificates with Shellstream extensions.
type Builder struct {
	config Config
}

// NewBuilder creates an SSH certificate builder.
func NewBuilder(cfg Config) (*Builder, error) {
	if cfg.TrustDomain == "" {
		return nil, fmt.Errorf("sshcert: trust domain is required")
	}
	// Apply spec defaults if not configured.
	if cfg.MaxValidSeconds == 0 {
		cfg.MaxValidSeconds = 3600 // 1 hour per SSH-SVID spec
	}
	if cfg.MinValidSeconds == 0 {
		cfg.MinValidSeconds = 30 // 30 seconds per SSH-SVID spec
	}
	if cfg.DefaultValidSeconds == 0 {
		cfg.DefaultValidSeconds = 300 // 5 minutes per SSH-SVID spec
	}
	if cfg.MinValidSeconds > cfg.MaxValidSeconds {
		return nil, fmt.Errorf("sshcert: min TTL %d exceeds max TTL %d", cfg.MinValidSeconds, cfg.MaxValidSeconds)
	}
	if cfg.DefaultValidSeconds < cfg.MinValidSeconds || cfg.DefaultValidSeconds > cfg.MaxValidSeconds {
		return nil, fmt.Errorf("sshcert: default TTL %d outside [%d, %d]",
			cfg.DefaultValidSeconds, cfg.MinValidSeconds, cfg.MaxValidSeconds)
	}
	return &Builder{config: cfg}, nil
}

// Build creates an SSH certificate from the request.
// Generates an Ed25519 user keypair, builds the certificate with Shellstream
// governance extensions, and signs it with the CA key.
func (b *Builder) Build(req *CertRequest) ([]byte, error) {
	if req == nil {
		return nil, fmt.Errorf("sshcert: nil request")
	}
	if req.SpiffeID == "" {
		return nil, fmt.Errorf("sshcert: spiffe ID is required")
	}

	// TTL enforcement.
	ttl := req.ValidSeconds
	if ttl == 0 {
		ttl = b.config.DefaultValidSeconds
	}
	if ttl < b.config.MinValidSeconds {
		return nil, fmt.Errorf("sshcert: requested TTL %d below minimum %d", ttl, b.config.MinValidSeconds)
	}
	if ttl > b.config.MaxValidSeconds {
		return nil, fmt.Errorf("sshcert: requested TTL %d exceeds maximum %d", ttl, b.config.MaxValidSeconds)
	}

	// Principal validation.
	principals := req.Principals
	if len(principals) == 0 {
		leaf, err := SpiffeIDLeaf(req.SpiffeID)
		if err != nil {
			return nil, fmt.Errorf("sshcert: invalid SPIFFE ID: %w", err)
		}
		principals = []string{leaf}
	}

	if len(b.config.AllowedPrincipals) > 0 {
		for _, p := range principals {
			if !containsString(b.config.AllowedPrincipals, p) {
				return nil, fmt.Errorf("sshcert: principal %q not in allowed list", p)
			}
		}
	}

	// Validate SPIFFE ID format for all paths.
	if _, err := SpiffeIDLeaf(req.SpiffeID); err != nil {
		return nil, fmt.Errorf("sshcert: invalid SPIFFE ID: %w", err)
	}

	// Generate user keypair.
	pubKey, _, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("sshcert: generate user key: %w", err)
	}

	sshPubKey, err := ssh.NewPublicKey(pubKey)
	if err != nil {
		return nil, fmt.Errorf("sshcert: convert public key: %w", err)
	}

	// Build extensions map with Shellstream governance metadata.
	extensions := make(map[string]string)
	extensions["permit-pty"] = ""
	extensions["permit-user-rc"] = ""

	// Embed the SPIFFE ID as a critical option for verification.
	criticalOptions := map[string]string{
		"spiffe-id": req.SpiffeID,
	}

	// Encode Shellstream extensions if present.
	if req.Extensions != nil {
		ssExtensions, err := shellstream.Encode(req.Extensions)
		if err != nil {
			return nil, fmt.Errorf("sshcert: encode shellstream extensions: %w", err)
		}
		for k, v := range ssExtensions {
			extensions[k] = v
		}
	}

	// Generate a random serial number.
	serialBytes := make([]byte, 8)
	if _, err := rand.Read(serialBytes); err != nil {
		return nil, fmt.Errorf("sshcert: generate serial: %w", err)
	}
	serial := binary.BigEndian.Uint64(serialBytes)

	now := time.Now()
	cert := &ssh.Certificate{
		CertType:        ssh.UserCert,
		Key:             sshPubKey,
		Serial:          serial,
		KeyId:           req.SpiffeID,
		ValidPrincipals: principals,
		ValidAfter:      uint64(now.Unix()),
		ValidBefore:     uint64(now.Add(time.Duration(ttl) * time.Second).Unix()),
		Permissions: ssh.Permissions{
			CriticalOptions: criticalOptions,
			Extensions:      extensions,
		},
		Nonce: serialBytes, // Reuse random bytes as nonce.
	}

	// Sign with CA key.
	caKey := b.config.CAKey
	if caKey == nil {
		// Generate ephemeral CA key for development/testing.
		_, caPriv, err := ed25519.GenerateKey(rand.Reader)
		if err != nil {
			return nil, fmt.Errorf("sshcert: generate ephemeral CA: %w", err)
		}
		caKey, err = ssh.NewSignerFromKey(caPriv)
		if err != nil {
			return nil, fmt.Errorf("sshcert: create CA signer: %w", err)
		}
	}

	if err := cert.SignCert(rand.Reader, caKey); err != nil {
		return nil, fmt.Errorf("sshcert: sign certificate: %w", err)
	}

	return ssh.MarshalAuthorizedKey(cert), nil
}

// BuildWithSerial is like Build but also returns the generated serial number.
func (b *Builder) BuildWithSerial(req *CertRequest) (certBytes []byte, serial *big.Int, err error) {
	certBytes, err = b.Build(req)
	if err != nil {
		return nil, nil, err
	}
	// Parse back to extract serial.
	key, _, _, _, parseErr := ssh.ParseAuthorizedKey(certBytes)
	if parseErr != nil {
		return certBytes, nil, nil // Return cert even if we can't extract serial.
	}
	if cert, ok := key.(*ssh.Certificate); ok {
		serial = new(big.Int).SetUint64(cert.Serial)
	}
	return certBytes, serial, nil
}

// SpiffeIDLeaf extracts the last path component from a SPIFFE ID URI.
//
//	spiffe://guildhouse.dev/ns/prod/sa/web-server → web-server
//	spiffe://guildhouse.dev/operator/tyler → tyler
func SpiffeIDLeaf(spiffeID string) (string, error) {
	u, err := url.Parse(spiffeID)
	if err != nil {
		return "", fmt.Errorf("not a valid URI: %w", err)
	}
	if u.Scheme != "spiffe" {
		return "", fmt.Errorf("expected spiffe:// scheme, got %q", u.Scheme)
	}
	if u.Host == "" {
		return "", fmt.Errorf("missing trust domain")
	}
	path := strings.TrimPrefix(u.Path, "/")
	if path == "" {
		return "", fmt.Errorf("empty path")
	}
	parts := strings.Split(path, "/")
	leaf := parts[len(parts)-1]
	if leaf == "" {
		return "", fmt.Errorf("trailing slash in path")
	}
	return leaf, nil
}

func containsString(haystack []string, needle string) bool {
	for _, s := range haystack {
		if s == needle {
			return true
		}
	}
	return false
}