package acmeserver

import (
	"context"
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"errors"
	"io"
	"log"
	"sync"
	"time"

	"git.autistici.org/ai3/replds"
	"github.com/prometheus/client_golang/prometheus"
)

// certInfo represents what we know about the state of the certificate
// at runtime.
type certInfo struct {
	config        *certConfig
	retryDeadline time.Time

	// Write-only attributes (not part of the renewal logic) to
	// indicate whether we think we have a valid certificate or
	// not. Used for monitoring and debugging.
	valid     bool
	expiresAt time.Time
}

func (i *certInfo) cn() string {
	return i.config.Names[0]
}

func (i *certInfo) setOk(cert *x509.Certificate, renewalDays int) {
	i.retryDeadline = cert.NotAfter.AddDate(0, 0, -renewalDays)
	i.expiresAt = cert.NotAfter
	i.valid = true

	l := prometheus.Labels{"cn": i.cn()}
	certOk.With(l).Set(1)
}

func (i *certInfo) setErr() {
	now := time.Now()
	i.retryDeadline = now.Add(errorRetryTimeout)
	i.valid = false
	if !i.expiresAt.IsZero() {
		i.valid = i.expiresAt.After(now)
	}

	l := prometheus.Labels{"cn": i.cn()}
	certRenewalErrorCount.With(l).Inc()
	certOk.With(l).Set(b2f(i.valid))
}

type certStorage interface {
	GetCert(string) ([][]byte, crypto.Signer, error)
	PutCert(string, [][]byte, crypto.Signer) error
}

// CertGenerator is an interface to something that can generate a new
// certificate, given a list of domains and a private key. The ACME
// type implements this interface.
type CertGenerator interface {
	GetCertificate(context.Context, crypto.Signer, *certConfig) ([][]byte, *x509.Certificate, error)
}

// Manager periodically renews certificates before they expire, and
// responds to http-01 validation requests.
type Manager struct {
	configDirs  []string
	useRSA      bool
	storage     certStorage
	certs       []*certInfo
	certGen     CertGenerator
	renewalDays int

	configCh chan []*certConfig
	doneCh   chan bool
}

// NewManager creates a new Manager with the given configuration.
func NewManager(config *Config, certGen CertGenerator) (*Manager, error) {
	// Validate the configuration.
	if len(config.Dirs) == 0 {
		return nil, errors.New("configuration parameter 'config_dirs' is unset")
	}
	if config.Output.Path == "" {
		return nil, errors.New("'output.path' is unset")
	}

	m := &Manager{
		useRSA:      config.UseRSA,
		configDirs:  config.Dirs,
		doneCh:      make(chan bool),
		configCh:    make(chan []*certConfig, 1),
		certGen:     certGen,
		renewalDays: config.RenewalDays,
	}
	if m.renewalDays <= 0 {
		m.renewalDays = 15
	}

	ds := &dirStorage{root: config.Output.Path}
	if config.Output.ReplDS == nil {
		m.storage = ds
	} else {
		r, err := replds.NewPublicClient(config.Output.ReplDS)
		if err != nil {
			return nil, err
		}
		m.storage = &replStorage{
			dirStorage: ds,
			replClient: r,
		}
	}

	return m, nil
}

// Start the renewal processes. Canceling the provided context will
// cause background processing to stop, interrupting all running
// updates.
func (m *Manager) Start(ctx context.Context) error {
	domains, err := readCertConfigsFromDirs(m.configDirs)
	if err != nil {
		return err
	}
	m.configCh <- domains
	go func() {
		m.loop(ctx)
		close(m.doneCh)
	}()
	return nil
}

// Wait for the Manager to terminate once it has been canceled.
func (m *Manager) Wait() {
	<-m.doneCh
}

// Reload configuration.
func (m *Manager) Reload() {
	domains, err := readCertConfigsFromDirs(m.configDirs)
	if err != nil {
		log.Printf("error reading config: %v", err)
		return
	}
	m.configCh <- domains
	log.Printf("configuration reloaded")
}

var (
	renewalTimeout    = 10 * time.Minute
	errorRetryTimeout = 10 * time.Minute
)

func (m *Manager) updateAllCerts(ctx context.Context, certs []*certInfo) {
	for _, certInfo := range certs {
		if certInfo.retryDeadline.After(time.Now()) {
			continue
		}

		// Abort the loop if our context is canceled.
		select {
		case <-ctx.Done():
			return
		default:
		}

		certRenewalCount.With(prometheus.Labels{"cn": certInfo.cn()}).Inc()
		uctx, cancel := context.WithTimeout(ctx, renewalTimeout)
		leaf, err := m.updateCert(uctx, certInfo)
		cancel()
		if err == nil {
			log.Printf("successfully renewed %s", certInfo.cn())
			certInfo.setOk(leaf, m.renewalDays)
		} else {
			log.Printf("error updating %s: %v", certInfo.cn(), err)
			certInfo.setErr()
		}
	}
}

