package tlsconf

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"net"
	"os"
	"path/filepath"
	"time"
)

// CA is a certificate authority for one llmgw cluster.
//
// It exists because the alternative is worse. Public authorities no longer
// issue certificates carrying the clientAuth usage, so a Let's Encrypt
// certificate cannot authenticate an agent however convenient it would be. And
// a self-signed certificate per node would mean every machine trusting every
// other machine's certificate individually, which stops being manageable the
// moment the fleet changes.
//
// A private CA is the ordinary answer: one root everyone trusts, leaves issued
// from it, and identity carried in the subject where the control plane can read
// it. Nothing outside llmgw ever sees these certificates — the client-facing
// API is fronted by a publicly trusted certificate instead.
type CA struct {
	Cert *x509.Certificate
	Key  *ecdsa.PrivateKey
	// PEM is the encoded certificate, which is what goes on every machine.
	PEM []byte
}

// CAValidity is how long a new authority is good for.
//
// Long, deliberately. Rotating a CA means touching every machine at once, so it
// should be a planned event rather than something that surprises a cluster.
const CAValidity = 10 * 365 * 24 * time.Hour

// LeafValidity is how long an issued certificate is good for.
//
// A year rather than 90 days: there is no ACME here, renewal is a manual step,
// and a short lifetime would turn that into recurring unplanned work. Long
// enough to be calm, short enough that a leaked key is not permanent.
const LeafValidity = 365 * 24 * time.Hour

// Backdate is how far before its issue a certificate is valid from, to
// tolerate clock skew across the fleet.
const Backdate = time.Hour

// IssuedAt is when this authority issued a certificate: its NotBefore, less
// the backdating. What a node's revocation is compared against.
func IssuedAt(cert *x509.Certificate) time.Time { return cert.NotBefore.Add(Backdate) }

// NewCA creates an authority.
func NewCA(commonName string) (*CA, error) {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}
	serial, err := randomSerial()
	if err != nil {
		return nil, err
	}

	tmpl := &x509.Certificate{
		SerialNumber:          serial,
		Subject:               pkix.Name{CommonName: commonName},
		NotBefore:             time.Now().Add(-time.Hour), // tolerate clock skew across the fleet
		NotAfter:              time.Now().Add(CAValidity),
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		BasicConstraintsValid: true,
		IsCA:                  true,
		MaxPathLenZero:        true, // leaves only; this CA never signs another CA
	}

	der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
	if err != nil {
		return nil, err
	}
	cert, err := x509.ParseCertificate(der)
	if err != nil {
		return nil, err
	}
	return &CA{Cert: cert, Key: key, PEM: pemEncode("CERTIFICATE", der)}, nil
}

// LoadCA reads an authority from disk.
func LoadCA(certPath, keyPath string) (*CA, error) {
	certPEM, err := os.ReadFile(certPath)
	if err != nil {
		return nil, err
	}
	block, _ := pem.Decode(certPEM)
	if block == nil {
		return nil, fmt.Errorf("%s is not PEM", certPath)
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, err
	}

	keyPEM, err := os.ReadFile(keyPath)
	if err != nil {
		return nil, err
	}
	kb, _ := pem.Decode(keyPEM)
	if kb == nil {
		return nil, fmt.Errorf("%s is not PEM", keyPath)
	}
	key, err := x509.ParseECPrivateKey(kb.Bytes)
	if err != nil {
		return nil, err
	}
	return &CA{Cert: cert, Key: key, PEM: certPEM}, nil
}

// Leaf is an issued certificate and its key, both PEM encoded.
type Leaf struct {
	CertPEM []byte
	KeyPEM  []byte
}

