package protocol

import (
	"bytes"
	"encoding/base64"
	"fmt"
	stdhtml "html"
	"io"
	"mime"
	"mime/multipart"
	"mime/quotedprintable"
	netmail "net/mail"
	"net/textproto"
	"strings"
	"unicode/utf8"

	"github.com/emersion/go-imap"
	messagemail "github.com/emersion/go-message/mail"
	xhtml "golang.org/x/net/html"
)

const (
	maxDraftBodyBytes   = 512 * 1024
	maxBodyParts        = 256
	maxAttachments      = 100
	maxAddresses        = 50
	maxAddressBytes     = 320
	maxSubjectBytes     = 998
	maxFilenameBytes    = 1024
	maxMediaTypeBytes   = 255
	maxMailboxNameBytes = 1024
)

// Attachment-carrying messages are bounded on five independent layers: the
// attachment count, the single decoded part size, the decoded total, the
// encoded SMTP message size, and the MIME tree size (maxBodyParts).
const (
	MaxMessageAttachments   = 10
	MaxAttachmentPartBytes  = 10 << 20
	MaxTotalAttachmentBytes = 20 << 20
	// maxAttachmentFetchBytes is the encoded octet budget for one IMAP part
	// fetch. It exceeds the decoded limit enough for the worst transfer
	// encodings: base64 inflates by 4/3 (about 13.3 MiB for 10 MiB decoded)
	// and quoted-printable can inflate a 10 MiB body towards 30+ MiB. The
	// decoded size is enforced separately during transfer-encoding decode.
	maxAttachmentFetchBytes = 40 << 20
	// maxDraftBodyFetchBytes is the encoded octet budget for one draft body
	// fetch. Quoted-printable can inflate a 512 KiB body towards 1.5 MiB, so
	// the encoded window must exceed the decoded body limit; the decoded size
	// and exact-length completeness are enforced separately.
	maxDraftBodyFetchBytes = 2 << 20
)

// DraftAttachment is one fully materialized outbound attachment. Blob is the
// decoded content; the MIME writer re-encodes it for transport.
type DraftAttachment struct {
	Filename  string
	MediaType string
	Blob      []byte
}

type DraftContent struct {
	From        string
	To          []string
	CC          []string
	BCC         []string
	Subject     string
	BodyFormat  string
	Body        string
	Attachments []DraftAttachment
}

func BuildDraftMessage(content DraftContent) ([]byte, error) {
	message, err := prepareDraftMessage(content)
	if err != nil {
		return nil, err
	}
	return writeDraftMessage(message, true)
}

func buildSMTPMessage(message FrozenMessage) ([]byte, []string, error) {
	content := DraftContent{
		From:       message.SenderAddress,
		To:         append([]string(nil), message.To...),
		CC:         append([]string(nil), message.CC...),
		BCC:        append([]string(nil), message.BCC...),
		Subject:    message.Subject,
		BodyFormat: message.BodyFormat,
		Body:       message.Body,
	}
	for _, attachment := range message.Attachments {
		content.Attachments = append(content.Attachments, DraftAttachment{
			Filename:  attachment.Filename,
			MediaType: attachment.MediaType,
			Blob:      attachment.Blob,
		})
	}
	prepared, err := prepareDraftMessage(content)
	if err != nil {
		return nil, nil, err
	}
	raw, err := writeDraftMessage(prepared, false)
	if err != nil {
		return nil, nil, err
	}
	recipients := make([]string, 0, len(content.To)+len(content.CC)+len(content.BCC))
	recipients = append(recipients, content.To...)
	recipients = append(recipients, content.CC...)
	recipients = append(recipients, content.BCC...)
	return raw, recipients, nil
}

type preparedDraftMessage struct {
	content DraftContent
	from    []*messagemail.Address
	to      []*messagemail.Address
	cc      []*messagemail.Address
	bcc     []*messagemail.Address
}

