package controlplane

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/hex"
	"encoding/pem"
	"math/big"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func TestStoreLocksDataAndKeepsBoundedBackups(t *testing.T) {
	directory := t.TempDir()
	path := filepath.Join(directory, "state.json")
	store, err := OpenStore(path, StoreOptions{BackupDir: filepath.Join(directory, "backups"), BackupCount: 2})
	if err != nil {
		t.Fatal(err)
	}
	if err := store.Mutate(func(state *PersistedState) error {
		state.Nodes["node-1"] = Node{NodeID: "node-1"}
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	if err := store.Mutate(func(state *PersistedState) error {
		state.Nodes["node-1"] = Node{NodeID: "node-1", AgentVersion: "2"}
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	if err := store.Mutate(func(state *PersistedState) error {
		state.Nodes["node-1"] = Node{NodeID: "node-1", AgentVersion: "3"}
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	if _, err := OpenStore(path, StoreOptions{BackupDir: filepath.Join(directory, "backups"), BackupCount: 2}); err == nil || !strings.Contains(err.Error(), "already in use") {
		t.Fatalf("second writer was not rejected: %v", err)
	}
	if err := store.Close(); err != nil {
		t.Fatal(err)
	}
	entries, err := os.ReadDir(filepath.Join(directory, "backups"))
	if err != nil {
		t.Fatal(err)
	}
	if len(entries) != 2 {
		t.Fatalf("expected two backups, got %d", len(entries))
	}
	for _, entry := range entries {
		if entry.Type().Perm() != 0o600 {
			info, infoErr := entry.Info()
			if infoErr != nil {
				t.Fatal(infoErr)
			}
			if info.Mode().Perm() != 0o600 {
				t.Fatalf("backup permissions = %o", info.Mode().Perm())
			}
		}
	}
	second, err := OpenStore(path, StoreOptions{BackupDir: filepath.Join(directory, "backups"), BackupCount: 2})
	if err != nil {
		t.Fatal(err)
	}
	defer second.Close()
}

func TestStorePrunesOnlyTerminalRecordsByRetention(t *testing.T) {
	directory := t.TempDir()
	path := filepath.Join(directory, "state.json")
	seed, err := OpenStore(path, StoreOptions{BackupDir: "", BackupCount: -1})
	if err != nil {
		t.Fatal(err)
	}
	old := time.Now().UTC().Add(-2 * time.Hour)
	if err := seed.Mutate(func(state *PersistedState) error {
		state.Tasks["old-terminal"] = Task{TaskID: "old-terminal", Status: TaskSucceeded, UpdatedAt: old}
		state.Tasks["old-running"] = Task{TaskID: "old-running", Status: TaskRunning, UpdatedAt: old}
		state.Events = append(state.Events, StoredEvent{ReceivedAt: old})
		state.Audit = append(state.Audit, AuditEntry{At: old})
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	if err := seed.Close(); err != nil {
		t.Fatal(err)
	}

	store, err := OpenStore(path, StoreOptions{
		BackupDir: filepath.Join(directory, "backups"), BackupCount: -1,
		TaskRetention: time.Hour, EventRetention: time.Hour, AuditRetention: time.Hour,
	})
	if err != nil {
		t.Fatal(err)
	}
	defer store.Close()
	if err := store.Mutate(func(state *PersistedState) error {
		state.Nodes["node-1"] = Node{NodeID: "node-1"}
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	snapshot := store.Snapshot()
	if _, ok := snapshot.Tasks["old-terminal"]; ok {
		t.Fatal("old terminal task was retained")
	}
	if _, ok := snapshot.Tasks["old-running"]; !ok {
		t.Fatal("old running task was pruned")
	}
	if len(snapshot.Events) != 0 || len(snapshot.Audit) != 0 {
		t.Fatalf("old events/audit were retained: events=%d audit=%d", len(snapshot.Events), len(snapshot.Audit))
	}
}

func TestTLSConfigSupportsVerifiedNodeClientCertificates(t *testing.T) {
	directory := t.TempDir()
	certFile, keyFile, caFile := writeTestCertificateMaterial(t, directory)
	config := ServerConfig{
		TLSCertFile:          certFile,
		TLSKeyFile:           keyFile,
		MTLSClientCAFile:     caFile,
		MTLSRequireNodeCert:  true,
		MTLSRevokedCertsFile: filepath.Join(directory, "revoked.txt"),
	}
	if err := os.WriteFile(config.MTLSRevokedCertsFile, []byte("# initially empty\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	tlsConfig, err := newTLSConfig(config)
	if err != nil {
		t.Fatal(err)
	}
	if tlsConfig.MinVersion != tls.VersionTLS13 || tlsConfig.ClientAuth != tls.VerifyClientCertIfGiven || len(tlsConfig.Certificates) != 1 {
		t.Fatalf("unexpected TLS config: min=%d auth=%d certs=%d", tlsConfig.MinVersion, tlsConfig.ClientAuth, len(tlsConfig.Certificates))
	}

	rawCertificate := []byte("client-cert")
	digest := sha256.Sum256(rawCertificate)
	fingerprint := hex.EncodeToString(digest[:])
	if err := os.WriteFile(config.MTLSRevokedCertsFile, []byte(fingerprint+"\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	server := &Server{config: config}
	request := &http.Request{TLS: &tls.ConnectionState{PeerCertificates: []*x509.Certificate{{Raw: rawCertificate}}}}
	if server.clientCertificateAllowed(request) {
		t.Fatal("revoked certificate was accepted")
	}
}

func TestTLSConfigRejectsIncompleteSecuritySettings(t *testing.T) {
	if _, err := newTLSConfig(ServerConfig{TLSCertFile: "cert.pem"}); err == nil {
		t.Fatal("incomplete TLS certificate settings were accepted")
	}
	if _, err := newTLSConfig(ServerConfig{MTLSRequireNodeCert: true}); err == nil {
		t.Fatal("mTLS without TLS material was accepted")
	}
}

func writeTestCertificateMaterial(t *testing.T, directory string) (string, string, string) {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	template := &x509.Certificate{
		SerialNumber:          big.NewInt(1),
		Subject:               pkix.Name{CommonName: "wxagent-test"},
		NotBefore:             time.Now().Add(-time.Minute),
		NotAfter:              time.Now().Add(time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment | x509.KeyUsageCertSign,
		IsCA:                  true,
		BasicConstraintsValid: true,
	}
	der, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
	if err != nil {
		t.Fatal(err)
	}
	certFile := filepath.Join(directory, "server.pem")
	keyFile := filepath.Join(directory, "server-key.pem")
	caFile := filepath.Join(directory, "client-ca.pem")
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der})
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})
	for path, data := range map[string][]byte{certFile: certPEM, keyFile: keyPEM, caFile: certPEM} {
		if err := os.WriteFile(path, data, 0o600); err != nil {
			t.Fatal(err)
		}
	}
	return certFile, keyFile, caFile
}