// IssueServer issues a certificate for something agents connect to.
//
// The hosts are what matter here, not the common name: since Go 1.15 hostname
// verification reads subject alternative names only and ignores the common name
// entirely. A certificate without the address agents dial — including the bare
// IP, if that is what they dial — is rejected no matter what its subject says.
func (ca *CA) IssueServer(commonName string, hosts []string) (Leaf, error) {
	if len(hosts) == 0 {
		return Leaf{}, fmt.Errorf("a server certificate needs at least one host; " +
			"without a matching SAN every client rejects it")
	}
	return ca.issue(commonName, hosts, x509.ExtKeyUsageServerAuth)
}

// IssueClient issues a certificate identifying one node.
//
// The common name is the node id, and the control plane refuses any agent whose
// claimed id does not match it. That is the whole point of the exercise: with
// it, a node's identity is something it proves rather than something it picks.
func (ca *CA) IssueClient(nodeID string) (Leaf, error) {
	if nodeID == "" {
		return Leaf{}, fmt.Errorf("a client certificate needs a node id as its common name")
	}
	return ca.issue(nodeID, nil, x509.ExtKeyUsageClientAuth)
}

// IssuePeer issues a certificate that serves and also identifies its holder
// as a client: what an etcd member presents to the others, both ways, and
// to the clients that must show one.
func (ca *CA) IssuePeer(commonName string, hosts []string) (Leaf, error) {
	if len(hosts) == 0 {
		return Leaf{}, fmt.Errorf("a peer certificate needs at least one host")
	}
	return ca.issue(commonName, hosts, x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth)
}

func (ca *CA) issue(commonName string, hosts []string, usage ...x509.ExtKeyUsage) (Leaf, error) {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return Leaf{}, err
	}
	serial, err := randomSerial()
	if err != nil {
		return Leaf{}, err
	}

	tmpl := &x509.Certificate{
		SerialNumber: serial,
		Subject:      pkix.Name{CommonName: commonName},
		NotBefore:    time.Now().Add(-Backdate),
		NotAfter:     time.Now().Add(LeafValidity),
		KeyUsage:     x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:  usage,
	}
	for _, h := range hosts {
		if ip := net.ParseIP(h); ip != nil {
			tmpl.IPAddresses = append(tmpl.IPAddresses, ip)
		} else {
			tmpl.DNSNames = append(tmpl.DNSNames, h)
		}
	}

	der, err := x509.CreateCertificate(rand.Reader, tmpl, ca.Cert, &key.PublicKey, ca.Key)
	if err != nil {
		return Leaf{}, err
	}
	keyDER, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return Leaf{}, err
	}
	return Leaf{
		CertPEM: pemEncode("CERTIFICATE", der),
		KeyPEM:  pemEncode("EC PRIVATE KEY", keyDER),
	}, nil
}

// WriteCA writes the authority to a directory: the certificate world-readable,
// the key not.
func (ca *CA) Write(dir string) error {
	keyDER, err := x509.MarshalECPrivateKey(ca.Key)
	if err != nil {
		return err
	}
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	if err := os.WriteFile(filepath.Join(dir, "ca.pem"), ca.PEM, 0o644); err != nil {
		return err
	}
	// The CA key signs every identity in the cluster. It belongs on whichever
	// machine issues certificates and nowhere else.
	return os.WriteFile(filepath.Join(dir, "ca-key.pem"), pemEncode("EC PRIVATE KEY", keyDER), 0o600)
}

// Write saves a leaf under dir as <name>.pem and <name>-key.pem.
func (l Leaf) Write(dir, name string) error {
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	if err := os.WriteFile(filepath.Join(dir, name+".pem"), l.CertPEM, 0o644); err != nil {
		return err
	}
	return os.WriteFile(filepath.Join(dir, name+"-key.pem"), l.KeyPEM, 0o600)
}

func pemEncode(kind string, der []byte) []byte {
	return pem.EncodeToMemory(&pem.Block{Type: kind, Bytes: der})
}

func randomSerial() (*big.Int, error) {
	// 128 bits, as the CA/Browser Forum requires and every sane issuer does:
	// a predictable serial is one of the ingredients of a chosen-prefix attack.
	return rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
}