gpt-plus-gpt/pkg/provider/email/mailgateway.go

package email

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"regexp"
	"strings"
	"time"
)

var codeRegexp = regexp.MustCompile(`\b(\d{6})\b`)

// MailGatewayProvider implements EmailProvider using the Mail Gateway API.
// API docs: https://regmail.zhengmi.org/
type MailGatewayProvider struct {
	baseURL  string
	apiKey   string
	provider string
}

// NewMailGateway creates a new MailGatewayProvider.
func NewMailGateway(baseURL, apiKey, provider string) *MailGatewayProvider {
	return &MailGatewayProvider{
		baseURL:  strings.TrimRight(baseURL, "/"),
		apiKey:   apiKey,
		provider: provider,
	}
}

type createMailboxResponse struct {
	OK    bool   `json:"ok"`
	Error string `json:"error,omitempty"`
	Data  struct {
		MailboxID string `json:"mailbox_id"`
		Email     string `json:"email"`
		Provider  string `json:"provider"`
	} `json:"data"`
}

type emailEntry struct {
	ID         string `json:"id"`
	From       string `json:"from"`
	Subject    string `json:"subject"`
	Content    string `json:"content"`
	Body       string `json:"body"`
	Text       string `json:"text"`
	ReceivedAt string `json:"received_at"`
}

type listEmailsResponse struct {
	OK    bool   `json:"ok"`
	Error string `json:"error,omitempty"`
	Data  struct {
		Emails []emailEntry `json:"emails"`
	} `json:"data"`
}

type verificationResponse struct {
	OK    bool            `json:"ok"`
	Error string          `json:"error,omitempty"`
	Data  json.RawMessage `json:"data"`
}

func (m *MailGatewayProvider) doRequest(req *http.Request) ([]byte, error) {
	req.Header.Set("X-API-Key", m.apiKey)

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("request %s: %w", req.URL.Path, err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("read response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("%s returned %d: %s", req.URL.Path, resp.StatusCode, string(body))
	}

	return body, nil
}

// CreateMailbox creates a temporary mailbox via the Mail Gateway API.
// POST /api/v1/mailboxes {"provider": "gptmail"}
func (m *MailGatewayProvider) CreateMailbox(ctx context.Context) (string, string, error) {
	payload := fmt.Sprintf(`{"provider":"%s"}`, m.provider)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, m.baseURL+"/api/v1/mailboxes", strings.NewReader(payload))
	if err != nil {
		return "", "", fmt.Errorf("build create mailbox request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

	body, err := m.doRequest(req)
	if err != nil {
		return "", "", fmt.Errorf("create mailbox: %w", err)
	}

	var result createMailboxResponse
	if err := json.Unmarshal(body, &result); err != nil {
		return "", "", fmt.Errorf("parse create mailbox response: %w", err)
	}

	if !result.OK {
		return "", "", fmt.Errorf("create mailbox failed: %s", result.Error)
	}

	log.Printf("[email] mailbox created: %s (provider: %s, id: %s)", result.Data.Email, result.Data.Provider, result.Data.MailboxID)
	return result.Data.Email, result.Data.MailboxID, nil
}

// WaitForVerificationCode polls for an OTP code, ignoring emails that existed before notBefore.
// If notBefore is zero, all emails are considered (for fresh mailboxes).
func (m *MailGatewayProvider) WaitForVerificationCode(ctx context.Context, mailboxID string, timeout time.Duration, notBefore time.Time) (string, error) {
	// Snapshot existing email IDs so we can skip them when filtering
	knownIDs := make(map[string]bool)
	if !notBefore.IsZero() {
		existing, _ := m.ListEmailIDs(ctx, mailboxID)
		for _, id := range existing {
			knownIDs[id] = true
		}
		log.Printf("[email] snapshotted %d existing emails to skip", len(knownIDs))
	}

	deadline := time.Now().Add(timeout)
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	for attempt := 1; ; attempt++ {
		log.Printf("[email] polling for verification code (attempt %d)", attempt)

		code, err := m.pollNewEmails(ctx, mailboxID, knownIDs)
		if err == nil && code != "" {
			log.Printf("[email] verification code found: %s", code)
			return code, nil
		}

		if time.Now().After(deadline) {
			return "", fmt.Errorf("timeout waiting for verification code after %s", timeout)
		}

		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case <-ticker.C:
		}
	}
}

// ListEmailIDs returns the IDs of all current emails in the mailbox.
func (m *MailGatewayProvider) ListEmailIDs(ctx context.Context, mailboxID string) ([]string, error) {
	reqURL := fmt.Sprintf("%s/api/v1/emails?mailbox_id=%s&limit=20", m.baseURL, mailboxID)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil)
	if err != nil {
		return nil, err
	}
	body, err := m.doRequest(req)
	if err != nil {
		return nil, err
	}
	var result listEmailsResponse
	if err := json.Unmarshal(body, &result); err != nil {
		return nil, err
	}
	var ids []string
	for _, e := range result.Data.Emails {
		ids = append(ids, e.ID)
	}
	return ids, nil
}

