package ca

import (
	"crypto"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"fmt"
	"math/big"
	"net"
	"strings"
	"sync"
	"time"

	"git.maze.io/maze/styx/internal/cryptutil"
	"git.maze.io/maze/styx/logger"
	"github.com/miekg/dns"
)

type CertificateAuthority interface {
	GetCertificate(commonName string, dnsNames []string, ips []net.IP) (*tls.Certificate, error)
}

type ca struct {
	cert  *x509.Certificate
	key   crypto.PrivateKey
	cache sync.Map
}

func Open(certData, keyData string) (CertificateAuthority, error) {
	cert, key, err := cryptutil.LoadKeyPair(certData, keyData)
	if err != nil {
		return nil, err
	} else if !cert.IsCA {
		return nil, fmt.Errorf("ca: certificate for %s is not a certificate authority", cert.Subject.String())
	}

	return &ca{
		cert: cert,
		key:  key,
	}, nil
}

func (ca *ca) GetCertificate(cn string, names []string, ips []net.IP) (*tls.Certificate, error) {
	var (
		log = logger.StandardLog.Values(logger.Values{
			"cn":    cn,
			"names": names,
			"ips":   ips,
		})
		now    = time.Now().UTC()
		parent = parentDomain(cn)
	)
	if cn == parent {
		names = append(names, "*."+cn)
	} else {
		names = append(names, "*."+parent, cn)
		cn = parent
		log = log.Value("cn", cn)
	}
	if v, ok := ca.cache.Load(parent); ok {
		if cert, ok := v.(*tls.Certificate); ok && now.After(cert.Leaf.NotBefore) && now.Before(cert.Leaf.NotAfter.Add(-time.Hour)) {
			log.Value("valid", cert.Leaf.NotAfter.Sub(now)).Debug("Using cached certificate")
			return cert, nil
		}
		log.Debug("Cached certificate invalid")
		ca.cache.Delete(parent)
	}

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return nil, fmt.Errorf("ca: failed to generate serial number: %w", err)
	}

	notBefore := now.Round(24 * time.Hour)
	notAfter := notBefore.Add(48 * time.Hour)

	log.Values(logger.Values{
		"serial":  serialNumber.String(),
		"subject": pkix.Name{CommonName: cn}.String(),
	}).Debug("Generating certificate")
	template := &x509.Certificate{
		SerialNumber: serialNumber,
		KeyUsage:     x509.KeyUsageDataEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		Subject:      pkix.Name{CommonName: cn},
		DNSNames:     names,
		IPAddresses:  ips,
		PublicKey:    cryptutil.PublicKey(ca.key),
		NotBefore:    notBefore,
		NotAfter:     notAfter,
	}
	der, err := x509.CreateCertificate(rand.Reader, template, ca.cert, template.PublicKey, ca.key)
	if err != nil {
		return nil, err
	}
	cert, err := x509.ParseCertificate(der)
	if err != nil {
		return nil, err
	}

	output := &tls.Certificate{
		Certificate: [][]byte{der},
		Leaf:        cert,
		PrivateKey:  ca.key,
	}
	ca.cache.Store(parent, output)
	return output, nil
}

func parentDomain(name string) string {
	part := dns.SplitDomainName(name)
	if len(part) <= 2 {
		return name
	}
	return strings.Join(part[1:], ".")
}