package protocol

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"net"
	"sync"
	"testing"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/extension/providers/mail/transport"
)

type staticConnector struct {
	mu         sync.Mutex
	connection net.Conn
	tlsConfig  *tls.Config
	calls      int
}

func (connector *staticConnector) ResolveAndDial(
	_ context.Context,
	target transport.Target,
) (transport.PinnedConnection, error) {
	connector.mu.Lock()
	defer connector.mu.Unlock()
	connector.calls++
	if connector.connection == nil {
		return transport.PinnedConnection{},
			fmt.Errorf("scripted connection was already consumed")
	}
	connection := connector.connection
	connector.connection = nil
	return transport.PinnedConnection{
		Conn:      connection,
		Host:      target.Host,
		Address:   "scripted:1",
		TLSConfig: connector.tlsConfig.Clone(),
		TLSMode:   target.TLSMode,
	}, nil
}

func newScriptedConnector(
	t *testing.T,
) (*staticConnector, net.Conn, tls.Certificate) {
	t.Helper()
	certificate, roots := newTestCertificate(t, "mail.example.com")
	client, server := net.Pipe()
	return &staticConnector{
		connection: client,
		tlsConfig: &tls.Config{
			MinVersion: tls.VersionTLS12,
			ServerName: "mail.example.com",
			RootCAs:    roots,
		},
	}, server, certificate
}

func newTestCertificate(
	t *testing.T,
	host string,
) (tls.Certificate, *x509.CertPool) {
	t.Helper()
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 120))
	if err != nil {
		t.Fatal(err)
	}
	template := &x509.Certificate{
		SerialNumber: serial,
		Subject:      pkix.Name{CommonName: host},
		DNSNames:     []string{host},
		NotBefore:    time.Now().Add(-time.Minute),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}
	der, err := x509.CreateCertificate(
		rand.Reader,
		template,
		template,
		&privateKey.PublicKey,
		privateKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	keyDER, err := x509.MarshalPKCS8PrivateKey(privateKey)
	if err != nil {
		t.Fatal(err)
	}
	certificate, err := tls.X509KeyPair(
		pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}),
		pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: keyDER}),
	)
	if err != nil {
		t.Fatal(err)
	}
	parsedCertificate, err := x509.ParseCertificate(der)
	if err != nil {
		t.Fatal(err)
	}
	roots := x509.NewCertPool()
	roots.AddCert(parsedCertificate)
	return certificate, roots
}