warmbox/src/jobs/rescuer.job.ts

import { simpleParser } from 'mailparser';
import { ImapFlow } from 'imapflow';
import type { Inbox, WarmupLog } from '../generated/prisma/client';
import { InboxStatus, WarmupStatus } from '../lib/prisma';
import { prisma } from '../lib/prisma';
import { withInboxGuard } from '../lib/inboxGuard';
import { logger } from '../lib/logger';
import { SPAM_FOLDERS } from '../models/enums';
import * as aiService from '../services/ai.service';
import * as emailService from '../services/email.service';
import { withImapClient } from '../services/imap.service';

type PendingLog = WarmupLog & { sender: Inbox };

function normalizeMessageId(id: string | undefined | null): string | null {
  if (!id) return null;
  return id.replace(/^<|>$/g, '').trim();
}

function extractFromAddress(fromHeader: string | undefined): string | null {
  if (!fromHeader) return null;
  const match = fromHeader.match(/<([^>]+)>/) ?? fromHeader.match(/([\w.+-]+@[\w.-]+)/);
  return match?.[1]?.toLowerCase() ?? null;
}

async function findMatchingLog(
  logs: PendingLog[],
  messageId: string | null,
  fromEmail: string | null,
  subject: string | undefined,
): Promise<PendingLog | undefined> {
  if (messageId) {
    const normalized = normalizeMessageId(messageId);
    const byId = logs.find(
      (l) => l.messageId && normalizeMessageId(l.messageId) === normalized,
    );
    if (byId) return byId;
  }

  if (fromEmail && subject) {
    return logs.find(
      (l) =>
        l.sender.email.toLowerCase() === fromEmail &&
        l.subject.trim() === subject.replace(/^Re:\s*/i, '').trim(),
    );
  }

  return undefined;
}

async function processMessage(
  receiver: Inbox,
  uid: number,
  client: ImapFlow,
  pendingLogs: PendingLog[],
): Promise<void> {
  const msg = await client.fetchOne(String(uid), {
    envelope: true,
    source: true,
    headers: ['message-id', 'from', 'subject'],
    uid: true,
  }, { uid: true });

  if (!msg || !msg.source) return;

  const parsed = await simpleParser(msg.source);
  const fromEmail = extractFromAddress(parsed.from?.text ?? msg.envelope?.from?.[0]?.address);
  const messageId = parsed.messageId ?? msg.envelope?.messageId ?? null;
  const subject = parsed.subject ?? msg.envelope?.subject ?? '';

  const log = await findMatchingLog(pendingLogs, messageId, fromEmail, subject);
  if (!log) return;

  await client.messageFlagsAdd(String(uid), ['\\Seen', '\\Flagged'], { uid: true });

  try {
    await client.messageFlagsAdd(String(uid), ['\\Important'], { uid: true, useLabels: true });
  } catch {
    logger.debug({ inboxId: receiver.id }, 'Important flag not supported');
  }

  const incomingBody = parsed.text ?? parsed.html?.toString() ?? '';
  const { subject: replySubject, body: replyBody } = await aiService.generateReply(
    receiver.industry,
    incomingBody,
    log.subject,
  );

  const originalMessageId = log.messageId ?? messageId ?? undefined;

  await withInboxGuard(receiver, async () => {
    const { messageId: replyMessageId } = await emailService.sendMail(receiver, {
      to: log.sender.email,
      subject: replySubject,
      body: replyBody,
      inReplyTo: originalMessageId,
      references: originalMessageId,
    });

    await prisma.warmupLog.update({
      where: { id: log.id },
      data: { status: WarmupStatus.replied, inReplyTo: replyMessageId },
    });

    await prisma.warmupLog.update({
      where: { id: log.id },
      data: { status: WarmupStatus.complete },
    });

    logger.info(
      { receiver: receiver.email, sender: log.sender.email, logId: log.id },
      'Warmup reply sent',
    );
  });

  // Remove from pending so we don't process twice
  const idx = pendingLogs.findIndex((l) => l.id === log.id);
  if (idx >= 0) pendingLogs.splice(idx, 1);
}

async function rescueFromSpam(
  receiver: Inbox,
  client: ImapFlow,
  pendingLogs: PendingLog[],
): Promise<void> {
  const senderEmails = [...new Set(pendingLogs.map((l) => l.sender.email))];

  for (const folder of SPAM_FOLDERS) {
    try {
      await client.mailboxOpen(folder);
    } catch {
      continue;
    }

    for (const senderEmail of senderEmails) {
      const uids = await client.search({ seen: false, from: senderEmail }, { uid: true });
      if (!uids || uids.length === 0) continue;

      await client.messageMove(uids, 'INBOX', { uid: true });

      const matchingLogs = pendingLogs.filter(
        (l) => l.sender.email.toLowerCase() === senderEmail.toLowerCase(),
      );

      for (const log of matchingLogs) {
        await prisma.warmupLog.update({
          where: { id: log.id },
          data: {
            status: WarmupStatus.rescued_from_spam,
            rescuedFromSpam: true,
          },
        });
        log.status = WarmupStatus.rescued_from_spam;
      }

      logger.info(
        { receiver: receiver.email, folder, count: uids.length, senderEmail },
        'Rescued emails from spam',
      );
    }
  }
}

async function processInbox(receiver: Inbox): Promise<void> {
  const pendingLogs = (await prisma.warmupLog.findMany({
    where: {
      receiverId: receiver.id,
      status: { in: [WarmupStatus.sent, WarmupStatus.rescued_from_spam] },
    },
    include: { sender: true },
  })) as PendingLog[];

  if (pendingLogs.length === 0) return;

  await withInboxGuard(receiver, async () => {
    await withImapClient(receiver, async (client) => {
      await rescueFromSpam(receiver, client, pendingLogs);

      await client.mailboxOpen('INBOX');
      const senderEmails: string[] = [
        ...new Set(pendingLogs.map((l: PendingLog) => l.sender.email)),
      ];

      for (const senderEmail of senderEmails) {
        const uids = await client.search({ seen: false, from: senderEmail }, { uid: true });
        if (!uids || uids.length === 0) continue;

        for (const uid of uids) {
          try {
            await processMessage(receiver, uid, client, pendingLogs);
          } catch (err) {
            logger.error({ err, receiverId: receiver.id, uid }, 'Failed to process message');
          }
        }
      }
    });
  });
}

export async function runRescuerJob(): Promise<void> {
  logger.info('Rescuer job started');

  const inboxes = await prisma.inbox.findMany({
    where: { status: InboxStatus.active },
  });

  for (const inbox of inboxes) {
    try {
      await processInbox(inbox);
    } catch (err) {
      logger.error({ err, inboxId: inbox.id }, 'Rescuer failed for inbox');
    }
  }

  logger.info('Rescuer job finished');
}