package mail

import (
	"bufio"
	"bytes"
	"errors"
	"io"
	"mime"
	"mime/multipart"
	"mime/quotedprintable"
	"net"
	"net/mail"
	"net/textproto"
	"strings"
	"testing"
)

// fakeSMTP is a tiny scripted SMTP server. script maps an uppercased command
// verb (or "DATA_BODY" for the reply after the body's terminating dot, or
// "CLOSE_AFTER_BODY" to drop the connection with no reply at all) to the
// line(s) to write back. Anything not listed gets a plain 250 OK.
func fakeSMTP(t *testing.T, script map[string]string) string {
	t.Helper()
	port, _ := fakeSMTPData(t, script)
	return port
}

// fakeSMTPData is fakeSMTP that also hands back the raw DATA the client sent,
// once the body's terminating dot arrives.
func fakeSMTPData(t *testing.T, script map[string]string) (string, <-chan []byte) {
	t.Helper()
	got := make(chan []byte, 1)
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { ln.Close() })

	go func() {
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		defer conn.Close()
		w := bufio.NewWriter(conn)
		r := bufio.NewReader(conn)
		writeLine := func(s string) {
			w.WriteString(s + "\r\n")
			w.Flush()
		}
		writeLine("220 fake.example ESMTP")
		inData := false
		var data bytes.Buffer
		for {
			line, err := r.ReadString('\n')
			if err != nil {
				return
			}
			raw := line
			line = strings.TrimRight(line, "\r\n")
			if inData {
				if line != "." {
					data.WriteString(raw)
				}
				if line == "." {
					inData = false
					got <- data.Bytes()
					if reply, ok := script["DATA_BODY"]; ok {
						if reply == "CLOSE" {
							return
						}
						writeLine(reply)
					} else {
						writeLine("250 OK")
					}
				}
				continue
			}
			verb := strings.ToUpper(strings.Fields(line)[0])
			if verb == "EHLO" || verb == "HELO" {
				writeLine("250 fake.example")
				continue
			}
			if verb == "DATA" {
				if reply, ok := script["DATA"]; ok {
					if reply == "CLOSE" {
						return
					}
					writeLine(reply)
					continue
				}
				inData = true
				writeLine("354 go ahead")
				continue
			}
			if verb == "QUIT" {
				if reply, ok := script["QUIT"]; ok {
					if reply == "CLOSE" {
						return
					}
					writeLine(reply)
				} else {
					writeLine("221 bye")
				}
				return
			}
			if reply, ok := script[verb]; ok {
				if reply == "CLOSE" {
					return
				}
				writeLine(reply)
				continue
			}
			writeLine("250 OK")
		}
	}()

	host, port, _ := net.SplitHostPort(ln.Addr().String())
	_ = host
	return port, got
}

func testMsg() message {
	return message{To: "rcpt@example.com", Subject: "s", Text: "t", HTML: "<p>h</p>"}
}

func TestSendPhaseOnRecipientFailure(t *testing.T) {
	port := fakeSMTP(t, map[string]string{"RCPT": "550 no such user"})
	s := Sender{Host: "127.0.0.1", Port: port, From: "updates@example.com"}
	err := s.send(testMsg())
	var se *SendError
	if !errors.As(err, &se) {
		t.Fatalf("err = %v, want *SendError", err)
	}
	if se.Phase != PhaseRecipient {
		t.Fatalf("phase = %q, want %q", se.Phase, PhaseRecipient)
	}
	var tp *textproto.Error
	if !errors.As(err, &tp) || tp.Code != 550 {
		t.Fatalf("textproto reply = %+v", tp)
	}
}

func TestSendPhaseOnMailFromFailure(t *testing.T) {
	// No AUTH configured, so a MAIL-stage 550 exercises the connect phase
	// without needing to script a real AUTH challenge/response.
	port := fakeSMTP(t, map[string]string{"MAIL": "550 relay denied"})
	s := Sender{Host: "127.0.0.1", Port: port, From: "updates@example.com"}
	err := s.send(testMsg())
	var se *SendError
	if !errors.As(err, &se) {
		t.Fatalf("err = %v, want *SendError", err)
	}
	if se.Phase != PhaseConnect {
		t.Fatalf("phase = %q, want %q", se.Phase, PhaseConnect)
	}
	var tp *textproto.Error
	if !errors.As(err, &tp) || tp.Code != 550 {
		t.Fatalf("textproto reply = %+v", tp)
	}
}

func TestSendPhaseOnDataCommandFailure(t *testing.T) {
	port := fakeSMTP(t, map[string]string{"DATA": "554 no thanks"})
	s := Sender{Host: "127.0.0.1", Port: port, From: "updates@example.com"}
	err := s.send(testMsg())
	var se *SendError
	if !errors.As(err, &se) {
		t.Fatalf("err = %v, want *SendError", err)
	}
	if se.Phase != PhaseData {
		t.Fatalf("phase = %q, want %q", se.Phase, PhaseData)
	}
	var tp *textproto.Error
	if !errors.As(err, &tp) || tp.Code != 554 {
		t.Fatalf("textproto reply = %+v", tp)
	}
}

// The connection drops after the body's terminating dot but before any reply
// is read: the data phase, but with no textproto error, since the server
// never spoke back at all.
func TestSendPhaseOnDataNoReply(t *testing.T) {
	port := fakeSMTP(t, map[string]string{"DATA_BODY": "CLOSE"})
	s := Sender{Host: "127.0.0.1", Port: port, From: "updates@example.com"}
	err := s.send(testMsg())
	var se *SendError
	if !errors.As(err, &se) {
		t.Fatalf("err = %v, want *SendError", err)
	}
	if se.Phase != PhaseData {
		t.Fatalf("phase = %q, want %q", se.Phase, PhaseData)
	}
	var tp *textproto.Error
	if errors.As(err, &tp) {
		t.Fatalf("expected no textproto reply, got %+v", tp)
	}
}

