package protocol

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/extension/providers/mail/transport"
)

const defaultCommandTimeout = 15 * time.Second

// Bounded bulk transfers get their own operation-level budgets and deadlines;
// metadata-sized commands keep the tight defaults. Every budget still counts
// TLS record bytes because the byte counter wraps the raw connection below
// the TLS layer, so each limit carries handshake and record overhead slack.
const (
	maxProtocolBytes = 2 * 1024 * 1024

	protocolOverheadBytes = 1 << 20
	// bulkTransferTimeout bounds one attachment-sized IMAP/SMTP transfer on a
	// slow link while staying finite.
	bulkTransferTimeout = 10 * time.Minute
	// draftBodyTransferTimeout bounds one bounded draft body fetch.
	draftBodyTransferTimeout = 2 * time.Minute
)

var errProtocolBudget = errors.New("mail protocol byte budget exceeded")

// protocolBudget is the per-operation read/write allowance and deadline
// installed on one pinned connection.
type protocolBudget struct {
	read     int64
	write    int64
	deadline time.Duration
}

// defaultProtocolBudget fits every metadata-sized command including probes,
// searches, envelope fetches, and draft revision checks.
func defaultProtocolBudget(options ClientOptions) protocolBudget {
	deadline := options.CommandTimeout
	if deadline <= 0 {
		deadline = defaultCommandTimeout
	}
	return protocolBudget{
		read:     maxProtocolBytes,
		write:    maxProtocolBytes,
		deadline: deadline,
	}
}

// attachmentFetchBudget fits one bounded encoded attachment part plus the
// session handshake and protocol chatter.
func attachmentFetchBudget(options ClientOptions) protocolBudget {
	return protocolBudget{
		read:     maxAttachmentFetchBytes + protocolOverheadBytes,
		write:    maxProtocolBytes,
		deadline: effectiveBulkTimeout(options),
	}
}

// draftBodyFetchBudget fits one bounded encoded draft body plus handshake
// and protocol chatter.
func draftBodyFetchBudget(options ClientOptions) protocolBudget {
	return protocolBudget{
		read:     maxDraftBodyFetchBytes + protocolOverheadBytes,
		write:    maxProtocolBytes,
		deadline: effectiveDraftBodyTimeout(options),
	}
}

// draftAppendBudget fits appending one fully encoded draft message.
func draftAppendBudget(options ClientOptions, messageBytes int) protocolBudget {
	write := int64(messageBytes) + protocolOverheadBytes
	if write < maxProtocolBytes {
		write = maxProtocolBytes
	}
	return protocolBudget{
		read:     maxProtocolBytes,
		write:    write,
		deadline: effectiveBulkTimeout(options),
	}
}

// smtpSendBudget fits submitting one fully encoded message.
func smtpSendBudget(options ClientOptions) protocolBudget {
	return protocolBudget{
		read:     maxProtocolBytes,
		write:    maxSMTPMessageBytes + protocolOverheadBytes,
		deadline: effectiveBulkTimeout(options),
	}
}

func effectiveBulkTimeout(options ClientOptions) time.Duration {
	if options.BulkTransferTimeout > 0 {
		return options.BulkTransferTimeout
	}
	return bulkTransferTimeout
}

func effectiveDraftBodyTimeout(options ClientOptions) time.Duration {
	if options.DraftBodyTransferTimeout > 0 {
		return options.DraftBodyTransferTimeout
	}
	return draftBodyTransferTimeout
}

func openPinned(
	ctx context.Context,
	options ClientOptions,
	protocol string,
	endpoint Endpoint,
	budget protocolBudget,
) (transport.PinnedConnection, error) {
	if options.Connector == nil {
		return transport.PinnedConnection{},
			fmt.Errorf("mail protocol connector is not configured")
	}
	connection, err := options.Connector.ResolveAndDial(
		ctx,
		transport.Target{
			Protocol: protocol,
			Host:     endpoint.Host,
			Port:     endpoint.Port,
			TLSMode:  endpoint.TLSMode,
		},
	)
	if err != nil {
		return transport.PinnedConnection{}, err
	}
	deadline := time.Now().Add(budget.deadline)
	if contextDeadline, ok := ctx.Deadline(); ok &&
		contextDeadline.Before(deadline) {
		deadline = contextDeadline
	}
	if err := connection.Conn.SetDeadline(deadline); err != nil {
		_ = connection.Conn.Close()
		return transport.PinnedConnection{},
			fmt.Errorf("set mail protocol deadline: %w", err)
	}
	connection.Conn = &budgetConn{
		Conn:            connection.Conn,
		readRemaining:   budget.read,
		writeRemaining:  budget.write,
		deadlineCeiling: deadline,
	}
	return connection, nil
}

