package protocol

import (
	"encoding/base64"
	"net"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func attachmentRef() MessageRef {
	return MessageRef{Mailbox: "INBOX", UIDValidity: 777, UID: 9}
}

func TestIMAPGetAttachmentFetchesExactPartAndDecodes(t *testing.T) {
	t.Parallel()

	decoded := []byte("%PDF-1.4 scripted-attachment")
	encoded := []byte(base64.StdEncoding.EncodeToString(decoded))
	client, _, done := newHighLevelIMAPClient(t, highLevelIMAPScript{
		attachmentBody:    encoded,
		attachmentRawSize: uint32(len(encoded)),
	})
	content, err := client.GetAttachment(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
		attachmentRef(),
		[]int{2},
	)
	require.NoError(t, err)
	require.NoError(t, <-done)
	require.Equal(t, []int{2}, content.Part)
	require.Equal(t, "report.pdf", content.Filename)
	require.Equal(t, "application/pdf", content.MediaType)
	require.Equal(t, decoded, content.Blob)
}

func TestIMAPGetAttachmentRejectsSilentlyTruncatedPart(t *testing.T) {
	t.Parallel()

	// The declared encoded size is larger than what the server actually
	// returns; a truncated base64 block without padding would otherwise
	// decode silently into fewer bytes.
	decoded := []byte(strings.Repeat("A", 40))
	encoded := []byte(base64.StdEncoding.EncodeToString(decoded))
	client, _, done := newHighLevelIMAPClient(t, highLevelIMAPScript{
		attachmentBody:    encoded[:len(encoded)-16],
		attachmentRawSize: uint32(len(encoded)),
	})
	_, err := client.GetAttachment(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
		attachmentRef(),
		[]int{2},
	)
	require.Error(t, err)
	require.Equal(t, ErrorLimit, ErrorClass(err))
	require.NoError(t, <-done)
}

func TestIMAPGetAttachmentRejectsOversizedEncodedPart(t *testing.T) {
	t.Parallel()

	client, _, done := newHighLevelIMAPClient(t, highLevelIMAPScript{
		attachmentBody:    []byte("x"),
		attachmentRawSize: maxAttachmentFetchBytes + 1,
	})
	_, err := client.GetAttachment(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
		attachmentRef(),
		[]int{2},
	)
	require.Error(t, err)
	require.Equal(t, ErrorLimit, ErrorClass(err))
	require.NoError(t, <-done)
}

// TestDecodeAttachmentPartRejectsDecodedOverBudget pins the decoded-layer
// budget at unit level: a complete encoded part that decodes beyond the
// per-part limit must fail closed instead of returning truncated content.
func TestDecodeAttachmentPartRejectsDecodedOverBudget(t *testing.T) {
	t.Parallel()

	decoded := make([]byte, MaxAttachmentPartBytes+16)
	for index := range decoded {
		decoded[index] = 'A'
	}
	raw := []byte(base64.StdEncoding.EncodeToString(decoded))
	blob, err := decodeAttachmentPart(bodyPart{
		Path:     []int{2},
		Encoding: "BASE64",
		Size:     uint32(len(raw)),
	}, raw)
	require.Error(t, err)
	require.Nil(t, blob)
	require.Equal(t, ErrorLimit, ErrorClass(err))
}

func TestIMAPGetAttachmentRejectsNonAttachmentPart(t *testing.T) {
	t.Parallel()

	client, _, done := newHighLevelIMAPClient(t, highLevelIMAPScript{
		attachmentBody:    []byte("Hello"),
		attachmentRawSize: 5,
	})
	_, err := client.GetAttachment(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
		attachmentRef(),
		[]int{1},
	)
	require.Error(t, err)
	require.Equal(t, ErrorNotFound, ErrorClass(err))
	require.NoError(t, <-done)
}

// TestBudgetConnExhaustionClassifiesAsLimit pins the observable budget-exhaustion
// seam at unit level: protocol libraries swallow the underlying sentinel, so
// the connection tracks exhaustion as state and the session handle maps it to
// the typed limit classification.
func TestBudgetConnExhaustionClassifiesAsLimit(t *testing.T) {
	t.Parallel()

	client, server := net.Pipe()
	defer client.Close()
	defer server.Close()
	go func() {
		// Feed more bytes than the read budget allows, then drain.
		_, _ = server.Write(make([]byte, 64))
		scratch := make([]byte, 1024)
		for {
			if _, err := server.Read(scratch); err != nil {
				return
			}
		}
	}()
	ceiling := time.Now().Add(time.Minute)
	connection := &budgetConn{
		Conn:            client,
		readRemaining:   8,
		writeRemaining:  8,
		deadlineCeiling: ceiling,
	}
	payload := make([]byte, 16)
	count, err := connection.Read(payload)
	require.NoError(t, err)
	require.Equal(t, 8, count)
	// The next read exhausts the budget.
	_, err = connection.Read(payload)
	require.Error(t, err)
	require.ErrorIs(t, err, errProtocolBudget)

	handle := &imapSessionHandle{budget: connection}
	limitErr := handle.limitError("fetch attachment part")
	require.Error(t, limitErr)
	require.Equal(t, ErrorLimit, ErrorClass(limitErr))

	// A zero (clearing) or over-ceiling deadline request clamps down to the
	// absolute ceiling instead of lifting the bound.
	require.Equal(t, ceiling, connection.clampDeadlineToCeiling(time.Time{}))
	far := time.Now().Add(time.Hour)
	require.Equal(t, ceiling, connection.clampDeadlineToCeiling(far))
	near := time.Now().Add(time.Second)
	require.Equal(t, near, connection.clampDeadlineToCeiling(near))
}
