package checker

import (
	"bufio"
	"errors"
	"io"
	"net"
	"strings"
	"testing"
	"time"
)

// runStartTLS drives upgrader against a fake server. The server callback runs
// on the peer end of an in-memory pipe and may read/write the plaintext
// dialect transcript. The test deadline guards both ends from hanging.
func runStartTLS(t *testing.T, upgrader func(net.Conn, string) error, sni string, server func(net.Conn) error) error {
	t.Helper()
	clientConn, serverConn := net.Pipe()
	deadline := time.Now().Add(2 * time.Second)
	_ = clientConn.SetDeadline(deadline)
	_ = serverConn.SetDeadline(deadline)

	srvErr := make(chan error, 1)
	go func() {
		defer serverConn.Close()
		srvErr <- server(serverConn)
	}()

	clientErr := upgrader(clientConn, sni)
	clientConn.Close()

	if err := <-srvErr; err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrClosedPipe) {
		t.Logf("server side returned: %v", err)
	}
	return clientErr
}

// readLineCRLF reads one CRLF-terminated line.
func readLineCRLF(r *bufio.Reader) (string, error) {
	line, err := r.ReadString('\n')
	return strings.TrimRight(line, "\r\n"), err
}

func TestStartTLS_SMTP_OK(t *testing.T) {
	err := runStartTLS(t, starttlsSMTP, "mail.example.com", func(c net.Conn) error {
		br := bufio.NewReader(c)
		if _, err := io.WriteString(c, "220 mail.example.com ESMTP\r\n"); err != nil {
			return err
		}
		ehlo, err := readLineCRLF(br)
		if err != nil {
			return err
		}
		if !strings.HasPrefix(ehlo, "EHLO ") {
			return errors.New("expected EHLO")
		}
		if _, err := io.WriteString(c, "250-mail.example.com\r\n250-SIZE 10485760\r\n250 STARTTLS\r\n"); err != nil {
			return err
		}
		stls, err := readLineCRLF(br)
		if err != nil {
			return err
		}
		if stls != "STARTTLS" {
			return errors.New("expected STARTTLS")
		}
		_, err = io.WriteString(c, "220 ready\r\n")
		return err
	})
	if err != nil {
		t.Fatalf("expected success, got: %v", err)
	}
}

func TestStartTLS_SMTP_NotAdvertised(t *testing.T) {
	err := runStartTLS(t, starttlsSMTP, "mail.example.com", func(c net.Conn) error {
		br := bufio.NewReader(c)
		_, _ = io.WriteString(c, "220 mail.example.com ESMTP\r\n")
		if _, err := readLineCRLF(br); err != nil {
			return err
		}
		_, err := io.WriteString(c, "250-mail.example.com\r\n250 SIZE 10485760\r\n")
		return err
	})
	if !errors.Is(err, errStartTLSNotOffered) {
		t.Fatalf("expected errStartTLSNotOffered, got: %v", err)
	}
}

func TestStartTLS_SMTP_Refused(t *testing.T) {
	err := runStartTLS(t, starttlsSMTP, "mail.example.com", func(c net.Conn) error {
		br := bufio.NewReader(c)
		_, _ = io.WriteString(c, "220 mail.example.com ESMTP\r\n")
		_, _ = readLineCRLF(br)
		_, _ = io.WriteString(c, "250-mail.example.com\r\n250 STARTTLS\r\n")
		_, _ = readLineCRLF(br)
		_, err := io.WriteString(c, "454 TLS not available\r\n")
		return err
	})
	if err == nil {
		t.Fatal("expected refusal error")
	}
	if errors.Is(err, errStartTLSNotOffered) {
		t.Fatalf("refusal should not be classified as not-offered: %v", err)
	}
}

func TestStartTLS_IMAP_OK(t *testing.T) {
	err := runStartTLS(t, starttlsIMAP, "imap.example.com", func(c net.Conn) error {
		br := bufio.NewReader(c)
		_, _ = io.WriteString(c, "* OK IMAP4rev1 ready\r\n")
		cap1, err := readLineCRLF(br)
		if err != nil {
			return err
		}
		if !strings.HasSuffix(cap1, "CAPABILITY") {
			return errors.New("expected CAPABILITY")
		}
		_, _ = io.WriteString(c, "* CAPABILITY IMAP4rev1 STARTTLS LOGINDISABLED\r\nA001 OK CAPABILITY completed\r\n")
		stls, err := readLineCRLF(br)
		if err != nil {
			return err
		}
		if !strings.HasSuffix(stls, "STARTTLS") {
			return errors.New("expected STARTTLS")
		}
		_, err = io.WriteString(c, "A002 OK Begin TLS\r\n")
		return err
	})
	if err != nil {
		t.Fatalf("expected success, got: %v", err)
	}
}