func prepareDraftMessage(
	content DraftContent,
) (preparedDraftMessage, error) {
	if len(content.Body) > maxDraftBodyBytes {
		return preparedDraftMessage{}, fmt.Errorf("draft body is too large")
	}
	from, err := parseAddresses([]string{content.From})
	if err != nil {
		return preparedDraftMessage{}, err
	}
	if len(from) != 1 {
		return preparedDraftMessage{}, fmt.Errorf("draft sender is invalid")
	}
	to, err := parseAddresses(content.To)
	if err != nil {
		return preparedDraftMessage{}, err
	}
	cc, err := parseAddresses(content.CC)
	if err != nil {
		return preparedDraftMessage{}, err
	}
	bcc, err := parseAddresses(content.BCC)
	if err != nil {
		return preparedDraftMessage{}, err
	}
	if len(to)+len(cc)+len(bcc) == 0 {
		return preparedDraftMessage{}, fmt.Errorf("draft requires a recipient")
	}
	if len(to)+len(cc)+len(bcc) > maxSMTPRecipients {
		return preparedDraftMessage{},
			fmt.Errorf("draft recipient list exceeded limit")
	}
	if strings.ContainsAny(content.Subject, "\r\n") ||
		len(content.Subject) > 998 {
		return preparedDraftMessage{}, fmt.Errorf("draft subject is invalid")
	}
	switch content.BodyFormat {
	case "text", "html":
	default:
		return preparedDraftMessage{},
			fmt.Errorf("draft body format is invalid")
	}
	if err := validateDraftAttachments(content.Attachments); err != nil {
		return preparedDraftMessage{}, err
	}
	content.Attachments = normalizeDraftAttachments(content.Attachments)
	return preparedDraftMessage{
		content: content,
		from:    from,
		to:      to,
		cc:      cc,
		bcc:     bcc,
	}, nil
}

func validateDraftAttachments(attachments []DraftAttachment) error {
	if len(attachments) > MaxMessageAttachments {
		return fmt.Errorf("message attachments exceeded limit")
	}
	total := 0
	for _, attachment := range attachments {
		filename := strings.TrimSpace(attachment.Filename)
		if filename == "" ||
			len(filename) > maxFilenameBytes ||
			strings.ContainsAny(filename, "\r\n\x00/") {
			return fmt.Errorf("attachment filename is invalid")
		}
		mediaType := strings.TrimSpace(attachment.MediaType)
		if mediaType == "" || len(mediaType) > maxMediaTypeBytes {
			return fmt.Errorf("attachment media type is invalid")
		}
		if _, _, err := mime.ParseMediaType(mediaType); err != nil {
			return fmt.Errorf("attachment media type is invalid")
		}
		if len(attachment.Blob) == 0 ||
			len(attachment.Blob) > MaxAttachmentPartBytes {
			return fmt.Errorf("attachment size exceeded limit")
		}
		total += len(attachment.Blob)
		if total > MaxTotalAttachmentBytes {
			return fmt.Errorf("message attachments exceeded limit")
		}
	}
	return nil
}

// normalizeDraftAttachments copies the validated attachments and reduces every
// media type to its lowercase base type (for example "text/plain;
// charset=utf-8" becomes "text/plain"). The MIME writer must only ever see
// base types: passing a parameterized string to mime.FormatMediaType yields
// an empty value and a malformed Content-Type header.
func normalizeDraftAttachments(
	attachments []DraftAttachment,
) []DraftAttachment {
	normalized := make([]DraftAttachment, len(attachments))
	for index, attachment := range attachments {
		baseType, _, err := mime.ParseMediaType(
			strings.TrimSpace(attachment.MediaType),
		)
		if err != nil {
			// Unreachable after validateDraftAttachments; keep the value
			// only so failures surface at write time.
			baseType = attachment.MediaType
		}
		attachment.MediaType = strings.ToLower(baseType)
		attachment.Filename = strings.TrimSpace(attachment.Filename)
		normalized[index] = attachment
	}
	return normalized
}

