package checker

import (
	"bufio"
	"fmt"
	"net"
	"strings"
)

func init() {
	registerStartTLS("imap", starttlsIMAP)
}

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

	if _, err := readLineLimited(rw.Reader); 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 fmt.Errorf("flush CAPABILITY: %w", err)
	}

	supportsSTARTTLS := false
	for {
		line, err := readLineLimited(rw.Reader)
		if err != nil {
			return fmt.Errorf("read CAPABILITY: %w", err)
		}
		if strings.Contains(strings.ToUpper(line), "STARTTLS") {
			supportsSTARTTLS = true
		}
		if strings.HasPrefix(line, "A001 ") {
			rest := strings.TrimSpace(line[len("A001 "):])
			if !strings.HasPrefix(strings.ToUpper(rest), "OK") {
				return fmt.Errorf("CAPABILITY rejected by server: %s", rest)
			}
			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 fmt.Errorf("write STARTTLS: %w", err)
	}
	if err := rw.Flush(); err != nil {
		return fmt.Errorf("flush STARTTLS: %w", err)
	}
	for {
		line, err := readLineLimited(rw.Reader)
		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))
		}
	}
}