warmbox/apps/api/src/services/ai.service.ts

import OpenAI from 'openai';
import { config } from '../config';
import { prisma } from '../lib/prisma';
import { GeneratedEmail } from '../models/types';
import { logger } from '../lib/logger';
import {
  ContentValidationError,
  validateEmailContent,
  ValidationResult,
} from './content-validator.service';

const DEFAULT_MODEL = 'openrouter/free';
const MAX_VALIDATION_ATTEMPTS = 3;

export type MailboxProfile = {
  senderName?: string | null;
  companyName?: string | null;
};

async function getApiKey(): Promise<string> {
  if (config.AI_API_KEY) {
    return config.AI_API_KEY;
  }

  const setting = await prisma.setting.findUnique({
    where: { key: 'ai_api_key' },
  });

  if (setting?.value) {
    return setting.value;
  }

  const legacy = await prisma.setting.findUnique({
    where: { key: 'openai_api_key' },
  });
  if (legacy?.value) {
    return legacy.value;
  }

  throw new Error(
    'AI API key not configured. Set AI_API_KEY in .env (get a free key at https://openrouter.ai/keys)',
  );
}

async function getModel(): Promise<string> {
  if (config.AI_MODEL) {
    return config.AI_MODEL;
  }

  const setting = await prisma.setting.findUnique({
    where: { key: 'ai_model' },
  });

  return setting?.value ?? DEFAULT_MODEL;
}

async function createClient(): Promise<OpenAI> {
  const apiKey = await getApiKey();
  return new OpenAI({
    apiKey,
    baseURL: config.AI_BASE_URL,
    defaultHeaders: config.AI_BASE_URL.includes('openrouter')
      ? {
          'HTTP-Referer': 'https://github.com/warmbox',
          'X-Title': 'Warmbox Email Warmup',
        }
      : undefined,
  });
}

function extractJson(content: string): GeneratedEmail {
  const trimmed = content.trim();

  try {
    return parseEmailResponse(trimmed);
  } catch {
    const match = trimmed.match(/\{[\s\S]*\}/);
    if (match) {
      return parseEmailResponse(match[0]);
    }
    throw new Error('AI response was not valid JSON');
  }
}

function parseEmailResponse(content: string): GeneratedEmail {
  const parsed = JSON.parse(content) as { subject?: string; body?: string };
  if (!parsed.subject || !parsed.body) {
    throw new Error('AI response missing subject or body');
  }
  return { subject: parsed.subject.trim(), body: parsed.body.trim() };
}

function profileContext(profile?: MailboxProfile): string {
  const parts: string[] = [];
  if (profile?.senderName) parts.push(`Sender name: ${profile.senderName}`);
  if (profile?.companyName) parts.push(`Company: ${profile.companyName}`);
  if (parts.length === 0) {
    return 'Do not use placeholder names or bracket templates. Use generic professional language without [Name] style placeholders.';
  }
  return `${parts.join('. ')}. Use these real values in the signature if needed. Never use bracket placeholders.`;
}

async function chatOnce(
  client: OpenAI,
  model: string,
  messages: OpenAI.Chat.ChatCompletionMessageParam[],
): Promise<GeneratedEmail> {
  const request = async (useJsonMode: boolean) => {
    const response = await client.chat.completions.create({
      model,
      messages,
      temperature: 0.8,
      ...(useJsonMode ? { response_format: { type: 'json_object' as const } } : {}),
    });

    const content = response.choices[0]?.message?.content;
    if (!content) {
      throw new Error('Empty AI response');
    }
    return extractJson(content);
  };

  try {
    return await request(true);
  } catch (err) {
    logger.warn({ err, model }, 'AI request failed, retrying without json mode');
    await new Promise((r) => setTimeout(r, 1000));
    return request(false);
  }
}

async function generateValidated(
  buildMessages: (validationFeedback?: string) => OpenAI.Chat.ChatCompletionMessageParam[],
  options: { isReply?: boolean },
): Promise<GeneratedEmail & { validationResult: ValidationResult }> {
  const client = await createClient();
  const model = await getModel();
  let lastValidation: ValidationResult | null = null;

  for (let attempt = 1; attempt <= MAX_VALIDATION_ATTEMPTS; attempt++) {
    const feedback =
      lastValidation && !lastValidation.valid
        ? `Previous attempt failed validation: ${[...lastValidation.errors, ...lastValidation.warnings].join('; ')}. Fix these issues.`
        : undefined;

    const email = await chatOnce(client, model, buildMessages(feedback));
    const validation = validateEmailContent(email.subject, email.body, {
      isReply: options.isReply,
      failOnWarnings: attempt >= 2,
    });

    if (validation.valid) {
      return { ...email, validationResult: validation };
    }

    lastValidation = validation;
    logger.warn({ attempt, errors: validation.errors }, 'AI content validation failed');
  }

  throw new ContentValidationError(
    'AI content failed validation after maximum attempts',
    lastValidation ?? { valid: false, errors: ['unknown'], warnings: [] },
  );
}

export async function generateInitialEmail(
  industry: string,
  profile?: MailboxProfile,
): Promise<GeneratedEmail> {
  const result = await generateInitialEmailWithValidation(industry, profile);
  const { validationResult: _, ...email } = result;
  return email;
}

export async function generateInitialEmailWithValidation(
  industry: string,
  profile?: MailboxProfile,
): Promise<GeneratedEmail & { validationResult: ValidationResult }> {
  return generateValidated(
    (feedback) => [
      {
        role: 'system',
        content: `You write realistic B2B professional emails for the ${industry} industry.
Respond with ONLY valid JSON: {"subject":"...","body":"..."}
The email should sound natural, business-appropriate, and avoid spam trigger words.
Body should be 80-150 words. Subject should be concise and realistic.
No URLs, no bracket placeholders, no mustache templates.
${profileContext(profile)}
${feedback ?? ''}`,
      },
      {
        role: 'user',
        content: 'Generate a new outbound business email.',
      },
    ],
    { isReply: false },
  );
}

export async function generateReply(
  industry: string,
  incomingBody: string,
  threadSubject: string,
  profile?: MailboxProfile,
): Promise<GeneratedEmail> {
  const result = await generateValidated(
    (feedback) => [
      {
        role: 'system',
        content: `You write realistic B2B professional email replies for the ${industry} industry.
Respond with ONLY valid JSON: {"subject":"...","body":"..."}
Write a natural, short reply (40-100 words) that continues the conversation.
If subject does not start with "Re:", prefix it with "Re: ".
No URLs, no bracket placeholders.
${profileContext(profile)}
${feedback ?? ''}`,
      },
      {
        role: 'user',
        content: `Thread subject: ${threadSubject}\n\nIncoming message:\n${incomingBody}`,
      },
    ],
    { isReply: true },
  );

  let { subject, body } = result;
  if (!subject.toLowerCase().startsWith('re:')) {
    subject = `Re: ${threadSubject.replace(/^Re:\s*/i, '')}`;
  }

  const { validationResult: _, ...email } = { subject, body, validationResult: result.validationResult };
  return email;
}