// A 250 on DATA_BODY (the message is accepted) followed by a dropped
// connection on QUIT: the message was already delivered, so this must be the
// quit phase, not data or connect.
func TestSendPhaseOnQuitFailure(t *testing.T) {
	port := fakeSMTP(t, map[string]string{"QUIT": "CLOSE"})
	s := Sender{Host: "127.0.0.1", Port: port, From: "updates@example.com"}
	err := s.send(testMsg())
	var se *SendError
	if !errors.As(err, &se) {
		t.Fatalf("err = %v, want *SendError", err)
	}
	if se.Phase != PhaseQuit {
		t.Fatalf("phase = %q, want %q", se.Phase, PhaseQuit)
	}
}

func TestSendPrepareErrors(t *testing.T) {
	s := Sender{}
	err := s.send(testMsg())
	var se *SendError
	if !errors.As(err, &se) || se.Phase != PhasePrepare {
		t.Fatalf("not configured: err = %v", err)
	}

	s2 := Sender{Host: "127.0.0.1", Port: "0", From: "updates@example.com"}
	err2 := s2.send(message{Subject: "s", Text: "t", HTML: "<p>h</p>"})
	var se2 *SendError
	if !errors.As(err2, &se2) || se2.Phase != PhasePrepare {
		t.Fatalf("no recipient: err = %v", err2)
	}
}

// The envelope carries the bare address. "Vantage <x@y>" as MAIL FROM left
// rspamd with no envelope sender, and mailcow skipped DKIM signing.
func TestAddrSpecStripsDisplayName(t *testing.T) {
	for in, want := range map[string]string{
		"Vantage <support@example.com>": "support@example.com",
		"support@example.com":           "support@example.com",
		" a@example.com ":               "a@example.com",
		"not an address":                "not an address",
	} {
		if got := addrSpec(in); got != want {
			t.Errorf("addrSpec(%q) = %q, want %q", in, got, want)
		}
	}
}

// Every part must declare quoted-printable and put only 7-bit bytes on the
// wire. Raw UTF-8 under an implied 7bit encoding scored R_BAD_CTE_7BIT in
// rspamd.
func TestEnvelopePartsAreQuotedPrintable(t *testing.T) {
	s := Sender{From: "Vantage <support@example.com>"}
	m := message{
		To:      "a@example.com",
		Subject: "Smith & Sons · restored",
		Text:    "VANTAGE · Account\nBilling has resumed.",
		HTML:    "<p>VANTAGE · Account</p><p>" + strings.Repeat("x", 2000) + "</p>",
	}
	raw, err := s.envelope(m)
	if err != nil {
		t.Fatal(err)
	}
	for i, c := range raw {
		if c > 0x7e && c != '\r' && c != '\n' {
			t.Fatalf("byte %d is 0x%x: the wire must be 7-bit", i, c)
		}
	}
	for _, line := range strings.Split(string(raw), "\r\n") {
		if len(line) > 998 {
			t.Fatalf("line of %d bytes exceeds SMTP's 998", len(line))
		}
	}

	msg, err := mail.ReadMessage(bytes.NewReader(raw))
	if err != nil {
		t.Fatal(err)
	}
	_, params, err := mime.ParseMediaType(msg.Header.Get("Content-Type"))
	if err != nil {
		t.Fatal(err)
	}
	r := multipart.NewReader(msg.Body, params["boundary"])
	want := []string{m.Text, m.HTML}
	for i := range want {
		p, err := r.NextRawPart()
		if err != nil {
			t.Fatalf("part %d: %v", i, err)
		}
		if cte := p.Header.Get("Content-Transfer-Encoding"); cte != "quoted-printable" {
			t.Fatalf("part %d Content-Transfer-Encoding = %q", i, cte)
		}
		// NextPart would decode for us; NextRawPart keeps the check honest by
		// decoding here, so a missing header cannot pass by accident.
		body, err := io.ReadAll(quotedprintable.NewReader(p))
		if err != nil {
			t.Fatal(err)
		}
		// Quoted-printable text mode writes line breaks as CRLF, which is the
		// canonical form on the wire.
		if strings.ReplaceAll(string(body), "\r\n", "\n") != want[i] {
			t.Fatalf("part %d decodes to %q", i, body)
		}
	}
}

func TestEnvelopeWritesExtraHeadersSanitised(t *testing.T) {
	s := Sender{From: "Vantage <updates@example.com>"}
	raw, err := s.envelope(message{
		To: "a@example.com", Subject: "s", Text: "t", HTML: "<p>h</p>",
		Headers: map[string]string{
			"List-Unsubscribe-Post": "List-Unsubscribe=One-Click",
			"List-Unsubscribe":      "<https://x.example/u?t=1>\r\nBcc: evil@example.com",
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	msg, err := mail.ReadMessage(bytes.NewReader(raw))
	if err != nil {
		t.Fatal(err)
	}
	if got := msg.Header.Get("List-Unsubscribe-Post"); got != "List-Unsubscribe=One-Click" {
		t.Fatalf("List-Unsubscribe-Post = %q", got)
	}
	if got := msg.Header.Get("Bcc"); got != "" {
		t.Fatalf("header injection: Bcc = %q", got)
	}
	if !strings.HasPrefix(msg.Header.Get("List-Unsubscribe"), "<https://x.example/u?t=1>") {
		t.Fatalf("List-Unsubscribe = %q", msg.Header.Get("List-Unsubscribe"))
	}
}