func writeDraftMessage(
	message preparedDraftMessage,
	includeBCC bool,
) ([]byte, error) {
	var header messagemail.Header
	header.SetAddressList("From", message.from)
	header.SetAddressList("To", message.to)
	header.SetAddressList("Cc", message.cc)
	if includeBCC {
		header.SetAddressList("Bcc", message.bcc)
	}
	header.SetSubject(message.content.Subject)
	if len(message.content.Attachments) == 0 {
		switch message.content.BodyFormat {
		case "text":
			header.SetContentType("text/plain", map[string]string{
				"charset": "utf-8",
			})
		case "html":
			header.SetContentType("text/html", map[string]string{
				"charset": "utf-8",
			})
		}
		var output bytes.Buffer
		writer, err := messagemail.CreateSingleInlineWriter(&output, header)
		if err != nil {
			return nil, fmt.Errorf("create draft MIME writer: %w", err)
		}
		if _, err := writer.Write([]byte(message.content.Body)); err != nil {
			_ = writer.Close()
			return nil, fmt.Errorf("write draft MIME body: %w", err)
		}
		if err := writer.Close(); err != nil {
			return nil, fmt.Errorf("close draft MIME writer: %w", err)
		}
		return output.Bytes(), nil
	}
	return writeMultipartDraftMessage(header, message)
}

// writeMultipartDraftMessage emits multipart/mixed with the inline body as the
// first part followed by one base64 part per attachment. The header handed in
// already carries Bcc only for the IMAP draft projection; the SMTP projection
// never includes a Bcc header even though its envelope carries Bcc recipients.
func writeMultipartDraftMessage(
	header messagemail.Header,
	message preparedDraftMessage,
) ([]byte, error) {
	var output bytes.Buffer
	mixed := multipart.NewWriter(&output)
	header.Set("MIME-Version", "1.0")
	header.SetContentType("multipart/mixed", map[string]string{
		"boundary": mixed.Boundary(),
	})
	if err := writeRawHeader(&output, header); err != nil {
		return nil, err
	}
	bodyMediaType := "text/plain"
	if message.content.BodyFormat == "html" {
		bodyMediaType = "text/html"
	}
	bodyHeader := make(textproto.MIMEHeader)
	bodyHeader.Set("Content-Type", mime.FormatMediaType(
		bodyMediaType,
		map[string]string{"charset": "utf-8"},
	))
	bodyHeader.Set("Content-Transfer-Encoding", "quoted-printable")
	bodyPart, err := mixed.CreatePart(bodyHeader)
	if err != nil {
		return nil, fmt.Errorf("create draft MIME body part: %w", err)
	}
	quoted := quotedprintable.NewWriter(bodyPart)
	if _, err := quoted.Write([]byte(message.content.Body)); err != nil {
		_ = quoted.Close()
		return nil, fmt.Errorf("write draft MIME body: %w", err)
	}
	if err := quoted.Close(); err != nil {
		return nil, fmt.Errorf("close draft MIME body: %w", err)
	}
	for _, attachment := range message.content.Attachments {
		partHeader := make(textproto.MIMEHeader)
		partHeader.Set("Content-Type", mime.FormatMediaType(
			attachment.MediaType,
			map[string]string{"name": attachment.Filename},
		))
		partHeader.Set("Content-Disposition", mime.FormatMediaType(
			"attachment",
			map[string]string{"filename": attachment.Filename},
		))
		partHeader.Set("Content-Transfer-Encoding", "base64")
		part, err := mixed.CreatePart(partHeader)
		if err != nil {
			return nil, fmt.Errorf("create attachment MIME part: %w", err)
		}
		encoded := base64.NewEncoder(
			base64.StdEncoding,
			newCRLFWriter(part, 76),
		)
		if _, err := encoded.Write(attachment.Blob); err != nil {
			_ = encoded.Close()
			return nil, fmt.Errorf("encode attachment MIME part: %w", err)
		}
		if err := encoded.Close(); err != nil {
			return nil, fmt.Errorf("close attachment MIME part: %w", err)
		}
	}
	if err := mixed.Close(); err != nil {
		return nil, fmt.Errorf("close draft MIME writer: %w", err)
	}
	return output.Bytes(), nil
}

