package firewall_test

import (
	"crypto/x509"
	"encoding/pem"
	"os"
	"path/filepath"
	"testing"

	"github.com/schmitthub/clawker/internal/config"
	"github.com/schmitthub/clawker/internal/controlplane/firewall"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestEnsureCA_CreatesNew(t *testing.T) {
	certDir := t.TempDir()

	caCert, caKey, err := firewall.EnsureCA(certDir)
	require.NoError(t, err)

	assert.True(t, caCert.IsCA)
	assert.Equal(t, "Clawker Firewall CA", caCert.Subject.CommonName)
	assert.Equal(t, x509.KeyUsageCertSign|x509.KeyUsageCRLSign, caCert.KeyUsage)
	assert.NotNil(t, caKey)

	assert.FileExists(t, filepath.Join(certDir, "ca-cert.pem"))
	assert.FileExists(t, filepath.Join(certDir, "ca-key.pem"))
}

func TestEnsureCA_LoadsExisting(t *testing.T) {
	certDir := t.TempDir()

	cert1, key1, err := firewall.EnsureCA(certDir)
	require.NoError(t, err)

	cert2, key2, err := firewall.EnsureCA(certDir)
	require.NoError(t, err)

	// Same CA loaded — serial numbers and keys must match.
	assert.Equal(t, cert1.SerialNumber, cert2.SerialNumber)
	assert.True(t, key1.Equal(key2))
}

func TestGenerateDomainCert_Valid(t *testing.T) {
	certDir := t.TempDir()
	caCert, caKey, err := firewall.EnsureCA(certDir)
	require.NoError(t, err)

	certPEM, keyPEM, err := firewall.GenerateDomainCert(caCert, caKey, "github.com")
	require.NoError(t, err)
	require.NotEmpty(t, certPEM)
	require.NotEmpty(t, keyPEM)

	// Parse the domain cert.
	block, _ := pem.Decode(certPEM)
	require.NotNil(t, block)
	domainCert, err := x509.ParseCertificate(block.Bytes)
	require.NoError(t, err)

	assert.Equal(t, "github.com", domainCert.Subject.CommonName)
	assert.Contains(t, domainCert.DNSNames, "github.com")
	assert.False(t, domainCert.IsCA)
	assert.Contains(t, domainCert.ExtKeyUsage, x509.ExtKeyUsageServerAuth)

	// Verify the cert chain: domain cert signed by our CA.
	roots := x509.NewCertPool()
	roots.AddCert(caCert)
	_, err = domainCert.Verify(x509.VerifyOptions{
		Roots:     roots,
		KeyUsages: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	})
	require.NoError(t, err, "domain cert should verify against CA")

	// Verify key PEM is a valid EC private key.
	keyBlock, _ := pem.Decode(keyPEM)
	require.NotNil(t, keyBlock)
	_, err = x509.ParseECPrivateKey(keyBlock.Bytes)
	require.NoError(t, err, "key PEM should be a valid EC private key")
}

func TestRegenerateDomainCerts_AllTLSRules(t *testing.T) {
	certDir := t.TempDir()
	caCert, caKey, err := firewall.EnsureCA(certDir)
	require.NoError(t, err)

	rules := []config.EgressRule{
		{
			Dst:   "github.com",
			Proto: "tls",
			// No PathRules — still gets a cert for TLS inspection.
		},
		{
			Dst:   "api.openai.com",
			Proto: "tls",
			PathRules: []config.PathRule{
				{Path: "/v1/models", Action: "allow"},
			},
			PathDefault: "deny",
		},
		{
			Dst:   "storage.googleapis.com",
			Proto: "tls",
			PathRules: []config.PathRule{
				{Path: "/download/*", Action: "allow"},
			},
			PathDefault: "deny",
		},
		{
			Dst:   "git.example.com",
			Proto: "ssh",
			Port:  22,
			// SSH rules do NOT get certs.
		},
	}

	err = firewall.RegenerateDomainCerts(rules, certDir, caCert, caKey)
	require.NoError(t, err)

	// All TLS rules get certs — regardless of PathRules.
	assert.FileExists(t, filepath.Join(certDir, "github.com-cert.pem"))
	assert.FileExists(t, filepath.Join(certDir, "github.com-key.pem"))
	assert.FileExists(t, filepath.Join(certDir, "api.openai.com-cert.pem"))
	assert.FileExists(t, filepath.Join(certDir, "api.openai.com-key.pem"))
	assert.FileExists(t, filepath.Join(certDir, "storage.googleapis.com-cert.pem"))
	assert.FileExists(t, filepath.Join(certDir, "storage.googleapis.com-key.pem"))

	// SSH rules do NOT get certs.
	assert.NoFileExists(t, filepath.Join(certDir, "git.example.com-cert.pem"))
	assert.NoFileExists(t, filepath.Join(certDir, "git.example.com-key.pem"))

	// Verify one of the generated certs is valid and CA-signed.
	certPEM, err := os.ReadFile(filepath.Join(certDir, "api.openai.com-cert.pem"))
	require.NoError(t, err)
	block, _ := pem.Decode(certPEM)
	require.NotNil(t, block)
	cert, err := x509.ParseCertificate(block.Bytes)
	require.NoError(t, err)
	assert.Equal(t, "api.openai.com", cert.Subject.CommonName)

	roots := x509.NewCertPool()
	roots.AddCert(caCert)
	_, err = cert.Verify(x509.VerifyOptions{
		Roots:     roots,
		KeyUsages: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	})
	require.NoError(t, err)
}

func TestRotateCA_RegeneratesAll(t *testing.T) {
	certDir := t.TempDir()

	// Create initial CA and domain certs.
	oldCACert, oldCAKey, err := firewall.EnsureCA(certDir)
	require.NoError(t, err)

	rules := []config.EgressRule{
		{
			Dst:   "api.openai.com",
			Proto: "tls",
			PathRules: []config.PathRule{
				{Path: "/v1/models", Action: "allow"},
			},
			PathDefault: "deny",
		},
	}
	err = firewall.RegenerateDomainCerts(rules, certDir, oldCACert, oldCAKey)
	require.NoError(t, err)

	// Read old domain cert for comparison.
	oldDomainCertPEM, err := os.ReadFile(filepath.Join(certDir, "api.openai.com-cert.pem"))
	require.NoError(t, err)

	// Rotate.
	err = firewall.RotateCA(certDir, rules)
	require.NoError(t, err)

	// New CA should exist and be different.
	newCACert, _, err := firewall.EnsureCA(certDir)
	require.NoError(t, err)
	assert.NotEqual(t, oldCACert.SerialNumber, newCACert.SerialNumber, "CA serial should change after rotation")

	// New domain cert should exist and be different.
	newDomainCertPEM, err := os.ReadFile(filepath.Join(certDir, "api.openai.com-cert.pem"))
	require.NoError(t, err)
	assert.NotEqual(t, oldDomainCertPEM, newDomainCertPEM, "domain cert should be regenerated")

	// New domain cert should verify against the NEW CA.
	block, _ := pem.Decode(newDomainCertPEM)
	require.NotNil(t, block)
	newDomainCert, err := x509.ParseCertificate(block.Bytes)
	require.NoError(t, err)

	roots := x509.NewCertPool()
	roots.AddCert(newCACert)
	_, err = newDomainCert.Verify(x509.VerifyOptions{
		Roots:     roots,
		KeyUsages: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	})
	require.NoError(t, err, "new domain cert should verify against new CA")

	// Old CA should NOT verify new domain cert.
	oldRoots := x509.NewCertPool()
	oldRoots.AddCert(oldCACert)
	_, err = newDomainCert.Verify(x509.VerifyOptions{
		Roots:     oldRoots,
		KeyUsages: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	})
	assert.Error(t, err, "new domain cert should NOT verify against old CA")
}

func TestGenerateDomainCert_WildcardSANs(t *testing.T) {
	certDir := t.TempDir()
	caCert, caKey, err := firewall.EnsureCA(certDir)
	require.NoError(t, err)

	certPEM, _, err := firewall.GenerateDomainCert(caCert, caKey, ".datadoghq.com")
	require.NoError(t, err)

	block, _ := pem.Decode(certPEM)
	require.NotNil(t, block)
	cert, err := x509.ParseCertificate(block.Bytes)
	require.NoError(t, err)

	// CN should be the normalized domain (no leading dot).
	assert.Equal(t, "datadoghq.com", cert.Subject.CommonName)
	// SANs should include both apex and wildcard.
	assert.Contains(t, cert.DNSNames, "datadoghq.com")
	assert.Contains(t, cert.DNSNames, "*.datadoghq.com")
}

func TestGenerateDomainCert_ExactDomainNoWildcardSAN(t *testing.T) {
	certDir := t.TempDir()
	caCert, caKey, err := firewall.EnsureCA(certDir)
	require.NoError(t, err)

	certPEM, _, err := firewall.GenerateDomainCert(caCert, caKey, "api.openai.com")
	require.NoError(t, err)

	block, _ := pem.Decode(certPEM)
	require.NotNil(t, block)
	cert, err := x509.ParseCertificate(block.Bytes)
	require.NoError(t, err)

	assert.Equal(t, "api.openai.com", cert.Subject.CommonName)
	assert.Equal(t, []string{"api.openai.com"}, cert.DNSNames, "exact domain should not get wildcard SAN")
}

func TestRegenerateDomainCerts_WildcardAndExactDedup(t *testing.T) {
	// Both ".example.com" (wildcard) and "example.com" (exact) should produce
	// exactly one cert file with wildcard SANs, regardless of input order.
	tests := []struct {
		name  string
		rules []config.EgressRule
	}{
		{
			name: "wildcard first",
			rules: []config.EgressRule{
				{Dst: ".claude.ai", Proto: "tls"},
				{Dst: "claude.ai", Proto: "tls"},
			},
		},
		{
			name: "exact first",
			rules: []config.EgressRule{
				{Dst: "claude.ai", Proto: "tls"},
				{Dst: ".claude.ai", Proto: "tls"},
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			certDir := t.TempDir()
			caCert, caKey, err := firewall.EnsureCA(certDir)
			require.NoError(t, err)

			err = firewall.RegenerateDomainCerts(tc.rules, certDir, caCert, caKey)
			require.NoError(t, err)

			// Only one cert file pair — normalized to "claude.ai".
			certPath := filepath.Join(certDir, "claude.ai-cert.pem")
			assert.FileExists(t, certPath)
			assert.FileExists(t, filepath.Join(certDir, "claude.ai-key.pem"))
			assert.NoFileExists(t, filepath.Join(certDir, ".claude.ai-cert.pem"))

			// Cert must have both apex and wildcard SANs.
			certPEM, err := os.ReadFile(certPath)
			require.NoError(t, err)
			block, _ := pem.Decode(certPEM)
			require.NotNil(t, block)
			cert, err := x509.ParseCertificate(block.Bytes)
			require.NoError(t, err)

			assert.Equal(t, "claude.ai", cert.Subject.CommonName)
			assert.Contains(t, cert.DNSNames, "claude.ai", "apex SAN required")
			assert.Contains(t, cert.DNSNames, "*.claude.ai", "wildcard SAN required")

			// Count cert files — should be exactly 1 pair (+ CA pair).
			entries, err := os.ReadDir(certDir)
			require.NoError(t, err)
			certFiles := 0
			for _, e := range entries {
				if !e.IsDir() && filepath.Ext(e.Name()) == ".pem" {
					certFiles++
				}
			}
			// 2 CA files + 2 domain files = 4.
			assert.Equal(t, 4, certFiles, "should have exactly 1 domain cert pair + CA pair")
		})
	}
}

func TestRegenerateDomainCerts_WildcardFilenames(t *testing.T) {
	certDir := t.TempDir()
	caCert, caKey, err := firewall.EnsureCA(certDir)
	require.NoError(t, err)

	rules := []config.EgressRule{
		{
			Dst:   ".datadoghq.com",
			Proto: "tls",
			PathRules: []config.PathRule{
				{Path: "/api/v2/logs", Action: "allow"},
			},
			PathDefault: "deny",
		},
	}

	err = firewall.RegenerateDomainCerts(rules, certDir, caCert, caKey)
	require.NoError(t, err)

	// Files should use the normalized domain (no leading dot).
	assert.FileExists(t, filepath.Join(certDir, "datadoghq.com-cert.pem"))
	assert.FileExists(t, filepath.Join(certDir, "datadoghq.com-key.pem"))
	// No hidden files with leading dot.
	assert.NoFileExists(t, filepath.Join(certDir, ".datadoghq.com-cert.pem"))
	assert.NoFileExists(t, filepath.Join(certDir, ".datadoghq.com-key.pem"))

	// Verify the cert has wildcard SANs.
	certPEM, err := os.ReadFile(filepath.Join(certDir, "datadoghq.com-cert.pem"))
	require.NoError(t, err)
	block, _ := pem.Decode(certPEM)
	require.NotNil(t, block)
	cert, err := x509.ParseCertificate(block.Bytes)
	require.NoError(t, err)
	assert.Contains(t, cert.DNSNames, "datadoghq.com")
	assert.Contains(t, cert.DNSNames, "*.datadoghq.com")
}

func TestRegenerateDomainCerts_CleansStaleCerts(t *testing.T) {
	certDir := t.TempDir()
	caCert, caKey, err := firewall.EnsureCA(certDir)
	require.NoError(t, err)

	// Generate certs for two domains.
	rules := []config.EgressRule{
		{Dst: "old-domain.com", Proto: "tls"},
		{Dst: "kept-domain.com", Proto: "tls"},
	}
	err = firewall.RegenerateDomainCerts(rules, certDir, caCert, caKey)
	require.NoError(t, err)
	assert.FileExists(t, filepath.Join(certDir, "old-domain.com-cert.pem"))
	assert.FileExists(t, filepath.Join(certDir, "kept-domain.com-cert.pem"))

	// Regenerate with only the kept domain — stale cert should be cleaned.
	rules = []config.EgressRule{
		{Dst: "kept-domain.com", Proto: "tls"},
	}
	err = firewall.RegenerateDomainCerts(rules, certDir, caCert, caKey)
	require.NoError(t, err)
	assert.NoFileExists(t, filepath.Join(certDir, "old-domain.com-cert.pem"))
	assert.NoFileExists(t, filepath.Join(certDir, "old-domain.com-key.pem"))
	assert.FileExists(t, filepath.Join(certDir, "kept-domain.com-cert.pem"))
	// CA files preserved.
	assert.FileExists(t, filepath.Join(certDir, "ca-cert.pem"))
	assert.FileExists(t, filepath.Join(certDir, "ca-key.pem"))
}
