package protocol

import (
	"context"
	"fmt"
	"log/slog"
	"net"
	"strings"
	"time"

	"github.com/emersion/go-imap"
	imapclient "github.com/emersion/go-imap/client"

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

type IMAPClient struct {
	options ClientOptions
}

func NewIMAPClient(options ClientOptions) *IMAPClient {
	return &IMAPClient{options: options}
}

func (client *IMAPClient) Probe(
	ctx context.Context,
	endpoint Endpoint,
) (Capabilities, error) {
	session, connection, err := client.open(ctx, endpoint)
	if err != nil {
		return Capabilities{}, classifyReadError("open IMAP session", err)
	}
	defer connection.Close()
	defer session.Terminate()

	capabilities, err := discoverCapabilities(session)
	if err != nil {
		return Capabilities{},
			classifyReadError("probe IMAP capabilities", err)
	}
	_ = session.Logout()
	return capabilities, nil
}

func discoverCapabilities(
	session *imapclient.Client,
) (Capabilities, error) {
	values, err := readCapabilityValues(session)
	if err != nil {
		return Capabilities{}, err
	}
	capabilities := Capabilities{
		Values:    values,
		Mailboxes: map[string]string{FolderInbox: imap.InboxName},
	}
	mailboxes := make(chan *imap.MailboxInfo, 16)
	listErr := make(chan error, 1)
	logging.Go(slog.Default(), func() {
		listErr <- session.List("", "*", mailboxes)
	})
	count := 0
	var validationErr error
	for mailbox := range mailboxes {
		if validationErr != nil {
			continue
		}
		count++
		if count > 256 ||
			len(mailbox.Name) == 0 ||
			len(mailbox.Name) > maxMailboxNameBytes ||
			len(mailbox.Attributes) > 32 {
			validationErr = &ProtocolError{
				Class:     ErrorLimit,
				Operation: "list IMAP mailboxes",
			}
			continue
		}
		noSelect := false
		for _, attribute := range mailbox.Attributes {
			if len(attribute) > 64 {
				validationErr = &ProtocolError{
					Class:     ErrorLimit,
					Operation: "list IMAP mailboxes",
				}
				break
			}
			if strings.EqualFold(attribute, imap.NoSelectAttr) {
				noSelect = true
			}
		}
		if validationErr != nil {
			continue
		}
		if noSelect {
			continue
		}
		for _, attribute := range mailbox.Attributes {
			folder := specialUseFolder(attribute)
			if folder == "" {
				continue
			}
			existing := capabilities.Mailboxes[folder]
			if existing != "" && existing != mailbox.Name {
				validationErr = &ProtocolError{
					Class:     ErrorUnsupported,
					Operation: "map IMAP SPECIAL-USE mailbox",
				}
				break
			}
			capabilities.Mailboxes[folder] = mailbox.Name
		}
	}
	if err := <-listErr; err != nil {
		return Capabilities{},
			fmt.Errorf("list IMAP SPECIAL-USE folders: %w", err)
	}
	if validationErr != nil {
		return Capabilities{}, validationErr
	}
	capabilities.DraftMailbox = capabilities.Mailboxes[FolderDrafts]
	return capabilities, nil
}

func readCapabilityValues(
	session *imapclient.Client,
) (map[string]bool, error) {
	rawCapabilities, err := session.Capability()
	if err != nil {
		return nil, fmt.Errorf("read IMAP capabilities: %w", err)
	}
	if len(rawCapabilities) > 128 {
		return nil, &ProtocolError{
			Class:     ErrorLimit,
			Operation: "read IMAP capabilities",
		}
	}
	values := make(map[string]bool, len(rawCapabilities))
	for name, supported := range rawCapabilities {
		if !supported {
			continue
		}
		name = strings.ToUpper(name)
		if len(name) == 0 || len(name) > 64 {
			return nil, &ProtocolError{
				Class:     ErrorLimit,
				Operation: "read IMAP capabilities",
			}
		}
		values[name] = true
	}
	return values, nil
}

func specialUseFolder(attribute string) string {
	switch {
	case strings.EqualFold(attribute, imap.AllAttr):
		return FolderAll
	case strings.EqualFold(attribute, imap.ArchiveAttr):
		return FolderArchive
	case strings.EqualFold(attribute, imap.DraftsAttr):
		return FolderDrafts
	case strings.EqualFold(attribute, imap.FlaggedAttr):
		return FolderFlagged
	case strings.EqualFold(attribute, imap.JunkAttr):
		return FolderJunk
	case strings.EqualFold(attribute, imap.SentAttr):
		return FolderSent
	case strings.EqualFold(attribute, imap.TrashAttr):
		return FolderTrash
	default:
		return ""
	}
}

func (client *IMAPClient) open(
	ctx context.Context,
	endpoint Endpoint,
) (*imapclient.Client, net.Conn, error) {
	handle, err := client.openBudgeted(
		ctx,
		endpoint,
		defaultProtocolBudget(client.options),
	)
	if err != nil {
		return nil, nil, err
	}
	return handle.session, handle.connection, nil
}

// imapSessionHandle bundles one pinned IMAP session with its operation
// ceiling and budgeted connection so error paths can classify exhaustion.
type imapSessionHandle struct {
	session    *imapclient.Client
	connection net.Conn
	ceiling    time.Time
	budget     *budgetConn
}

// limitError maps an error to the typed limit classification when the
// observable budget state shows exhaustion; otherwise it returns nil so the
// caller falls back to the generic classification.
func (handle *imapSessionHandle) limitError(operation string) error {
	if handle == nil || !handle.budget.budgetExceeded() {
		return nil
	}
	return &ProtocolError{
		Class:     ErrorLimit,
		Operation: operation,
		Cause:     errProtocolBudget,
	}
}

// openBudgeted returns the session handle carrying the absolute operation
// deadline ceiling and the budgeted connection so bounded bulk commands can
// size their per-command timeout and classify exhaustion.
func (client *IMAPClient) openBudgeted(
	ctx context.Context,
	endpoint Endpoint,
	budget protocolBudget,
) (*imapSessionHandle, error) {
	pinned, err := openPinned(
		ctx,
		client.options,
		transport.ProtocolIMAP,
		endpoint,
		budget,
	)
	if err != nil {
		return nil, err
	}
	handle := &imapSessionHandle{connection: pinned.Conn}
	if budgeted, ok := pinned.Conn.(*budgetConn); ok {
		handle.budget = budgeted
	}
	connection := pinned.Conn
	if pinned.TLSMode == transport.TLSModeImplicit {
		connection, err = implicitTLS(ctx, pinned)
		if err != nil {
			if limitErr := handle.limitError("open IMAP session"); limitErr != nil {
				return nil, limitErr
			}
			return nil, err
		}
		handle.connection = connection
	}
	session, err := imapclient.New(connection)
	if err != nil {
		_ = connection.Close()
		if limitErr := handle.limitError("open IMAP session"); limitErr != nil {
			return nil, limitErr
		}
		return nil, fmt.Errorf("open IMAP session: %w", err)
	}
	handle.session = session
	session.ErrorLog = discardLogger{}
	metadataTimeout := client.options.CommandTimeout
	if metadataTimeout <= 0 {
		metadataTimeout = defaultCommandTimeout
	}
	session.Timeout = metadataTimeout
	ceiling := time.Now().Add(budget.deadline)
	if contextDeadline, ok := ctx.Deadline(); ok &&
		contextDeadline.Before(ceiling) {
		ceiling = contextDeadline
	}
	handle.ceiling = ceiling
	if pinned.TLSMode == transport.TLSModeStartTLS {
		supported, supportErr := session.SupportStartTLS()
		if supportErr != nil {
			_ = session.Terminate()
			_ = connection.Close()
			return nil,
				fmt.Errorf("read IMAP STARTTLS capability: %w", supportErr)
		}
		if !supported {
			_ = session.Terminate()
			_ = connection.Close()
			return nil,
				fmt.Errorf("IMAP STARTTLS is required")
		}
		if err := session.StartTLS(pinned.TLSConfig.Clone()); err != nil {
			_ = session.Terminate()
			_ = connection.Close()
			return nil,
				fmt.Errorf("verify IMAP STARTTLS: %w", err)
		}
	}
	capabilities, err := readCapabilityValues(session)
	if err != nil {
		_ = session.Terminate()
		_ = connection.Close()
		return nil, fmt.Errorf(
			"read post-TLS IMAP capabilities: %w",
			err,
		)
	}
	if !capabilities["IMAP4REV1"] && !capabilities["IMAP4REV2"] {
		_ = session.Terminate()
		_ = connection.Close()
		return nil, fmt.Errorf(
			"post-TLS IMAP capabilities are unsupported",
		)
	}
	if err := session.Login(
		endpoint.Username,
		endpoint.Password,
	); err != nil {
		_ = session.Terminate()
		_ = connection.Close()
		return nil,
			fmt.Errorf("authenticate IMAP session: %w", err)
	}
	return handle, nil
}

type discardLogger struct{}

func (discardLogger) Printf(string, ...interface{}) {}
func (discardLogger) Println(...interface{})        {}