package checker

import (
	"bufio"
	"encoding/xml"
	"fmt"
	"io"
	"net"
	"strings"
)

// starttlsUpgrader performs the plaintext portion of a STARTTLS upgrade on
// conn, leaving conn ready for tls.Client(conn, …).Handshake(). On success
// the returned function returns nil; on failure it returns a descriptive
// error (wrap errStartTLSNotOffered when the server advertises no STARTTLS).
type starttlsUpgrader func(conn net.Conn, sni string) error

var starttlsUpgraders = map[string]starttlsUpgrader{
	"smtp":        starttlsSMTP,
	"submission":  starttlsSMTP,
	"imap":        starttlsIMAP,
	"pop3":        starttlsPOP3,
	"xmpp-client": starttlsXMPPClient,
	"xmpp-server": starttlsXMPPServer,
}

// starttlsSMTP implements ESMTP EHLO + STARTTLS (RFC 3207).
func starttlsSMTP(conn net.Conn, sni string) error {
	rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))

	if err := readSMTPGreeting(rw.Reader); err != nil {
		return fmt.Errorf("read greeting: %w", err)
	}

	if _, err := rw.WriteString("EHLO checker.happydomain.org\r\n"); err != nil {
		return fmt.Errorf("write ehlo: %w", err)
	}
	if err := rw.Flush(); err != nil {
		return fmt.Errorf("flush ehlo: %w", err)
	}
	lines, err := readSMTPResponse(rw.Reader)
	if err != nil {
		return fmt.Errorf("read ehlo: %w", err)
	}
	if !hasSTARTTLSExt(lines) {
		return fmt.Errorf("%w: EHLO did not advertise STARTTLS", errStartTLSNotOffered)
	}

	if _, err := rw.WriteString("STARTTLS\r\n"); err != nil {
		return fmt.Errorf("write starttls: %w", err)
	}
	if err := rw.Flush(); err != nil {
		return fmt.Errorf("flush starttls: %w", err)
	}
	resp, err := readSMTPResponse(rw.Reader)
	if err != nil {
		return fmt.Errorf("read starttls: %w", err)
	}
	if len(resp) == 0 || !strings.HasPrefix(resp[0], "220") {
		return fmt.Errorf("server refused STARTTLS: %s", strings.Join(resp, " / "))
	}
	return nil
}

func readSMTPGreeting(r *bufio.Reader) error {
	_, err := readSMTPResponse(r)
	return err
}

// readSMTPResponse reads one multi-line SMTP response (lines with "NNN-" are
// continuation, "NNN " terminates).
func readSMTPResponse(r *bufio.Reader) ([]string, error) {
	var out []string
	for {
		line, err := r.ReadString('\n')
		if err != nil {
			return out, err
		}
		line = strings.TrimRight(line, "\r\n")
		out = append(out, line)
		if len(line) < 4 || line[3] == ' ' {
			return out, nil
		}
	}
}

func hasSTARTTLSExt(lines []string) bool {
	for _, l := range lines {
		if len(l) < 4 {
			continue
		}
		rest := strings.ToUpper(strings.TrimSpace(l[4:]))
		if rest == "STARTTLS" || strings.HasPrefix(rest, "STARTTLS ") {
			return true
		}
	}
	return false
}

// starttlsIMAP implements RFC 3501 STARTTLS.
func starttlsIMAP(conn net.Conn, sni string) error {
	rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))

	if _, err := rw.ReadString('\n'); err != nil {
		return fmt.Errorf("read greeting: %w", err)
	}

	if _, err := rw.WriteString("A001 CAPABILITY\r\n"); err != nil {
		return fmt.Errorf("write CAPABILITY: %w", err)
	}
	if err := rw.Flush(); err != nil {
		return err
	}

	supportsSTARTTLS := false
	for {
		line, err := rw.ReadString('\n')
		if err != nil {
			return fmt.Errorf("read CAPABILITY: %w", err)
		}
		if strings.Contains(strings.ToUpper(line), "STARTTLS") {
			supportsSTARTTLS = true
		}
		if strings.HasPrefix(line, "A001 ") {
			break
		}
	}
	if !supportsSTARTTLS {
		return fmt.Errorf("%w: IMAP CAPABILITY did not advertise STARTTLS", errStartTLSNotOffered)
	}

	if _, err := rw.WriteString("A002 STARTTLS\r\n"); err != nil {
		return err
	}
	if err := rw.Flush(); err != nil {
		return err
	}
	for {
		line, err := rw.ReadString('\n')
		if err != nil {
			return fmt.Errorf("read STARTTLS response: %w", err)
		}
		if strings.HasPrefix(line, "A002 OK") {
			return nil
		}
		if strings.HasPrefix(line, "A002 ") {
			return fmt.Errorf("server refused STARTTLS: %s", strings.TrimSpace(line))
		}
	}
}