func TestStartTLS_IMAP_Refused(t *testing.T) {
	err := runStartTLS(t, starttlsIMAP, "imap.example.com", func(c net.Conn) error {
		br := bufio.NewReader(c)
		_, _ = io.WriteString(c, "* OK IMAP4rev1 ready\r\n")
		_, _ = readLineCRLF(br)
		_, _ = io.WriteString(c, "* CAPABILITY IMAP4rev1 STARTTLS\r\nA001 OK CAPABILITY completed\r\n")
		_, _ = readLineCRLF(br)
		_, err := io.WriteString(c, "A002 NO STARTTLS unavailable\r\n")
		return err
	})
	if err == nil {
		t.Fatal("expected refusal error")
	}
	if errors.Is(err, errStartTLSNotOffered) {
		t.Fatalf("refusal should not be classified as not-offered: %v", err)
	}
}

func TestStartTLS_IMAP_NotAdvertised(t *testing.T) {
	err := runStartTLS(t, starttlsIMAP, "imap.example.com", func(c net.Conn) error {
		br := bufio.NewReader(c)
		_, _ = io.WriteString(c, "* OK IMAP4rev1 ready\r\n")
		_, _ = readLineCRLF(br)
		_, err := io.WriteString(c, "* CAPABILITY IMAP4rev1 LOGINDISABLED\r\nA001 OK CAPABILITY completed\r\n")
		return err
	})
	if !errors.Is(err, errStartTLSNotOffered) {
		t.Fatalf("expected errStartTLSNotOffered, got: %v", err)
	}
}

func TestStartTLS_POP3_OK(t *testing.T) {
	err := runStartTLS(t, starttlsPOP3, "pop.example.com", func(c net.Conn) error {
		br := bufio.NewReader(c)
		_, _ = io.WriteString(c, "+OK POP3 ready\r\n")
		capa, err := readLineCRLF(br)
		if err != nil {
			return err
		}
		if capa != "CAPA" {
			return errors.New("expected CAPA")
		}
		_, _ = io.WriteString(c, "+OK capa list\r\nUSER\r\nSTLS\r\n.\r\n")
		stls, err := readLineCRLF(br)
		if err != nil {
			return err
		}
		if stls != "STLS" {
			return errors.New("expected STLS")
		}
		_, err = io.WriteString(c, "+OK begin TLS\r\n")
		return err
	})
	if err != nil {
		t.Fatalf("expected success, got: %v", err)
	}
}

func TestStartTLS_POP3_NotAdvertised(t *testing.T) {
	err := runStartTLS(t, starttlsPOP3, "pop.example.com", func(c net.Conn) error {
		br := bufio.NewReader(c)
		_, _ = io.WriteString(c, "+OK POP3 ready\r\n")
		_, _ = readLineCRLF(br)
		_, err := io.WriteString(c, "+OK capa list\r\nUSER\r\n.\r\n")
		return err
	})
	if !errors.Is(err, errStartTLSNotOffered) {
		t.Fatalf("expected errStartTLSNotOffered, got: %v", err)
	}
}

func TestStartTLS_POP3_Refused(t *testing.T) {
	err := runStartTLS(t, starttlsPOP3, "pop.example.com", func(c net.Conn) error {
		br := bufio.NewReader(c)
		_, _ = io.WriteString(c, "+OK POP3 ready\r\n")
		_, _ = readLineCRLF(br)
		_, _ = io.WriteString(c, "+OK capa list\r\nUSER\r\nSTLS\r\n.\r\n")
		_, _ = readLineCRLF(br)
		_, err := io.WriteString(c, "-ERR STLS unavailable\r\n")
		return err
	})
	if err == nil {
		t.Fatal("expected refusal error")
	}
	if errors.Is(err, errStartTLSNotOffered) {
		t.Fatalf("refusal should not be classified as not-offered: %v", err)
	}
}