func implicitTLS(
	ctx context.Context,
	connection transport.PinnedConnection,
) (net.Conn, error) {
	if connection.TLSConfig == nil {
		_ = connection.Conn.Close()
		return nil, fmt.Errorf("mail TLS configuration is unavailable")
	}
	tlsConnection := tls.Client(
		connection.Conn,
		connection.TLSConfig.Clone(),
	)
	if err := tlsConnection.HandshakeContext(ctx); err != nil {
		_ = connection.Conn.Close()
		return nil, fmt.Errorf("verify implicit mail TLS: %w", err)
	}
	return tlsConnection, nil
}

type budgetConn struct {
	net.Conn

	readMu         sync.Mutex
	readRemaining  int64
	writeMu        sync.Mutex
	writeRemaining int64
	exhausted      atomic.Bool
	// deadlineCeiling is the absolute operation ceiling. Third-party protocol
	// clients reset connection deadlines before every command, so every
	// setter clamps its requested value to this ceiling: no library call can
	// extend a bounded transfer past its operation deadline.
	deadlineMu      sync.Mutex
	deadlineCeiling time.Time
}

func (connection *budgetConn) clampDeadline(requested time.Time) time.Time {
	connection.deadlineMu.Lock()
	defer connection.deadlineMu.Unlock()
	if connection.deadlineCeiling.IsZero() ||
		requested.Before(connection.deadlineCeiling) {
		return requested
	}
	return connection.deadlineCeiling
}

// setClampedDeadline applies one deadline kind through the ceiling clamp so
// third-party protocol clients cannot extend a bounded operation.
func (connection *budgetConn) setClampedDeadline(
	kind func(time.Time) error,
	deadline time.Time,
) error {
	return kind(connection.clampDeadline(deadline))
}

// clampDeadlineToCeiling returns the deadline a caller may actually apply.
// A zero value would CLEAR the deadline and lift the operation ceiling, and
// any value beyond the ceiling would extend it; both clamp down to the
// ceiling so no third-party protocol call can remove the bound.
func (connection *budgetConn) clampDeadlineToCeiling(deadline time.Time) time.Time {
	connection.deadlineMu.Lock()
	defer connection.deadlineMu.Unlock()
	if connection.deadlineCeiling.IsZero() {
		return deadline
	}
	if deadline.IsZero() || deadline.After(connection.deadlineCeiling) {
		return connection.deadlineCeiling
	}
	return deadline
}

func (connection *budgetConn) SetDeadline(deadline time.Time) error {
	return connection.Conn.SetDeadline(
		connection.clampDeadlineToCeiling(deadline),
	)
}

func (connection *budgetConn) SetReadDeadline(deadline time.Time) error {
	return connection.Conn.SetReadDeadline(
		connection.clampDeadlineToCeiling(deadline),
	)
}

func (connection *budgetConn) SetWriteDeadline(deadline time.Time) error {
	return connection.Conn.SetWriteDeadline(
		connection.clampDeadlineToCeiling(deadline),
	)
}

func (connection *budgetConn) Read(payload []byte) (int, error) {
	connection.readMu.Lock()
	defer connection.readMu.Unlock()
	if connection.readRemaining <= 0 {
		connection.exhausted.Store(true)
		return 0, errProtocolBudget
	}
	if int64(len(payload)) > connection.readRemaining {
		payload = payload[:connection.readRemaining]
	}
	count, err := connection.Conn.Read(payload)
	connection.readRemaining -= int64(count)
	return count, err
}

func (connection *budgetConn) Write(payload []byte) (int, error) {
	connection.writeMu.Lock()
	defer connection.writeMu.Unlock()
	if connection.writeRemaining <= 0 {
		connection.exhausted.Store(true)
		return 0, errProtocolBudget
	}
	truncated := false
	if int64(len(payload)) > connection.writeRemaining {
		payload = payload[:connection.writeRemaining]
		truncated = true
	}
	count, err := connection.Conn.Write(payload)
	connection.writeRemaining -= int64(count)
	if err == nil && (count < len(payload) || truncated) {
		connection.exhausted.Store(true)
		err = errProtocolBudget
	}
	return count, err
}

// budgetExceeded reports whether the byte budget ended the stream. Protocol
// libraries swallow the underlying error sentinel, so exhaustion is tracked as
// observable state and mapped to a typed limit error by the callers.
func (connection *budgetConn) budgetExceeded() bool {
	return connection != nil && connection.exhausted.Load()
}