func (m *Manager) updateCert(ctx context.Context, certInfo *certInfo) (*x509.Certificate, error) {
	// Create a new private key.
	var (
		err error
		key crypto.Signer
	)
	if m.useRSA {
		key, err = rsa.GenerateKey(rand.Reader, 2048)
	} else {
		key, err = ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	}
	if err != nil {
		return nil, err
	}

	der, leaf, err := m.certGen.GetCertificate(ctx, key, certInfo.config)
	if err != nil {
		return nil, err
	}

	if err := m.storage.PutCert(certInfo.cn(), der, key); err != nil {
		return nil, err
	}

	return leaf, nil
}

// Replace the current configuration.
func (m *Manager) loadConfig(certs []*certConfig) []*certInfo {
	var out []*certInfo
	for _, c := range certs {
		cn := c.Names[0]
		certInfo := &certInfo{
			config: c,
		}
		pub, priv, err := m.storage.GetCert(cn)
		if err != nil {
			log.Printf("cert for %s is missing", cn)
		} else {
			// By calling validCert we catch things like subjectAltName changes.
			leaf, err := validCert(c.Names, pub, priv)
			if err == nil {
				log.Printf("cert for %s loaded from storage", cn)
				certInfo.setOk(leaf, m.renewalDays)
			} else {
				log.Printf("cert for %s loaded from storage but parameters have changed", cn)
			}
		}
		out = append(out, certInfo)
	}
	return out
}

// This channel is used by the testing code to trigger an update,
// without having to wait for the timer to tick.
var testUpdateCh = make(chan bool)

func (m *Manager) loop(ctx context.Context) {
	// Updates are long-term jobs, so they should be
	// interruptible. We run updates in a separate goroutine, and
	// cancel them when the configuration is reloaded or on exit.
	var upCancel context.CancelFunc
	var wg sync.WaitGroup

	startUpdate := func(certs []*certInfo) context.CancelFunc {
		// Ensure the previous update has finished.
		wg.Wait()

		upCtx, cancel := context.WithCancel(ctx)
		wg.Add(1)
		go func() {
			m.updateAllCerts(upCtx, certs)
			wg.Done()
		}()
		return cancel
	}

	// Cancel the running update, if any. Called on config
	// updates, when exiting.
	cancelUpdate := func() {
		if upCancel != nil {
			upCancel()
		}
		wg.Wait()
	}
	defer cancelUpdate()

	tick := time.NewTicker(5 * time.Minute)
	defer tick.Stop()
	for {
		select {
		case <-tick.C:
			upCancel = startUpdate(m.certs)
		case <-testUpdateCh:
			upCancel = startUpdate(m.certs)
		case certDomains := <-m.configCh:
			cancelUpdate()
			m.certs = m.loadConfig(certDomains)
		case <-ctx.Done():
			return
		}
	}
}

func concatDER(der [][]byte) []byte {
	// Append DERs to a single []byte buffer and parse the results.
	var n int
	for _, b := range der {
		n += len(b)
	}
	out := make([]byte, n)
	n = 0
	for _, b := range der {
		n += copy(out[n:], b)
	}
	return out
}

func certRequest(key crypto.Signer, domains []string) ([]byte, error) {
	req := &x509.CertificateRequest{
		Subject: pkix.Name{CommonName: domains[0]},
	}
	if len(domains) > 1 {
		req.DNSNames = domains[1:]
	}
	return x509.CreateCertificateRequest(rand.Reader, req, key)
}

func parsePrivateKey(der []byte) (crypto.Signer, error) {
	if key, err := x509.ParsePKCS1PrivateKey(der); err == nil {
		return key, nil
	}
	if key, err := x509.ParsePKCS8PrivateKey(der); err == nil {
		switch key := key.(type) {
		case *rsa.PrivateKey:
			return key, nil
		case *ecdsa.PrivateKey:
			return key, nil
		default:
			return nil, errors.New("unknown private key type in PKCS#8 wrapping")
		}
	}
	if key, err := x509.ParseECPrivateKey(der); err == nil {
		return key, nil
	}

	return nil, errors.New("failed to parse private key")
}

func encodeECDSAKey(w io.Writer, key *ecdsa.PrivateKey) error {
	b, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return err
	}
	pb := &pem.Block{Type: "EC PRIVATE KEY", Bytes: b}
	return pem.Encode(w, pb)
}