package test

import (
	"bufio"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"math/big"
	"net"
	"strings"
	"sync"
	"testing"
	"time"
)

type fakeSMTPOptions struct {
	authMechanisms string
	startTLS       bool
	rejectRcpt     bool
	rejectData     bool
	implicitTLS    bool
	listenAddress  string
}

type fakeSMTPSession struct {
	commands []string
	from     string
	to       string
	data     string
	username string
	password string
	usedTLS  bool
}

type fakeSMTPServer struct {
	listener net.Listener
	options  fakeSMTPOptions
	cert     tls.Certificate
	roots    *x509.CertPool
	mu       sync.Mutex
	sessions []*fakeSMTPSession
	wg       sync.WaitGroup
}

func newFakeSMTPServer(t *testing.T, options fakeSMTPOptions) *fakeSMTPServer {
	t.Helper()
	address := options.listenAddress
	if address == "" {
		address = "127.0.0.1:0"
	}
	listener, err := net.Listen("tcp", address)
	if err != nil {
		t.Skipf("cannot listen on %s: %v", address, err)
	}
	cert, roots := selfSignedLoopbackCert(t)
	server := &fakeSMTPServer{listener: listener, options: options, cert: cert, roots: roots}
	server.wg.Add(1)
	go server.acceptLoop()
	t.Cleanup(func() {
		_ = listener.Close()
		server.wg.Wait()
	})
	return server
}

func selfSignedLoopbackCert(t *testing.T) (tls.Certificate, *x509.CertPool) {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	template := &x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject:      pkix.Name{CommonName: "127.0.0.1"},
		NotBefore:    time.Now().Add(-time.Hour),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		IPAddresses:  []net.IP{net.ParseIP("127.0.0.1")},
	}
	der, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
	if err != nil {
		t.Fatal(err)
	}
	parsed, err := x509.ParseCertificate(der)
	if err != nil {
		t.Fatal(err)
	}
	roots := x509.NewCertPool()
	roots.AddCert(parsed)
	return tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key}, roots
}

func (s *fakeSMTPServer) port() int {
	return s.listener.Addr().(*net.TCPAddr).Port
}

func (s *fakeSMTPServer) clientTLS() *tls.Config {
	return &tls.Config{RootCAs: s.roots}
}

func (s *fakeSMTPServer) received() []*fakeSMTPSession {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]*fakeSMTPSession(nil), s.sessions...)
}

func (s *fakeSMTPServer) acceptLoop() {
	defer s.wg.Done()
	for {
		connection, err := s.listener.Accept()
		if err != nil {
			return
		}
		s.wg.Add(1)
		go func() {
			defer s.wg.Done()
			s.serve(connection)
		}()
	}
}

func (s *fakeSMTPServer) serve(connection net.Conn) {
	defer func() { _ = connection.Close() }()
	_ = connection.SetDeadline(time.Now().Add(10 * time.Second))
	session := &fakeSMTPSession{}
	s.mu.Lock()
	s.sessions = append(s.sessions, session)
	s.mu.Unlock()

	if s.options.implicitTLS {
		connection = tls.Server(connection, &tls.Config{Certificates: []tls.Certificate{s.cert}})
		session.usedTLS = true
	}
	reader := bufio.NewReader(connection)
	reply := func(line string) { _, _ = connection.Write([]byte(line + "\r\n")) }
	reply("220 fake ESMTP")

	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			return
		}
		line = strings.TrimRight(line, "\r\n")
		s.record(session, line)
		verb := strings.ToUpper(strings.Fields(line + " ")[0])

		switch verb {
		case "EHLO":
			s.replyEhlo(reply, session)
		case "STARTTLS":
			reply("220 ready")
			connection = tls.Server(connection, &tls.Config{Certificates: []tls.Certificate{s.cert}})
			reader = bufio.NewReader(connection)
			reply = func(line string) { _, _ = connection.Write([]byte(line + "\r\n")) }
			session.usedTLS = true
		case "AUTH":
			if !s.authenticate(line, reader, reply, session) {
				return
			}
		case "MAIL":
			session.from = smtpAngleContent(line, "<", ">")
			reply("250 ok")
		case "RCPT":
			if s.options.rejectRcpt {
				reply("550 no such user")
				continue
			}
			session.to = smtpAngleContent(line, "<", ">")
			reply("250 ok")
		case "DATA":
			if s.options.rejectData {
				reply("554 refused")
				continue
			}
			reply("354 go ahead")
			session.data = readSMTPData(reader)
			reply("250 queued")
		case "QUIT":
			reply("221 bye")
			return
		default:
			reply("250 ok")
		}
	}
}

func (s *fakeSMTPServer) record(session *fakeSMTPSession, line string) {
	s.mu.Lock()
	session.commands = append(session.commands, line)
	s.mu.Unlock()
}

func (s *fakeSMTPServer) replyEhlo(reply func(string), session *fakeSMTPSession) {
	lines := []string{"250-fake"}
	if s.options.startTLS && !session.usedTLS {
		lines = append(lines, "250-STARTTLS")
	}
	if s.options.authMechanisms != "" {
		lines = append(lines, "250-AUTH "+s.options.authMechanisms)
	}
	lines = append(lines, "250 8BITMIME")
	for _, line := range lines {
		reply(line)
	}
}

func (s *fakeSMTPServer) authenticate(line string, reader *bufio.Reader, reply func(string), session *fakeSMTPSession) bool {
	fields := strings.Fields(line)
	if len(fields) < 2 {
		reply("501 syntax")
		return false
	}
	switch strings.ToUpper(fields[1]) {
	case "PLAIN":
		raw, _ := base64.StdEncoding.DecodeString(fields[2])
		parts := strings.Split(string(raw), "\x00")
		session.username, session.password = parts[1], parts[2]
	case "LOGIN":
		reply("334 " + base64.StdEncoding.EncodeToString([]byte("Username:")))
		username, _ := reader.ReadString('\n')
		reply("334 " + base64.StdEncoding.EncodeToString([]byte("Password:")))
		password, _ := reader.ReadString('\n')
		session.username = decodeBase64Line(username)
		session.password = decodeBase64Line(password)
	default:
		reply("504 unsupported")
		return false
	}
	reply("235 ok")
	return true
}

func decodeBase64Line(line string) string {
	raw, _ := base64.StdEncoding.DecodeString(strings.TrimSpace(line))
	return string(raw)
}

func smtpAngleContent(line string, open string, close string) string {
	start := strings.Index(line, open)
	end := strings.LastIndex(line, close)
	if start < 0 || end <= start {
		return ""
	}
	return line[start+1 : end]
}

func readSMTPData(reader *bufio.Reader) string {
	var data strings.Builder
	for {
		line, err := reader.ReadString('\n')
		if err != nil || line == ".\r\n" {
			return data.String()
		}
		data.WriteString(strings.TrimPrefix(line, "."))
	}
}