package tlsreload

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

func TestTrustProviderReloadsRootsAndKeepsLastKnownGoodBundle(t *testing.T) {
	directory := t.TempDir()
	trustPath := filepath.Join(directory, "ca.crt")
	firstAuthority := newTestAuthority(t, "first")
	firstServer := firstAuthority.server(t, "controller")
	writeTrustBundle(t, trustPath, firstAuthority.certificate)

	provider, err := NewTrustProvider(trustPath)
	if err != nil {
		t.Fatalf("NewTrustProvider() error = %v", err)
	}
	verify := provider.VerifyServer("controller")
	if err := verify(tls.ConnectionState{PeerCertificates: []*x509.Certificate{firstServer}}); err != nil {
		t.Fatalf("verify first server with first root: %v", err)
	}

	secondAuthority := newTestAuthority(t, "second")
	secondServer := secondAuthority.server(t, "controller")
	writeTrustBundle(t, trustPath, secondAuthority.certificate)
	if err := verify(tls.ConnectionState{PeerCertificates: []*x509.Certificate{firstServer}}); err == nil {
		t.Fatal("verify old server after trust rotation error = nil")
	}
	if err := verify(tls.ConnectionState{PeerCertificates: []*x509.Certificate{secondServer}}); err != nil {
		t.Fatalf("verify second server with second root: %v", err)
	}

	if err := os.WriteFile(trustPath, []byte("incomplete"), 0o600); err != nil {
		t.Fatalf("break trust bundle: %v", err)
	}
	if err := verify(tls.ConnectionState{PeerCertificates: []*x509.Certificate{secondServer}}); err != nil {
		t.Fatalf("verify with last known good trust bundle: %v", err)
	}
}

func TestTrustProviderRejectsEmptyPath(t *testing.T) {
	if _, err := NewTrustProvider(""); err != ErrInvalidTrustPath {
		t.Fatalf("NewTrustProvider(empty path) error = %v, want %v", err, ErrInvalidTrustPath)
	}
}

type testAuthority struct {
	certificate *x509.Certificate
	privateKey  *ecdsa.PrivateKey
}

func newTestAuthority(t *testing.T, commonName string) testAuthority {
	t.Helper()
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("generate authority key: %v", err)
	}
	now := time.Now()
	template := &x509.Certificate{
		SerialNumber: big.NewInt(now.UnixNano()), Subject: pkix.Name{CommonName: commonName},
		NotBefore: now.Add(-time.Minute), NotAfter: now.Add(time.Hour),
		IsCA: true, BasicConstraintsValid: true,
		KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature,
	}
	der, err := x509.CreateCertificate(rand.Reader, template, template, &privateKey.PublicKey, privateKey)
	if err != nil {
		t.Fatalf("create authority certificate: %v", err)
	}
	certificate, err := x509.ParseCertificate(der)
	if err != nil {
		t.Fatalf("parse authority certificate: %v", err)
	}
	return testAuthority{certificate: certificate, privateKey: privateKey}
}

func (authority testAuthority) server(t *testing.T, serverName string) *x509.Certificate {
	t.Helper()
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("generate server key: %v", err)
	}
	now := time.Now()
	template := &x509.Certificate{
		SerialNumber: big.NewInt(now.UnixNano()), Subject: pkix.Name{CommonName: serverName},
		DNSNames: []string{serverName}, NotBefore: now.Add(-time.Minute), NotAfter: now.Add(time.Hour),
		KeyUsage:    x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}
	der, err := x509.CreateCertificate(rand.Reader, template, authority.certificate, &privateKey.PublicKey, authority.privateKey)
	if err != nil {
		t.Fatalf("create server certificate: %v", err)
	}
	certificate, err := x509.ParseCertificate(der)
	if err != nil {
		t.Fatalf("parse server certificate: %v", err)
	}
	return certificate
}

func writeTrustBundle(t *testing.T, path string, certificate *x509.Certificate) {
	t.Helper()
	if certificate == nil {
		t.Fatal("write trust bundle: certificate is nil")
	}
	payload := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certificate.Raw})
	if err := os.WriteFile(path, payload, 0o600); err != nil {
		t.Fatalf("write trust bundle: %v", err)
	}
}