// writeRawHeader emits the already-encoded mail header fields followed by the
// blank line separating the header from the multipart body.
func writeRawHeader(
	output *bytes.Buffer,
	header messagemail.Header,
) error {
	fields := header.Fields()
	for fields.Next() {
		raw, err := fields.Raw()
		if err != nil || len(raw) == 0 {
			raw = []byte(fields.Key() + ": " + fields.Value() + "\r\n")
		}
		if _, err := output.Write(raw); err != nil {
			return fmt.Errorf("write draft MIME header: %w", err)
		}
	}
	if _, err := output.Write([]byte("\r\n")); err != nil {
		return fmt.Errorf("write draft MIME header: %w", err)
	}
	return nil
}

// crlfWriter inserts a CRLF after every lineLength output bytes so base64
// parts stay within the SMTP 998-octet line limit.
type crlfWriter struct {
	inner       io.Writer
	lineLength  int
	currentLine int
}

func newCRLFWriter(inner io.Writer, lineLength int) *crlfWriter {
	return &crlfWriter{inner: inner, lineLength: lineLength}
}

func (writer *crlfWriter) Write(chunk []byte) (int, error) {
	written := 0
	for len(chunk) > 0 {
		remaining := writer.lineLength - writer.currentLine
		if remaining == 0 {
			if _, err := writer.inner.Write([]byte("\r\n")); err != nil {
				return written, err
			}
			writer.currentLine = 0
			remaining = writer.lineLength
		}
		if len(chunk) < remaining {
			remaining = len(chunk)
		}
		count, err := writer.inner.Write(chunk[:remaining])
		written += count
		if err != nil {
			return written, err
		}
		chunk = chunk[remaining:]
		writer.currentLine += remaining
	}
	return written, nil
}

func parseAddresses(values []string) ([]*messagemail.Address, error) {
	if len(values) > maxAddresses {
		return nil, fmt.Errorf("draft address list exceeded limit")
	}
	result := make([]*messagemail.Address, 0, len(values))
	for _, value := range values {
		trimmed := strings.TrimSpace(value)
		if trimmed == "" {
			continue
		}
		parsed, err := netmail.ParseAddress(trimmed)
		if err != nil || parsed.Address != trimmed {
			return nil, fmt.Errorf("draft address is invalid")
		}
		result = append(result, &messagemail.Address{
			Address: parsed.Address,
		})
	}
	return result, nil
}

type bodyPart struct {
	Path      []int
	MediaType string
	Encoding  string
	Size      uint32
	Filename  string
}