// pollNewEmails lists emails and returns a code only from emails NOT in knownIDs.
func (m *MailGatewayProvider) pollNewEmails(ctx context.Context, mailboxID string, knownIDs map[string]bool) (string, error) {
	reqURL := fmt.Sprintf("%s/api/v1/emails?mailbox_id=%s&limit=10", m.baseURL, mailboxID)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil)
	if err != nil {
		return "", err
	}
	body, err := m.doRequest(req)
	if err != nil {
		return "", err
	}
	var result listEmailsResponse
	if err := json.Unmarshal(body, &result); err != nil {
		return "", err
	}
	if !result.OK || len(result.Data.Emails) == 0 {
		return "", nil
	}
	for _, e := range result.Data.Emails {
		// Skip emails that existed before our OTP send
		if len(knownIDs) > 0 && knownIDs[e.ID] {
			continue
		}
		text := e.Content
		if text == "" {
			text = e.Body
		}
		if text == "" {
			text = e.Text
		}
		if text == "" {
			text = e.Subject
		}
		matches := codeRegexp.FindStringSubmatch(text)
		if len(matches) >= 2 {
			return matches[1], nil
		}
	}
	return "", nil
}

// pollVerificationEndpoint uses the dedicated verification email endpoint.
// GET /api/v1/emails/verification?mailbox_id=xxx&keyword=openai
func (m *MailGatewayProvider) pollVerificationEndpoint(ctx context.Context, mailboxID string) (string, error) {
	reqURL := fmt.Sprintf("%s/api/v1/emails/verification?mailbox_id=%s&keyword=openai", m.baseURL, mailboxID)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil)
	if err != nil {
		return "", err
	}

	body, err := m.doRequest(req)
	if err != nil {
		return "", err
	}

	var result verificationResponse
	if err := json.Unmarshal(body, &result); err != nil {
		return "", err
	}

	if !result.OK || len(result.Data) == 0 {
		return "", nil // no verification email yet
	}

	// Data can be a single object or an array — handle both
	var entries []emailEntry
	if err := json.Unmarshal(result.Data, &entries); err != nil {
		// Try as single object
		var single emailEntry
		if err2 := json.Unmarshal(result.Data, &single); err2 != nil {
			return "", fmt.Errorf("parse verification data: %w", err2)
		}
		entries = []emailEntry{single}
	}

	// Extract 6-digit OTP from email content/body/text/subject
	for _, e := range entries {
		text := e.Content
		if text == "" {
			text = e.Body
		}
		if text == "" {
			text = e.Text
		}
		if text == "" {
			text = e.Subject
		}
		matches := codeRegexp.FindStringSubmatch(text)
		if len(matches) >= 2 {
			return matches[1], nil
		}
	}

	return "", nil
}

// WaitForTeamAccountID polls for a Team workspace creation email and extracts account_id.
// TODO: implement proper email parsing for team workspace emails.
func (m *MailGatewayProvider) WaitForTeamAccountID(ctx context.Context, mailboxID string, timeout time.Duration, notBefore time.Time) (string, error) {
	return "", fmt.Errorf("WaitForTeamAccountID not implemented for MailGateway provider")
}