package checker

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

	"git.happydns.org/checker-tls/contract"
)

// startEnumTestServer spins up a TCP listener that, for every accepted
// connection: (1) optionally drives a fake STARTTLS dialect handshake, then
// (2) lets the standard library terminate TLS with the provided cert. It
// keeps accepting until the test closes the listener.
//
// We use the stdlib tls.Server (not utls) on the server side: the point of
// these tests is to exercise the *checker* glue (upgraderFor + enumerate)
// against the real client-side code, not to replay tlsenum's internals.
func startEnumTestServer(t *testing.T, withSTARTTLS bool, cert tls.Certificate) net.Listener {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			go handleEnumConn(c, withSTARTTLS, cert)
		}
	}()
	return ln
}

func handleEnumConn(c net.Conn, withSTARTTLS bool, cert tls.Certificate) {
	defer c.Close()
	if withSTARTTLS {
		// Pretend to be SMTP: 220 banner, EHLO ack, STARTTLS ack. The
		// implementation of starttlsSMTP only requires the server to
		// advertise STARTTLS in its EHLO response and to reply with a 2xx
		// to the STARTTLS verb — exact verbs come from RFC 3207.
		if _, err := io.WriteString(c, "220 enum.test ESMTP\r\n"); err != nil {
			return
		}
		buf := make([]byte, 1024)
		// EHLO line
		if _, err := c.Read(buf); err != nil {
			return
		}
		if _, err := io.WriteString(c, "250-enum.test\r\n250 STARTTLS\r\n"); err != nil {
			return
		}
		// STARTTLS line
		if _, err := c.Read(buf); err != nil {
			return
		}
		if _, err := io.WriteString(c, "220 ready\r\n"); err != nil {
			return
		}
	}
	tc := tls.Server(c, &tls.Config{
		Certificates: []tls.Certificate{cert},
		MinVersion:   tls.VersionTLS12,
		MaxVersion:   tls.VersionTLS12, // narrow surface so the sweep is fast
	})
	defer tc.Close()
	_ = tc.Handshake()
}

// enumTestCert is a one-time self-signed ECDSA cert reused across tests.
func enumTestCert(t *testing.T) tls.Certificate {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("genkey: %v", err)
	}
	tmpl := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject:      pkix.Name{CommonName: "enum.test"},
		NotBefore:    time.Now().Add(-time.Hour),
		NotAfter:     time.Now().Add(time.Hour),
		DNSNames:     []string{"enum.test"},
		KeyUsage:     x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}
	der, err := x509.CreateCertificate(rand.Reader, &tmpl, &tmpl, &key.PublicKey, key)
	if err != nil {
		t.Fatalf("createcert: %v", err)
	}
	keyDER, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		t.Fatalf("marshal key: %v", err)
	}
	c, err := tls.X509KeyPair(
		pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}),
		pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER}),
	)
	if err != nil {
		t.Fatalf("keypair: %v", err)
	}
	return c
}

func portOf(t *testing.T, ln net.Listener) uint16 {
	t.Helper()
	_, p, err := net.SplitHostPort(ln.Addr().String())
	if err != nil {
		t.Fatalf("split addr: %v", err)
	}
	n, err := strconv.ParseUint(p, 10, 16)
	if err != nil {
		t.Fatalf("parse port: %v", err)
	}
	return uint16(n)
}

// TestEnumerateEndpoint_DirectTLS asserts the sweep returns at least one
// supported version + cipher when the endpoint is plain TLS — proving the
// nil-upgrader path of upgraderFor wires correctly.
func TestEnumerateEndpoint_DirectTLS(t *testing.T) {
	cert := enumTestCert(t)
	ln := startEnumTestServer(t, false, cert)
	defer ln.Close()

	res, skip := enumerateEndpoint(context.Background(), contract.TLSEndpoint{
		Host: "127.0.0.1",
		Port: portOf(t, ln),
		SNI:  "enum.test",
	}, 30*time.Second)
	if skip != "" {
		t.Fatalf("unexpected skip reason: %q", skip)
	}
	if res == nil || len(res.Versions) == 0 {
		t.Fatalf("expected at least one supported version, got %+v", res)
	}
	gotTLS12 := false
	for _, v := range res.Versions {
		if v.Version == tls.VersionTLS12 && len(v.Ciphers) > 0 {
			gotTLS12 = true
		}
	}
	if !gotTLS12 {
		t.Fatalf("expected TLS 1.2 with at least one cipher, got %+v", res.Versions)
	}
}

// TestEnumerateEndpoint_SMTP_STARTTLS asserts the sweep drives the SMTP
// dialect upgrade on every sub-probe and still discovers ciphers — proving
// the upgraderFor("smtp", sni) path is wired into Enumerate.
func TestEnumerateEndpoint_SMTP_STARTTLS(t *testing.T) {
	cert := enumTestCert(t)
	ln := startEnumTestServer(t, true, cert)
	defer ln.Close()

	res, skip := enumerateEndpoint(context.Background(), contract.TLSEndpoint{
		Host:     "127.0.0.1",
		Port:     portOf(t, ln),
		SNI:      "enum.test",
		STARTTLS: "smtp",
	}, 60*time.Second)
	if skip != "" {
		t.Fatalf("unexpected skip reason: %q", skip)
	}
	if res == nil || len(res.Versions) == 0 {
		t.Fatalf("expected at least one supported version through STARTTLS, got %+v", res)
	}
}

// TestEnumerateEndpoint_UnknownDialect asserts an unsupported STARTTLS
// dialect is rejected with a non-empty skip reason and no result — the
// observation must record *why* enumeration didn't run, not silently report
// "no versions accepted".
func TestEnumerateEndpoint_UnknownDialect(t *testing.T) {
	res, skip := enumerateEndpoint(context.Background(), contract.TLSEndpoint{
		Host:     "127.0.0.1",
		Port:     1, // unreachable on purpose; we never get past the dialect check
		STARTTLS: "no-such-dialect",
	}, time.Second)
	if res != nil {
		t.Fatalf("expected nil result for unknown dialect, got %+v", res)
	}
	if skip == "" {
		t.Fatalf("expected non-empty skip reason for unknown dialect")
	}
}