// starttlsPOP3 implements RFC 2595 STLS.
func starttlsPOP3(conn net.Conn, sni string) error {
	rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))

	greeting, err := rw.ReadString('\n')
	if err != nil {
		return fmt.Errorf("read greeting: %w", err)
	}
	if !strings.HasPrefix(greeting, "+OK") {
		return fmt.Errorf("unexpected POP3 greeting: %s", strings.TrimSpace(greeting))
	}

	if _, err := rw.WriteString("CAPA\r\n"); err != nil {
		return err
	}
	if err := rw.Flush(); err != nil {
		return err
	}
	first, err := rw.ReadString('\n')
	if err != nil {
		return fmt.Errorf("read CAPA: %w", err)
	}
	supportsSTLS := false
	if strings.HasPrefix(first, "+OK") {
		for {
			line, err := rw.ReadString('\n')
			if err != nil {
				return fmt.Errorf("read CAPA body: %w", err)
			}
			line = strings.TrimRight(line, "\r\n")
			if line == "." {
				break
			}
			if strings.EqualFold(line, "STLS") {
				supportsSTLS = true
			}
		}
	}
	if !supportsSTLS {
		return fmt.Errorf("%w: POP3 CAPA did not advertise STLS", errStartTLSNotOffered)
	}

	if _, err := rw.WriteString("STLS\r\n"); err != nil {
		return err
	}
	if err := rw.Flush(); err != nil {
		return err
	}
	resp, err := rw.ReadString('\n')
	if err != nil {
		return fmt.Errorf("read STLS response: %w", err)
	}
	if !strings.HasPrefix(resp, "+OK") {
		return fmt.Errorf("server refused STLS: %s", strings.TrimSpace(resp))
	}
	return nil
}

// starttlsXMPPClient implements RFC 6120 STARTTLS for c2s streams.
func starttlsXMPPClient(conn net.Conn, sni string) error {
	return starttlsXMPP(conn, sni, "jabber:client")
}

// starttlsXMPPServer implements RFC 6120 STARTTLS for s2s streams.
func starttlsXMPPServer(conn net.Conn, sni string) error {
	return starttlsXMPP(conn, sni, "jabber:server")
}

func starttlsXMPP(conn net.Conn, sni, ns string) error {
	header := fmt.Sprintf(`<?xml version='1.0'?><stream:stream xmlns='%s' xmlns:stream='http://etherx.jabber.org/streams' version='1.0' to='%s'>`, ns, sni)
	if _, err := io.WriteString(conn, header); err != nil {
		return fmt.Errorf("write stream header: %w", err)
	}

	dec := xml.NewDecoder(conn)

	// Read the inbound <stream:stream> opening and its <stream:features>.
	hasStartTLS := false
	for {
		tok, err := dec.Token()
		if err != nil {
			return fmt.Errorf("read stream features: %w", err)
		}
		if se, ok := tok.(xml.StartElement); ok {
			if se.Name.Local == "features" {
				// Scan features children.
				for {
					t2, err := dec.Token()
					if err != nil {
						return fmt.Errorf("read features body: %w", err)
					}
					switch ee := t2.(type) {
					case xml.StartElement:
						if ee.Name.Local == "starttls" {
							hasStartTLS = true
						}
						_ = dec.Skip()
					case xml.EndElement:
						if ee.Name.Local == "features" {
							goto doneFeatures
						}
					}
				}
			}
		}
	}
doneFeatures:
	if !hasStartTLS {
		return fmt.Errorf("%w: XMPP features did not advertise starttls", errStartTLSNotOffered)
	}

	if _, err := io.WriteString(conn, `<starttls xmlns='urn:ietf:params:xml:ns:xmpp-tls'/>`); err != nil {
		return fmt.Errorf("write starttls: %w", err)
	}

	for {
		tok, err := dec.Token()
		if err != nil {
			return fmt.Errorf("read proceed: %w", err)
		}
		if se, ok := tok.(xml.StartElement); ok {
			switch se.Name.Local {
			case "proceed":
				return nil
			case "failure":
				return fmt.Errorf("server refused STARTTLS (<failure/>)")
			}
		}
	}
}