package cql

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"math/big"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"testing"
	"time"

	"go.uber.org/zap"
)

type noopHandler struct{}

func (noopHandler) HandleMessage(ctx context.Context, frame *Frame) (*Frame, error) {
	return nil, nil
}

func TestServerTLSEnabledConnectionInfo(t *testing.T) {
	certFile, keyFile := writeTestTLSCert(t)

	server := NewServer("127.0.0.1:0", noopHandler{}, zap.NewNop(), &ServerOptions{
		ProtocolVersion: ProtoVersion4,
		TLS: &TLSConfig{
			Mode:     "manual",
			CertFile: certFile,
			KeyFile:  keyFile,
		},
	})
	if err := server.Start(); err != nil {
		t.Fatalf("server start failed: %v", err)
	}
	defer server.Stop()

	addr := server.Addr()
	if addr == nil {
		t.Fatalf("server addr is nil")
	}
	tcpAddr, ok := addr.(*net.TCPAddr)
	if !ok || tcpAddr == nil {
		t.Fatalf("unexpected addr type %T", addr)
	}
	client, err := tls.Dial("tcp", net.JoinHostPort("127.0.0.1", strconv.Itoa(tcpAddr.Port)), &tls.Config{
		InsecureSkipVerify: true,
	})
	if err != nil {
		t.Fatalf("tls dial failed: %v", err)
	}
	defer client.Close()

	info := waitForConnectionInfo(t, server, 2*time.Second)
	if !info.SSL {
		t.Fatalf("expected SSL enabled")
	}
	if info.SSLProtocol == "" {
		t.Fatalf("expected SSL protocol")
	}
	if info.SSLCipherSuite == "" {
		t.Fatalf("expected SSL cipher suite")
	}
}

func waitForConnectionInfo(t *testing.T, server *Server, timeout time.Duration) ConnectionInfo {
	t.Helper()
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		infos := server.GetActiveConnections()
		if len(infos) > 0 {
			return infos[0]
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for active connection")
	return ConnectionInfo{}
}

func writeTestTLSCert(t *testing.T) (string, string) {
	t.Helper()
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("generate key: %v", err)
	}
	serial, err := rand.Int(rand.Reader, big.NewInt(1<<62))
	if err != nil {
		t.Fatalf("generate serial: %v", err)
	}
	now := time.Now()
	template := x509.Certificate{
		SerialNumber: serial,
		NotBefore:    now.Add(-time.Hour),
		NotAfter:     now.Add(time.Hour),
		KeyUsage:     x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		DNSNames:     []string{"localhost"},
		IPAddresses:  []net.IP{net.ParseIP("127.0.0.1")},
	}
	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		t.Fatalf("create cert: %v", err)
	}
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	keyBytes, err := x509.MarshalECPrivateKey(priv)
	if err != nil {
		t.Fatalf("marshal key: %v", err)
	}
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyBytes})

	dir := t.TempDir()
	certFile := filepath.Join(dir, "server-cert.pem")
	keyFile := filepath.Join(dir, "server-key.pem")
	if err := os.WriteFile(certFile, certPEM, 0600); err != nil {
		t.Fatalf("write cert: %v", err)
	}
	if err := os.WriteFile(keyFile, keyPEM, 0600); err != nil {
		t.Fatalf("write key: %v", err)
	}
	return certFile, keyFile
}