func selectBodyPart(
	structure *imap.BodyStructure,
) (bodyPart, []Attachment, error) {
	if structure == nil {
		return bodyPart{}, nil, &ProtocolError{
			Class:     ErrorFailed,
			Operation: "parse MIME tree",
		}
	}
	type pendingPart struct {
		path []int
		part *imap.BodyStructure
	}
	pending := []pendingPart{{path: []int{1}, part: structure}}
	if len(structure.Parts) > 0 {
		pending[0].path = nil
	}
	var plain bodyPart
	var html bodyPart
	attachments := make([]Attachment, 0)
	visited := 0
	for len(pending) > 0 {
		current := pending[0]
		pending = pending[1:]
		visited++
		if visited > maxBodyParts {
			return bodyPart{}, nil,
				&ProtocolError{Class: ErrorLimit, Operation: "parse MIME tree"}
		}
		part := current.part
		for index, child := range part.Parts {
			path := append([]int(nil), current.path...)
			path = append(path, index+1)
			pending = append(pending, pendingPart{
				path: path,
				part: child,
			})
		}
		if len(part.Parts) > 0 {
			continue
		}
		mediaType := strings.ToLower(part.MIMEType + "/" + part.MIMESubType)
		filename, _ := part.Filename()
		isAttachment := strings.EqualFold(part.Disposition, "attachment") ||
			filename != ""
		if isAttachment {
			if len(attachments) >= maxAttachments {
				return bodyPart{}, nil,
					&ProtocolError{
						Class:     ErrorLimit,
						Operation: "list attachments",
					}
			}
			boundedFilename := boundedUTF8(filename, maxFilenameBytes)
			boundedMediaType := boundedUTF8(
				mediaType,
				maxMediaTypeBytes,
			)
			attachments = append(attachments, Attachment{
				Part:      append([]int(nil), current.path...),
				Filename:  boundedFilename,
				MediaType: boundedMediaType,
				SizeBytes: part.Size,
				Truncated: boundedFilename != filename ||
					boundedMediaType != mediaType,
			})
			continue
		}
		candidate := bodyPart{
			Path:      append([]int(nil), current.path...),
			MediaType: mediaType,
			Encoding:  part.Encoding,
			Size:      part.Size,
		}
		switch mediaType {
		case "text/plain":
			if plain.MediaType == "" {
				plain = candidate
			}
		case "text/html":
			if html.MediaType == "" {
				html = candidate
			}
		}
	}
	if plain.MediaType != "" {
		return plain, attachments, nil
	}
	return html, attachments, nil
}

// decodeTransferEncoding returns a reader that undoes the part's
// Content-Transfer-Encoding. Unknown encodings are rejected.
func decodeTransferEncoding(
	part bodyPart,
	raw []byte,
	operation string,
) (io.Reader, error) {
	var reader io.Reader = bytes.NewReader(raw)
	switch strings.ToUpper(strings.TrimSpace(part.Encoding)) {
	case "", "7BIT", "8BIT", "BINARY":
	case "BASE64":
		reader = base64.NewDecoder(base64.StdEncoding, reader)
	case "QUOTED-PRINTABLE":
		reader = quotedprintable.NewReader(reader)
	default:
		return nil, &ProtocolError{
			Class:     ErrorFailed,
			Operation: operation,
		}
	}
	return reader, nil
}

func decodeBodyPart(
	part bodyPart,
	raw []byte,
) (Body, error) {
	const operation = "decode message body"
	if len(raw) > maxDraftBodyBytes {
		return Body{}, &ProtocolError{
			Class:     ErrorLimit,
			Operation: "read message body",
		}
	}
	reader, err := decodeTransferEncoding(part, raw, operation)
	if err != nil {
		return Body{}, err
	}
	decoded, err := io.ReadAll(io.LimitReader(reader, maxDraftBodyBytes+1))
	if err != nil {
		return Body{}, &ProtocolError{
			Class:     ErrorFailed,
			Operation: operation,
			Cause:     err,
		}
	}
	truncated := part.Size > maxDraftBodyBytes ||
		len(decoded) > maxDraftBodyBytes
	if len(decoded) > maxDraftBodyBytes {
		decoded = decoded[:maxDraftBodyBytes]
	}
	content := string(bytes.ToValidUTF8(decoded, []byte("\uFFFD")))
	format := "text"
	if part.MediaType == "text/html" {
		content, err = htmlToPlainText(content)
		if err != nil {
			return Body{}, &ProtocolError{
				Class:     ErrorFailed,
				Operation: "sanitize message HTML",
				Cause:     err,
			}
		}
		content = "<pre>" + stdhtml.EscapeString(content) + "</pre>"
		format = "sanitized_html"
	}
	boundedContent := boundedUTF8(content, maxDraftBodyBytes)
	truncated = truncated || boundedContent != content
	return Body{
		Format:    format,
		Content:   boundedContent,
		Truncated: truncated,
	}, nil
}

