package blockpage

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"fmt"
	"math/big"
	"net"
	"sync"
	"time"
)

// certCap bounds how many per-hostname certificates are kept in memory. A
// client that churns through SNI names (deliberately or not) could otherwise
// grow the cache without limit; when the cap is hit the whole cache is
// cleared and rebuilt on demand, which just costs a few CPU-cheap regenerations.
const certCap = 20_000

// certStore lazily generates and caches a self-signed TLS certificate for
// each hostname a client connects with over SNI. Signing each leaf with a
// name that matches what was requested means the browser's warning is the
// expected "this certificate is not trusted" one, not also a confusing
// hostname mismatch.
type certStore struct {
	mu    sync.Mutex
	certs map[string]*tls.Certificate
}

func newCertStore() certStore {
	return certStore{certs: make(map[string]*tls.Certificate)}
}

// get implements tls.Config.GetCertificate.
func (c *certStore) get(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
	name := hello.ServerName
	if name == "" {
		name = "block-page.invalid"
	}

	c.mu.Lock()
	if cert, ok := c.certs[name]; ok {
		c.mu.Unlock()
		return cert, nil
	}
	c.mu.Unlock()

	cert, err := generateCert(name)
	if err != nil {
		return nil, err
	}

	c.mu.Lock()
	if len(c.certs) >= certCap {
		c.certs = make(map[string]*tls.Certificate)
	}
	c.certs[name] = cert
	c.mu.Unlock()
	return cert, nil
}

// generateCert creates a fresh, self-signed ECDSA leaf certificate for name.
func generateCert(name string) (*tls.Certificate, error) {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("generate block page certificate key: %w", err)
	}
	serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 62))
	if err != nil {
		return nil, fmt.Errorf("generate block page certificate serial: %w", err)
	}

	tmpl := &x509.Certificate{
		SerialNumber:          serial,
		Subject:               pkix.Name{CommonName: name, Organization: []string{"VibeDNS Block Page"}},
		NotBefore:             time.Now().Add(-time.Hour),
		NotAfter:              time.Now().AddDate(10, 0, 0),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		IsCA:                  true,
	}
	if ip := net.ParseIP(name); ip != nil {
		tmpl.IPAddresses = []net.IP{ip}
	} else {
		tmpl.DNSNames = []string{name}
	}

	der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
	if err != nil {
		return nil, fmt.Errorf("create block page certificate: %w", err)
	}
	return &tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key}, nil
}

// parseLeaf parses the leaf certificate's DER bytes, for tests that need to
// inspect what generateCert produced.
func parseLeaf(cert *tls.Certificate) (*x509.Certificate, error) {
	return x509.ParseCertificate(cert.Certificate[0])
}