func TestStartTLS_XMPP_OK(t *testing.T) {
	err := runStartTLS(t, starttlsXMPPClient, "xmpp.example.com", func(c net.Conn) error {
		br := bufio.NewReader(c)
		// Read the client's stream header (one line is enough for our writer).
		buf := make([]byte, 1024)
		if _, err := br.Read(buf); err != nil {
			return err
		}
		_, _ = io.WriteString(c,
			`<?xml version='1.0'?><stream:stream xmlns='jabber:client' xmlns:stream='http://etherx.jabber.org/streams' id='1' from='xmpp.example.com' version='1.0'>`+
				`<stream:features><starttls xmlns='urn:ietf:params:xml:ns:xmpp-tls'><required/></starttls></stream:features>`)
		// Read the <starttls/> request from the client.
		if _, err := br.Read(buf); err != nil {
			return err
		}
		_, err := io.WriteString(c, `<proceed xmlns='urn:ietf:params:xml:ns:xmpp-tls'/>`)
		return err
	})
	if err != nil {
		t.Fatalf("expected success, got: %v", err)
	}
}

func TestStartTLS_XMPP_NotAdvertised(t *testing.T) {
	err := runStartTLS(t, starttlsXMPPClient, "xmpp.example.com", func(c net.Conn) error {
		br := bufio.NewReader(c)
		buf := make([]byte, 1024)
		if _, err := br.Read(buf); err != nil {
			return err
		}
		_, err := io.WriteString(c,
			`<?xml version='1.0'?><stream:stream xmlns='jabber:client' xmlns:stream='http://etherx.jabber.org/streams' id='1' from='xmpp.example.com' version='1.0'>`+
				`<stream:features><mechanisms xmlns='urn:ietf:params:xml:ns:xmpp-sasl'><mechanism>PLAIN</mechanism></mechanisms></stream:features>`)
		return err
	})
	if !errors.Is(err, errStartTLSNotOffered) {
		t.Fatalf("expected errStartTLSNotOffered, got: %v", err)
	}
}

func TestStartTLS_XMPP_Refused(t *testing.T) {
	err := runStartTLS(t, starttlsXMPPClient, "xmpp.example.com", func(c net.Conn) error {
		br := bufio.NewReader(c)
		buf := make([]byte, 1024)
		if _, err := br.Read(buf); err != nil {
			return err
		}
		_, _ = io.WriteString(c,
			`<?xml version='1.0'?><stream:stream xmlns='jabber:client' xmlns:stream='http://etherx.jabber.org/streams' id='1' from='xmpp.example.com' version='1.0'>`+
				`<stream:features><starttls xmlns='urn:ietf:params:xml:ns:xmpp-tls'/></stream:features>`)
		if _, err := br.Read(buf); err != nil {
			return err
		}
		_, err := io.WriteString(c, `<failure xmlns='urn:ietf:params:xml:ns:xmpp-tls'/>`)
		return err
	})
	if err == nil {
		t.Fatal("expected failure error")
	}
	if errors.Is(err, errStartTLSNotOffered) {
		t.Fatalf("<failure/> should not be classified as not-offered: %v", err)
	}
}

func TestStartTLS_XMPP_StreamError(t *testing.T) {
	err := runStartTLS(t, starttlsXMPPClient, "xmpp.example.com", func(c net.Conn) error {
		br := bufio.NewReader(c)
		buf := make([]byte, 1024)
		if _, err := br.Read(buf); err != nil {
			return err
		}
		_, err := io.WriteString(c,
			`<?xml version='1.0'?><stream:stream xmlns='jabber:client' xmlns:stream='http://etherx.jabber.org/streams' id='1' from='xmpp.example.com' version='1.0'>`+
				`<stream:error><host-unknown xmlns='urn:ietf:params:xml:ns:xmpp-streams'/></stream:error>`)
		return err
	})
	if err == nil {
		t.Fatal("expected stream:error to surface as error")
	}
}

func TestStartTLS_LDAP_OK(t *testing.T) {
	err := runStartTLS(t, starttlsLDAP, "ldap.example.com", func(c net.Conn) error {
		// Drain the StartTLS request (fixed 31 bytes: 0x30 0x1d + 29 bytes).
		req := make([]byte, 31)
		if _, err := io.ReadFull(c, req); err != nil {
			return err
		}
		// Build a minimal ExtendedResponse with resultCode=0.
		// LDAPMessage SEQUENCE { messageID INTEGER 1, [APPLICATION 24] SEQUENCE { resultCode ENUMERATED 0, matchedDN "", diagnosticMessage "" } }
		resp := []byte{
			0x30, 0x0c, // SEQUENCE, length 12
			0x02, 0x01, 0x01, // messageID = 1
			0x78, 0x07, // [APPLICATION 24], length 7
			0x0a, 0x01, 0x00, // resultCode ENUMERATED 0
			0x04, 0x00, // matchedDN ""
			0x04, 0x00, // diagnosticMessage ""
		}
		_, err := c.Write(resp)
		return err
	})
	if err != nil {
		t.Fatalf("expected success, got: %v", err)
	}
}