// decodeDraftBodyPart returns the EXACT bounded body of a draft part. Draft
// bodies are frozen input the server itself wrote, so they are never HTML
// sanitized or rewritten: send must submit byte-for-byte what the draft says.
// Anything over the decoded budget is marked truncated so callers fail closed
// instead of sending partial content.
func decodeDraftBodyPart(
	part bodyPart,
	raw []byte,
) (Body, error) {
	const operation = "decode draft body"
	reader, err := decodeTransferEncoding(part, raw, operation)
	if err != nil {
		return Body{}, err
	}
	decoded, err := io.ReadAll(io.LimitReader(reader, maxDraftBodyBytes+1))
	if err != nil {
		return Body{}, &ProtocolError{
			Class:     ErrorFailed,
			Operation: operation,
			Cause:     err,
		}
	}
	format := "text"
	if part.MediaType == "text/html" {
		format = "html"
	}
	return Body{
		Format:    format,
		Content:   string(decoded),
		Truncated: len(decoded) > maxDraftBodyBytes,
	}, nil
}

// decodeAttachmentPart returns the raw decoded bytes of one bounded MIME part.
// Unlike decodeBodyPart it never sanitizes or truncates: an over-budget
// attachment is a hard limit error so callers fail closed instead of sending
// or materializing partial content.
func decodeAttachmentPart(
	part bodyPart,
	raw []byte,
) ([]byte, error) {
	var reader io.Reader = bytes.NewReader(raw)
	switch strings.ToUpper(strings.TrimSpace(part.Encoding)) {
	case "", "7BIT", "8BIT", "BINARY":
	case "BASE64":
		reader = base64.NewDecoder(base64.StdEncoding, reader)
	case "QUOTED-PRINTABLE":
		reader = quotedprintable.NewReader(reader)
	default:
		return nil, &ProtocolError{
			Class:     ErrorFailed,
			Operation: "decode attachment part",
		}
	}
	decoded, err := io.ReadAll(io.LimitReader(reader, MaxAttachmentPartBytes+1))
	if err != nil {
		return nil, &ProtocolError{
			Class:     ErrorFailed,
			Operation: "decode attachment part",
			Cause:     err,
		}
	}
	if len(decoded) > MaxAttachmentPartBytes {
		return nil, &ProtocolError{
			Class:     ErrorLimit,
			Operation: "read attachment part",
		}
	}
	if len(decoded) == 0 {
		return nil, &ProtocolError{
			Class:     ErrorFailed,
			Operation: "read attachment part",
		}
	}
	return decoded, nil
}

func htmlToPlainText(raw string) (string, error) {
	tokenizer := xhtml.NewTokenizer(strings.NewReader(raw))
	var output strings.Builder
	for {
		tokenType := tokenizer.Next()
		switch tokenType {
		case xhtml.ErrorToken:
			if tokenizer.Err() == io.EOF {
				return strings.TrimSpace(output.String()), nil
			}
			return "", tokenizer.Err()
		case xhtml.TextToken:
			if output.Len() < maxDraftBodyBytes {
				remaining := maxDraftBodyBytes - output.Len()
				output.WriteString(boundedUTF8(
					string(tokenizer.Text()),
					remaining,
				))
			}
		case xhtml.StartTagToken, xhtml.EndTagToken,
			xhtml.SelfClosingTagToken:
			name, _ := tokenizer.TagName()
			switch string(name) {
			case "br", "p", "div", "li", "tr", "blockquote", "pre":
				if output.Len() < maxDraftBodyBytes {
					output.WriteByte('\n')
				}
			}
		}
	}
}

func boundedUTF8(value string, maxBytes int) string {
	if maxBytes <= 0 {
		return ""
	}
	if len(value) <= maxBytes {
		return value
	}
	value = value[:maxBytes]
	for len(value) > 0 && !utf8.ValidString(value) {
		value = value[:len(value)-1]
	}
	return value
}