func TestStartTLS_LDAP_WrongTag(t *testing.T) {
	err := runStartTLS(t, starttlsLDAP, "ldap.example.com", func(c net.Conn) error {
		req := make([]byte, 31)
		if _, err := io.ReadFull(c, req); err != nil {
			return err
		}
		_, err := c.Write([]byte{0x42, 0x00})
		return err
	})
	if err == nil {
		t.Fatal("expected error for wrong tag")
	}
	if errors.Is(err, errStartTLSNotOffered) {
		t.Fatalf("malformed response should not be classified as not-offered: %v", err)
	}
}

func TestStartTLS_LDAP_OversizedLength(t *testing.T) {
	err := runStartTLS(t, starttlsLDAP, "ldap.example.com", func(c net.Conn) error {
		req := make([]byte, 31)
		if _, err := io.ReadFull(c, req); err != nil {
			return err
		}
		// SEQUENCE with long-form length = 0x10000 (64 KiB) — beyond our 16 KiB cap.
		_, err := c.Write([]byte{0x30, 0x83, 0x01, 0x00, 0x00})
		return err
	})
	if err == nil {
		t.Fatal("expected oversized-length error")
	}
}

func TestStartTLS_LDAP_TruncatedBody(t *testing.T) {
	err := runStartTLS(t, starttlsLDAP, "ldap.example.com", func(c net.Conn) error {
		req := make([]byte, 31)
		if _, err := io.ReadFull(c, req); err != nil {
			return err
		}
		// Announce 12 bytes of body, only send 5 then close.
		_, err := c.Write([]byte{0x30, 0x0c, 0x02, 0x01, 0x01, 0x78, 0x07})
		return err
	})
	if err == nil {
		t.Fatal("expected error on truncated body")
	}
}

func TestStartTLS_LDAP_DiagnosticMessageOver4KiB(t *testing.T) {
	// A real-world response with a verbose diagnosticMessage can exceed the
	// previous 4 KiB cap. Confirm the bumped 16 KiB cap accepts it.
	const diagLen = 8000
	diag := make([]byte, diagLen)
	for i := range diag {
		diag[i] = 'x'
	}
	err := runStartTLS(t, starttlsLDAP, "ldap.example.com", func(c net.Conn) error {
		req := make([]byte, 31)
		if _, err := io.ReadFull(c, req); err != nil {
			return err
		}
		// Body: messageID(3) + extResp tag(1) + extResp len(3) + resultCode(3) + matchedDN(2) + diag tag+long-len(4) + diag bytes
		// extResp inner length = resultCode(3) + matchedDN(2) + diagTLV(4+diagLen) = 9 + diagLen
		extInner := 9 + diagLen
		// Outer SEQUENCE inner length = messageID(3) + extResp TLV(1+3+extInner)
		outerInner := 3 + 4 + extInner
		buf := []byte{0x30, 0x82, byte(outerInner >> 8), byte(outerInner & 0xff)}
		buf = append(buf, 0x02, 0x01, 0x01) // messageID
		buf = append(buf, 0x78, 0x82, byte(extInner>>8), byte(extInner&0xff))
		buf = append(buf, 0x0a, 0x01, 0x00) // resultCode = success
		buf = append(buf, 0x04, 0x00)       // matchedDN ""
		buf = append(buf, 0x04, 0x82, byte(diagLen>>8), byte(diagLen&0xff))
		buf = append(buf, diag...)
		_, err := c.Write(buf)
		return err
	})
	if err != nil {
		t.Fatalf("expected success with verbose diagnosticMessage, got: %v", err)
	}
}

func TestStartTLS_LDAP_Refused(t *testing.T) {
	err := runStartTLS(t, starttlsLDAP, "ldap.example.com", func(c net.Conn) error {
		req := make([]byte, 31)
		if _, err := io.ReadFull(c, req); err != nil {
			return err
		}
		// resultCode = 53 (unwillingToPerform) -> classified as not-offered.
		resp := []byte{
			0x30, 0x0c,
			0x02, 0x01, 0x01,
			0x78, 0x07,
			0x0a, 0x01, 0x35,
			0x04, 0x00,
			0x04, 0x00,
		}
		_, err := c.Write(resp)
		return err
	})
	if !errors.Is(err, errStartTLSNotOffered) {
		t.Fatalf("expected errStartTLSNotOffered for resultCode 53, got: %v", err)
	}
}