/**
 * Image Generation Routes
 *
 * CRUD + AI-powered prompt enhancement and image generation endpoints.
 * Flow: User Inputs → Ollama Prompt Enhancement → ChatGPT/DALL-E Image Generation
 * → Display Image → Save to Brand Assets → Regenerate/Manage
 */

import express, { Request, Response } from 'express';
import fs from 'fs';
import path from 'path';
import sharp from 'sharp';
import { body, validationResult } from 'express-validator';
import { getModels } from '../models';
import { readBrandAssetImageBytes, resolvePrimaryLogoAssetId, resolveBrandLogoDataUri } from '../utils/brandLogo';
import { authenticate } from '../middleware/auth';
import { requirePermission } from '../middleware/permissions';
import { generateWithAI, getAIConfig, generateZhipuToken, getImageProviderForModel, IMAGE_PROVIDER_NAMES, type ImageProvider } from '../utils/aiProvider';
import { buildEnhancementPrompts, buildFallbackPrompt, ASPECT_RATIO_SIZE_MAP, OPENAI_TO_COGVIEW_SIZE_MAP, COGVIEW_SIZE_MAP, isFullPageDocumentCategory, type ImageEnhancementInputs } from '../services/aiContext/imageGenerationPrompts';
import { getStyleGuidance, getPlatformGuidance, getAssetCategoryGuidance, getAssetContentElements, invalidatePromptCache } from '../services/aiContext/promptConfigLoader';
import { saveBrandAssetFile, saveHrAssetFile, saveStationeryFile, base64ToBuffer, getExtensionFromMime } from '../utils/fileStorage';
import { createJob, updateJobProgress, completeJob, failJob, getJob } from '../services/aiContext/aiJobManager';
import { stripReasoning } from '../utils/stripReasoning';
import { escapeRegex } from '../utils/escapeRegex';
import { notifyAiGenerationCompleted } from '../services/aiGenerationNotifications';

/**
 * Completion email for /generate/:id and /regenerate/:id.
 *
 * These two routes run synchronously and never reach aiJobManager.completeJob,
 * so unlike the job-based flows (logos, watermarks, backdrops, stationery) they
 * have no other path to the notifier. They are also shared — Brand Assets, HR
 * Assets and Stationery all post here — so the notification is opt-in rather
 * than automatic: only a request that declares an allow-listed `notifyModule`
 * triggers it. The value is checked against the allow-list so a caller cannot
 * emit a notification for some other module, and callers that omit the field
 * behave exactly as before.
 *
 * Called only after `status = 'completed'` and a successful save, so failed and
 * partial generations are never notified. Keyed to the generation AND its
 * version, so the notifier's de-dup makes a retry of the same version silent
 * while a genuine new version still notifies.
 */
const NOTIFY_MODULE_ALLOWLIST: Record<string, string> = {
  'brand-assets': 'brand-assets',
};

function notifyImageGenerationCompleted(req: Request, params: {
  generationId: string;
  version: number;
  companyId?: string;
  name?: string;
}): void {
  const source = NOTIFY_MODULE_ALLOWLIST[String(req.body?.notifyModule || '')];
  if (!source) return;
  notifyAiGenerationCompleted({
    jobId: `imggen:${params.generationId}:v${params.version}`,
    moduleSource: source,
    moduleId: source,
    companyId: params.companyId,
    userId: req.user?.id ? String(req.user.id) : undefined,
    autoFillData: params.name ? { name: params.name } : undefined,
    completedAt: Date.now(),
  });
}

const router = express.Router();

router.use(authenticate);

// ============================================
// HELPERS
// ============================================

/** Authorise company access — admin bypasses all checks */
const authorizeCompany = (req: Request, companyId: string): boolean => {
  return req.user?.companyIds?.includes(companyId) || req.user?.role === 'admin';
};

/** Handle errors with proper status codes */
const handleError = (res: Response, error: any) => {
  console.error('[ImageGenerations] Error:', error);
  if (error.name === 'ValidationError') {
    res.status(400).json({ error: error.message, details: Object.values(error.errors).map((e: any) => e.message) });
    return;
  }
  res.status(500).json({ error: error.message || 'Internal server error' });
};

/**
 * Parse an OpenAI API error message into a user-friendly string.
 * Raw errors look like:
 *   "OpenAI Images API error (400) with gpt-image-1: {"error":{"message":"Billing hard limit has been reached.","type":"billing_limit_user_error","code":"billing_hard_limit_reached"}}"
 * This function extracts the meaningful error code/message and returns a clear, actionable explanation.
 */
function parseOpenAIError(rawMessage: string): string {
  // Handle "no API key configured" case first (not JSON, just a plain string)
  if (rawMessage.includes('OpenAI API key not configured')) {
    return 'No OpenAI API key configured. Please add your API key in Super Admin settings or set OPENAI_API_KEY in your environment.';
  }

  // Try to extract the JSON error payload from the raw message
  const jsonMatch = rawMessage.match(/\{[\s\S]*\}/);
  if (jsonMatch) {
    try {
      const parsed = JSON.parse(jsonMatch[0]);
      const errObj = parsed.error || parsed;
      const code = errObj.code || '';
      const type = errObj.type || '';
      const message = (errObj.message || '').toLowerCase();

      // Billing limit / no credits
      if (code === 'billing_hard_limit_reached' || type === 'billing_limit_user_error' || message.includes('billing')) {
        return 'No credits available — your OpenAI billing limit has been reached. Please top up your credits at platform.openai.com and try again.';
      }

      // Insufficient quota (similar to billing limit, different code)
      if (code === 'insufficient_quota' || message.includes('insufficient') || message.includes('quota exceeded')) {
        return 'No credits available — your OpenAI account has insufficient quota. Please top up your credits at platform.openai.com and try again.';
      }

      // Authentication / key errors
      if (code === 'invalid_api_key' || type === 'authentication_error' || message.includes('invalid api key') || message.includes('incorrect api key')) {
        return 'Invalid OpenAI API key. Please check your API key configuration in Super Admin settings.';
      }

      // Rate limiting
      if (code === 'rate_limit_exceeded' || type === 'rate_limit_error' || message.includes('rate limit')) {
        return 'OpenAI rate limit reached — too many requests. Please wait a moment and try again.';
      }

      // Account suspended / deactivated
      if (message.includes('account') && (message.includes('deactivated') || message.includes('suspended'))) {
        return 'Your OpenAI account has been deactivated or suspended. Please check your account status at platform.openai.com.';
      }

      // Model access errors
      if (code === 'model_not_found' || message.includes('does not have access to model')) {
        return 'Your OpenAI plan does not include access to image generation models. Please upgrade your plan at platform.openai.com.';
      }

      // Content policy violations
      if (code === 'content_policy_violation' || message.includes('content policy') || message.includes('safety')) {
        return 'Your prompt was flagged by OpenAI content policy. Please modify your description and try again.';
      }

      // Server errors
      if (message.includes('server error') || message.includes('internal error')) {
        return 'OpenAI is experiencing server issues. Please try again in a few minutes.';
      }

      // If we extracted a readable message, use it with context
      const originalMessage = errObj.message || '';
      if (originalMessage) {
        return `Image generation failed: ${originalMessage}`;
      }
    } catch {
      // JSON parse failed — fall through to raw message pattern matching
    }
  }

  // Check for common patterns in the raw string even without JSON
  if (rawMessage.includes('billing_hard_limit_reached') || rawMessage.includes('billing_limit_user_error')) {
    return 'No credits available — your OpenAI billing limit has been reached. Please top up your credits at platform.openai.com and try again.';
  }
  if (rawMessage.includes('insufficient_quota')) {
    return 'No credits available — your OpenAI account has insufficient quota. Please top up your credits at platform.openai.com and try again.';
  }
  if (rawMessage.includes('invalid_api_key') || rawMessage.includes('Incorrect API key')) {
    return 'Invalid OpenAI API key. Please check your API key configuration in Super Admin settings.';
  }
  if (rawMessage.includes('rate_limit') || rawMessage.includes('Rate limit')) {
    return 'OpenAI rate limit reached — too many requests. Please wait a moment and try again.';
  }
  if (rawMessage.includes('OpenAI API key not configured')) {
    return 'No OpenAI API key configured. Please add your API key in Super Admin settings or set OPENAI_API_KEY in your environment.';
  }

  // Return the raw message as-is (already prefixed with "Image generation failed:" by the caller)
  return rawMessage;
}

/**
 * Fetch brand context for prompt enhancement.
 * Gathers brand data, business profile, and ICP info.
 */
async function fetchBrandContext(
  companyId: string,
  linkedData?: Record<string, string[] | string | undefined>,
): Promise<{
  brandName?: string;
  brandColors?: string[];
  brandTone?: string;
  brandPersonality?: string[];
  brandArchetype?: string;
  businessDescription?: string;
  businessIndustry?: string;
  icpDescription?: string;
  brandStrategy?: {
    mission?: string;
    vision?: string;
    values?: string[] | string;
    positioning?: string;
    differentiators?: string[] | string;
    personalityTraits?: string[] | string;
    voiceTone?: string;
  };
  visualIdentity?: {
    colorPalette?: any;
    typography?: any;
    designPrinciples?: string[] | string;
    moodDescription?: string;
    visualStyle?: string;
  };
  brandGuidelines?: {
    dosAndDonts?: any;
    voiceGuidelines?: any;
    designRules?: any;
  };
  brandManual?: {
    summary?: string;
    usageStandards?: any;
  };
  founderNames?: string[];
  founderBios?: string[];
  founderResponsibilityAreas?: string[];
  brandVisualDirection?: string;
  brandVisualTheme?: string;
  brandConsistencyGuardrails?: {
    cannotChange?: string[];
    canEvolve?: string[];
    misuseExamples?: string[];
  };
  brandForbiddenDesignPatterns?: string[];
}> {
  const { Brand, BusinessProfile, ICP, ModuleData, Founder, Company } = getModels();
  const context: any = {};

  // ── Company — the source of the company's actual name ──
  // Without this the generated designs fall back to the literal "Company Name".
  try {
    const company = await Company.findById(companyId).lean();
    if (company?.name) context.brandName = company.name;
  } catch { /* Company lookup is optional */ }

  // ── Brand model (singleton per company) ──
  // Field names must match models/Brand.ts exactly. getModels() loads models
  // through require(), so every model is `any` here and the compiler cannot
  // flag a field that does not exist — it just reads undefined forever, which
  // is how this block previously contributed nothing at all.
  try {
    const brand = await Brand.findOne({ companyId }).lean();
    if (brand) {
      const colors = [brand.primaryColor, brand.secondaryColor, brand.accentColor].filter(Boolean);
      if (colors.length) context.brandColors = colors;
      if (brand.voiceDescription) context.brandTone = brand.voiceDescription;
      const personality = [...(brand.personalityPrimary || []), ...(brand.personalitySecondary || [])].filter(Boolean);
      if (personality.length) context.brandPersonality = personality;
      // Brand SOP 1.7 — visual direction and guardrails
      if (brand.visualDescription) context.brandVisualDirection = brand.visualDescription;
      // Locked elements are, by definition, the things that must not change.
      // The schema has no equivalent of "can evolve" or "misuse examples", so
      // those stay empty rather than being invented from unrelated fields.
      if (brand.diffLockedElements?.length) {
        context.brandConsistencyGuardrails = {
          cannotChange: brand.diffLockedElements,
          canEvolve: [],
          misuseExamples: [],
        };
      }
      if (brand.rulesDesignForbiddenPatterns?.length) {
        context.brandForbiddenDesignPatterns = brand.rulesDesignForbiddenPatterns;
      }
    }
  } catch { /* Brand data is optional */ }

  // ── Business Profile (use linked ID if provided) ──
  try {
    const businessProfileId = linkedData?.businessProfileId as string | undefined;
    const bp = businessProfileId
      ? await BusinessProfile.findOne({ _id: businessProfileId, companyId }).lean()
      : await BusinessProfile.findOne({ companyId }).lean();
    if (bp) {
      context.businessDescription = bp.description || bp.descriptionLong;
      context.businessIndustry = bp.primaryIndustry;
      // Falls back to the business profile's own name when the company record
      // has none.
      if (!context.brandName && bp.name) context.brandName = bp.name;
    }
  } catch { /* Business profile is optional */ }

  // ── ICP ──
  try {
    const icps = await ICP.find({ companyId }).lean();
    if (icps.length > 0) {
      const icp = icps[0];
      context.icpDescription = [icp.name, icp.description, icp.industry, icp.companySize].filter(Boolean).join(' — ');
    }
  } catch { /* ICP data is optional */ }

  // ── Brand Strategy (use linked IDs if provided, otherwise fetch company default) ──
  try {
    const brandStrategyIds = linkedData?.brandStrategyIds as string[] | undefined;
    if (brandStrategyIds && brandStrategyIds.length > 0) {
      // Fetch specific linked brand strategy records
      const bsRecords = await ModuleData.find({
        moduleId: 'brand-strategy',
        companyId,
        _id: { $in: brandStrategyIds },
      }).lean();
      if (bsRecords.length > 0) {
        // Merge all linked strategy records into a unified context
        const merged: any = {};
        for (const record of bsRecords) {
          const d = record.data || {};
          if (d.mission || d.purposeStatement || d.brandMission) merged.mission = merged.mission || d.mission || d.purposeStatement || d.brandMission;
          if (d.vision || d.brandVision) merged.vision = merged.vision || d.vision || d.brandVision;
          if (d.coreValues || d.brandPillars || d.brandValues || d.values) merged.values = merged.values || d.coreValues || d.brandPillars || d.brandValues || d.values;
          if (d.positioning || d.brandPositioning || d.marketPositioning) merged.positioning = merged.positioning || d.positioning || d.brandPositioning || d.marketPositioning;
          if (d.differentiators || d.uniqueValueProposition || d.competitiveDifference) merged.differentiators = merged.differentiators || d.differentiators || d.uniqueValueProposition || d.competitiveDifference;
          if (d.personalityTraits || d.primaryPersonality || d.brandPersonality) merged.personalityTraits = merged.personalityTraits || d.personalityTraits || d.primaryPersonality || d.brandPersonality;
          if (d.voiceTone || d.brandVoice || d.toneOfVoice) merged.voiceTone = merged.voiceTone || d.voiceTone || d.brandVoice || d.toneOfVoice;
          // For array fields, merge unique values
          if (Array.isArray(d.brandValues)) merged.values = Array.from(new Set([...(Array.isArray(merged.values) ? merged.values : merged.values ? [merged.values] : []), ...d.brandValues]));
          if (Array.isArray(d.brandPillars)) merged.values = Array.from(new Set([...(Array.isArray(merged.values) ? merged.values : merged.values ? [merged.values] : []), ...d.brandPillars]));
          if (Array.isArray(d.personalityTraits)) merged.personalityTraits = Array.from(new Set([...(Array.isArray(merged.personalityTraits) ? merged.personalityTraits : []), ...d.personalityTraits]));
        }
        context.brandStrategy = merged;
      }
    } else {
      // No linked IDs — fetch company default
      const bsData = await ModuleData.findOne({ moduleId: 'brand-strategy', companyId }).lean();
      if (bsData?.data) {
        const d = bsData.data;
        context.brandStrategy = {
          mission: d.mission || d.purposeStatement || d.brandMission,
          vision: d.vision || d.brandVision,
          values: d.coreValues || d.brandPillars || d.brandValues || d.values,
          positioning: d.positioning || d.brandPositioning || d.marketPositioning,
          differentiators: d.differentiators || d.uniqueValueProposition || d.competitiveDifference,
          personalityTraits: d.personalityTraits || d.primaryPersonality || d.brandPersonality,
          voiceTone: d.voiceTone || d.brandVoice || d.toneOfVoice,
        };
      }
    }
  } catch { /* Brand Strategy data is optional */ }

  // ── Visual Identity (use linked IDs if provided, otherwise fetch company default) ──
  try {
    const visualIdentityIds = linkedData?.visualIdentityIds as string[] | undefined;
    if (visualIdentityIds && visualIdentityIds.length > 0) {
      const viRecords = await ModuleData.find({
        moduleId: 'visual-identity',
        companyId,
        _id: { $in: visualIdentityIds },
      }).lean();
      if (viRecords.length > 0) {
        // Use the first (or primary) linked visual identity record
        // For visual identity, we use the first matching record since colour palettes
        // and typography should come from a single coherent design system
        const d = viRecords[0].data || {};
        context.visualIdentity = {
          colorPalette: d.colorPalette || d.colors || { primary: d.primaryColor, secondary: d.secondaryColor, accent: d.accentColor, background: d.backgroundColor, surface: d.surfaceColor, text: d.textColor },
          typography: d.typography || d.fonts || { heading: d.headingFont, body: d.bodyFont, accent: d.accentFont },
          designPrinciples: d.designPrinciples || d.principles,
          moodDescription: d.mood || d.moodDescription || d.visualDescription,
          visualStyle: d.visualStyle || d.styleDirection || d.imageStyle?.description,
        };
      }
    } else {
      const viData = await ModuleData.findOne({ moduleId: 'visual-identity', companyId }).lean();
      if (viData?.data) {
        const d = viData.data;
        context.visualIdentity = {
          colorPalette: d.colorPalette || d.colors || { primary: d.primaryColor, secondary: d.secondaryColor, accent: d.accentColor, background: d.backgroundColor, surface: d.surfaceColor, text: d.textColor },
          typography: d.typography || d.fonts || { heading: d.headingFont, body: d.bodyFont, accent: d.accentFont },
          designPrinciples: d.designPrinciples || d.principles,
          moodDescription: d.mood || d.moodDescription || d.visualDescription,
          visualStyle: d.visualStyle || d.styleDirection || d.imageStyle?.description,
        };
      }
    }
  } catch { /* Visual Identity data is optional */ }

  // ── Brand Guidelines (always fetch company default — no longer in data source selector) ──
  try {
    const bgData = await ModuleData.findOne({ moduleId: 'brand-guidelines', companyId }).lean();
    if (bgData?.data) {
      const d = bgData.data;
      context.brandGuidelines = {
        dosAndDonts: d.dosAndDonts || d.guidelines || d.brandRules,
        voiceGuidelines: d.voiceGuidelines || d.toneGuidelines || d.brandVoice,
        designRules: d.designRules || d.designGuidelines || d.visualRules,
      };
    }
  } catch { /* Brand Guidelines data is optional */ }

  // ── Brand Manual (always fetch company default — no longer in data source selector) ──
  try {
    const bmData = await ModuleData.findOne({ moduleId: 'brand-manual', companyId }).lean();
    if (bmData?.data) {
      const d = bmData.data;
      context.brandManual = {
        summary: d.summary || d.overview || d.introduction,
        usageStandards: d.usageStandards || d.standards || d.brandStandards,
      };
    }
  } catch { /* Brand Manual data is optional */ }

  // ── Founders (use linked IDs if provided, otherwise fetch all for company) ──
  try {
    const founderIds = linkedData?.founderIds as string[] | undefined;
    let founders: any[];
    if (founderIds && founderIds.length > 0) {
      founders = await Founder.find({ _id: { $in: founderIds }, companyId }).lean();
    } else {
      founders = await Founder.find({ companyId }).limit(5).lean();
    }
    if (founders.length > 0) {
      context.founderNames = founders.map((f: any) => f.name).filter(Boolean);
      context.founderBios = founders.map((f: any) => f.bio).filter(Boolean);
      context.founderResponsibilityAreas = founders
        .map((f: any) => f.responsibilityArea || (f.expertise?.length ? f.expertise.join(', ') : ''))
        .filter(Boolean);
    }
  } catch { /* Founder data is optional */ }

  return context;
}

/**
 * Call OpenAI Image Generation API (DALL-E 3).
 * Different from chat completions — uses /v1/images/generations endpoint.
 */
export async function generateImageWithOpenAI(
  prompt: string,
  size: string,
  quality: string = 'standard',
  style: string = 'vivid',
  model: string = 'gpt-image-1',
  keyOverride?: any,
  userId?: string
): Promise<{
  base64Data: string;
  revisedPrompt: string;
  model: string;
  provider: string;
  tokenUsage: { inputTokens: number; outputTokens: number; totalTokens: number };
  latencyMs: number;
}> {
  const config = await getAIConfig(userId);
  const callStartTime = Date.now();

  // Determine the API key and base URL
  const openaiKey = keyOverride?.key || config.OPENAI_API_KEY;
  const openaiBaseUrl = (keyOverride?.url || config.OPENAI_API_URL || 'https://api.openai.com/v1/chat/completions')
    .replace('/chat/completions', '')
    .replace(/\/$/, '');

  if (!openaiKey) {
    throw new Error('OpenAI API key not configured. Set OPENAI_API_KEY in .env or add a key in Super Admin settings.');
  }

  // Use the images API endpoint
  const imagesUrl = `${openaiBaseUrl}/images/generations`;
  console.log(`[ImageGen] Calling OpenAI Images API: ${imagesUrl} | model=${model} | size=${size} | quality=${quality}`);

  const timeout = config.AI_TIMEOUT || 180000; // Image generation can take longer
  const controller = new AbortController();
  const timeoutId = setTimeout(() => controller.abort(), timeout);

  // Build request body based on model capabilities:
  // - dall-e-3: supports response_format, quality (standard/hd), style (vivid/natural)
  // - gpt-image-1: supports quality (low/medium/high/auto), does NOT support response_format or style
  // Project keys (sk-proj-*) typically route to gpt-image-1 regardless of the model param.
  const isGPTImageModel = model === 'gpt-image-1';

  const buildBody = (m: string, q: string, s: string): any => {
    // Project keys (sk-proj-*) force gpt-image-1 regardless of the model param,
    // so dall-e-3 specific params (style, n, response_format) will cause errors.
    // Use minimal safe params for all models to maximize compatibility.
    const qualityMap: Record<string, string> = { standard: 'auto', hd: 'high' };
    if (m === 'dall-e-3') {
      // For true dall-e-3 access (non-project keys), include dall-e-3 specific params
      return {
        model: m,
        prompt,
        n: 1,
        size,
        quality: q,
        style: s,
        response_format: 'b64_json',
      };
    }
    // gpt-image-1: only model, prompt, size, quality
    return {
      model: m,
      prompt,
      size,
      quality: qualityMap[q] || 'auto',
    };
  };

  const parseImageResponse = async (data: any, usedModel: string): Promise<{ base64Data: string; revisedPrompt: string; model: string; provider: string; tokenUsage: { inputTokens: number; outputTokens: number; totalTokens: number } }> => {
    // DALL-E 3 returns: { data: [{ b64_json, revised_prompt }] }
    // gpt-image-1 returns: { data: [{ b64_json }] } OR with url OR output_format-based
    // Extract token usage from OpenAI response if available
    const usage = data.usage;
    const tokenUsage = {
      inputTokens: usage?.input_tokens || usage?.prompt_tokens || 0,
      outputTokens: usage?.output_tokens || usage?.completion_tokens || 0,
      totalTokens: (usage?.input_tokens || usage?.prompt_tokens || 0) + (usage?.output_tokens || usage?.completion_tokens || 0),
    };
    // Log the response structure for debugging
    console.log(`[ImageGen] Response keys: ${Object.keys(data).join(', ')} | data length: ${data.data?.length || 0} | output length: ${data.output?.length || 0} | usage: ${JSON.stringify(tokenUsage)}`);
    if (data.data?.[0]) {
      const dataKeys = Object.keys(data.data[0]);
      console.log(`[ImageGen] data[0] keys: ${dataKeys.join(', ')}`);
      // Log first 50 chars of each field value to see what we're getting
      for (const key of dataKeys) {
        const val = data.data[0][key];
        if (typeof val === 'string') {
          console.log(`[ImageGen]   data[0].${key}: ${val.substring(0, 80)}${val.length > 80 ? '...' : ''} (length: ${val.length})`);
        } else {
          console.log(`[ImageGen]   data[0].${key}: ${JSON.stringify(val)}`);
        }
      }
    }

    // Try multiple response structures
    let base64 = '';
    let revisedPrompt = '';

    // Format 1: { data: [{ b64_json }] } — DALL-E 3 and some gpt-image-1 responses
    const image = data.data?.[0];
    if (image) {
      // Try all possible field names for base64 data
      base64 = image.b64_json || image.base64_data || image.base64Data || image.data || '';
      revisedPrompt = image.revised_prompt || '';
    }

    // Format 2: { output: [{ content: [{ type: "image_url", url: "data:image/png;base64,..." }] }] } — gpt-image-1 chat-style response
    if (!base64 && data.output) {
      for (const outputItem of data.output) {
        if (outputItem.content) {
          for (const contentItem of outputItem.content) {
            if (contentItem.type === 'image_url' && contentItem.url) {
              // Extract base64 from data URL: "data:image/png;base64,ABC123..."
              if (contentItem.url.startsWith('data:')) {
                const base64Match = contentItem.url.match(/^data:image\/\w+;base64,(.+)$/);
                if (base64Match) {
                  base64 = base64Match[1];
                }
              } else {
                // It's a regular URL — we'll download it below
                console.log(`[ImageGen] Found image URL in output format: ${contentItem.url.substring(0, 80)}...`);
                try {
                  const urlResponse = await fetch(contentItem.url, { signal: AbortSignal.timeout(30000) });
                  if (urlResponse.ok) {
                    const arrayBuffer = await urlResponse.arrayBuffer();
                    base64 = Buffer.from(arrayBuffer).toString('base64');
                    console.log(`[ImageGen] Downloaded image from output URL, base64 length: ${base64.length}`);
                  }
                } catch (e: any) {
                  console.error(`[ImageGen] Error downloading output URL: ${e.message}`);
                }
              }
            }
          }
        }
      }
    }

    // Format 3: { data: [{ url: "https://..." }] } — URL response (need to download)
    if (!base64 && image?.url) {
      console.log(`[ImageGen] Got URL response, downloading image from: ${image.url.substring(0, 80)}...`);
      try {
        const urlResponse = await fetch(image.url, { signal: AbortSignal.timeout(30000) });
        if (urlResponse.ok) {
          const arrayBuffer = await urlResponse.arrayBuffer();
          base64 = Buffer.from(arrayBuffer).toString('base64');
          console.log(`[ImageGen] Downloaded image from URL, base64 length: ${base64.length}`);
        } else {
          console.error(`[ImageGen] Failed to download image URL: ${urlResponse.status}`);
        }
      } catch (downloadError: any) {
        console.error(`[ImageGen] Error downloading image URL: ${downloadError.message}`);
      }
    }

    if (!base64) {
      console.error('[ImageGen] No base64 data found. Response structure:', JSON.stringify(data).substring(0, 1000));
      throw new Error('No base64 image data found in OpenAI response');
    }

    return {
      base64Data: base64,
      revisedPrompt,
      model: usedModel,
      provider: 'openai',
      tokenUsage,
    };
  };

  try {
    // Strategy: Try with the requested model first, then fallback to other models
    // Project keys (sk-proj-*) may force gpt-image-1 regardless of the model param,
    // so we also try with minimal params (no model-specific params) as a last resort.
    const attempts: { model: string; label: string; minimalParams: boolean }[] = [
      { model, label: model, minimalParams: false },
    ];

    // Add fallback models based on what was requested
    if (model === 'gpt-image-1') {
      // Try dall-e-3 with full params, then with minimal params
      attempts.push({ model: 'dall-e-3', label: 'dall-e-3 (fallback)', minimalParams: false });
      attempts.push({ model: 'dall-e-3', label: 'dall-e-3 minimal (fallback)', minimalParams: true });
    } else if (model === 'dall-e-3') {
      // Try gpt-image-1, then dall-e-3 with minimal params
      attempts.push({ model: 'gpt-image-1', label: 'gpt-image-1 (fallback)', minimalParams: false });
      attempts.push({ model: 'dall-e-3', label: 'dall-e-3 minimal (fallback)', minimalParams: true });
    }

    let lastError: string = '';

    for (const attempt of attempts) {
      let body = buildBody(attempt.model, quality, style);

      // Minimal params: strip model-specific params that project keys may reject
      if (attempt.minimalParams) {
        body = { model: attempt.model, prompt, size };
        console.log(`[ImageGen] Attempt with ${attempt.label} (minimal params): keys=${Object.keys(body).join(', ')}`);
      } else {
        console.log(`[ImageGen] Attempt with ${attempt.label}: keys=${Object.keys(body).join(', ')}`);
      }

      // Retry up to 3 times for transient server errors (500, 502, 503)
      const maxRetries = 3;
      for (let retry = 0; retry < maxRetries; retry++) {
        try {
          const attemptController = new AbortController();
          const attemptTimeoutId = setTimeout(() => attemptController.abort(), timeout);

          const response = await fetch(imagesUrl, {
            method: 'POST',
            headers: {
              'Content-Type': 'application/json',
              'Authorization': `Bearer ${openaiKey}`,
            },
            body: JSON.stringify(body),
            signal: attemptController.signal,
          });

          clearTimeout(attemptTimeoutId);

          if (!response.ok) {
            const errorText = await response.text();
            lastError = `OpenAI Images API error (${response.status}) with ${attempt.label}: ${errorText}`;
            console.warn(`[ImageGen] ${lastError}`);

            // If this is an account-level error (billing limit, invalid key, insufficient quota, account suspended),
            // there's no point trying other models — they will all fail with the same error.
            // Bail out immediately with a clear message.
            const isAccountError =
              errorText.includes('billing_hard_limit_reached') ||
              errorText.includes('billing_limit_user_error') ||
              errorText.includes('insufficient_quota') ||
              errorText.includes('invalid_api_key') ||
              errorText.includes('Incorrect API key') ||
              errorText.includes('account_deactivated') ||
              errorText.includes('account_suspended');
            if (isAccountError) {
              console.warn(`[ImageGen] Account-level error detected (${response.status}), skipping model fallbacks.`);
              throw new Error(lastError);
            }

            // If this is a param error or model access error, try the next model (not a retry)
            if (errorText.includes('Unknown parameter') || errorText.includes('unknown_parameter') || errorText.includes('model_not_found') || errorText.includes('does not have access to model')) {
              console.log(`[ImageGen] ${attempt.label} not available, trying next model...`);
              break; // Break out of retry loop, continue to next model
            }

            // Transient server errors (500, 502, 503) — retry with backoff
            if (response.status >= 500 && retry < maxRetries - 1) {
              const backoffMs = (retry + 1) * 3000; // 3s, 6s, 9s
              console.log(`[ImageGen] Transient error ${response.status}, retrying in ${backoffMs}ms (attempt ${retry + 1}/${maxRetries})...`);
              await new Promise(resolve => setTimeout(resolve, backoffMs));
              continue;
            }

            // Rate limit (429) — retry once after a longer delay
            if (response.status === 429 && retry < 1) {
              console.log(`[ImageGen] Rate limited, retrying in 10s...`);
              await new Promise(resolve => setTimeout(resolve, 10000));
              continue;
            }

            // Non-retryable errors (400 auth, billing, etc.) — don't retry
            throw new Error(lastError);
          }

          const data = await response.json();
          const latencyMs = Date.now() - callStartTime;
          const result = await parseImageResponse(data, attempt.model);
          console.log(`[ImageGen] Image generated successfully with ${attempt.label} in ${latencyMs}ms | revised prompt length=${result.revisedPrompt.length} | tokens=${JSON.stringify(result.tokenUsage)}`);
          return { ...result, latencyMs };
        } catch (error: any) {
          // Account-level errors (billing, quota, auth) — bail out entirely, no point trying other models
          const isAccountLevelError =
            error.message?.includes('billing_hard_limit_reached') ||
            error.message?.includes('billing_limit_user_error') ||
            error.message?.includes('insufficient_quota') ||
            error.message?.includes('invalid_api_key') ||
            error.message?.includes('Incorrect API key') ||
            error.message?.includes('account_deactivated') ||
            error.message?.includes('account_suspended');
          if (isAccountLevelError) {
            throw error; // Don't try fallback models — same account, same error
          }

          if (error.message?.includes('Unknown parameter') || error.message?.includes('unknown_parameter') || error.message?.includes('model_not_found') || error.message?.includes('does not have access to model')) {
            break; // Break retry loop, try next model
          }
          // Network/timeout errors — retry
          if (retry < maxRetries - 1 && !error.message?.includes('OpenAI Images API error')) {
            const backoffMs = (retry + 1) * 3000;
            console.log(`[ImageGen] Network error, retrying in ${backoffMs}ms...`);
            await new Promise(resolve => setTimeout(resolve, backoffMs));
            continue;
          }
          throw error; // Re-throw after exhausting retries
        }
      }
    }

    // All attempts failed
    throw new Error(lastError || 'All image generation attempts failed');
  } catch (error: any) {
    clearTimeout(timeoutId);
    console.error(`[ImageGen] Error generating image: ${error.message}`);
    throw error;
  }
}

/**
 * Generate an image using Zhipu CogView-3 as a fallback provider.
 * CogView-3 returns image URLs which are downloaded and converted to base64.
 */
export async function generateImageWithZhipuCogView(
  prompt: string,
  size: string,
  userId?: string,
): Promise<{
  base64Data: string;
  revisedPrompt: string;
  model: string;
  provider: string;
  tokenUsage: { inputTokens: number; outputTokens: number; totalTokens: number };
  latencyMs: number;
}> {
  const config = await getAIConfig(userId);
  const callStartTime = Date.now();

  // Map the OpenAI size to a CogView-supported size
  const cogviewSize = OPENAI_TO_COGVIEW_SIZE_MAP[size] || COGVIEW_SIZE_MAP['1:1'] || '1024x1024';

  // Build the CogView endpoint from the Zhipu API URL
  const zhipuBaseUrl = (config.ZHIPU_API_URL || 'https://open.bigmodel.cn/api/paas/v4/chat/completions')
    .replace('/chat/completions', '')
    .replace(/\/$/, '');
  const cogviewUrl = `${zhipuBaseUrl}/images/generations`;

  // Collect available Zhipu keys (DB keys + env fallback)
  // Include all keys in id.secret format that could be Zhipu keys
  const zhipuKeys: string[] = [];
  const seenKeys = new Set<string>();

  // DB-configured keys
  if (config.ZHIPU_KEY_LIST && config.ZHIPU_KEY_LIST.length > 0) {
    for (const keyEntry of config.ZHIPU_KEY_LIST) {
      if (keyEntry.isActive && keyEntry.key && !seenKeys.has(keyEntry.key)) {
        zhipuKeys.push(keyEntry.key);
        seenKeys.add(keyEntry.key);
      }
    }
  }
  // DB/config Zhipu key
  if (config.ZHIPU_API_KEY && !seenKeys.has(config.ZHIPU_API_KEY)) {
    zhipuKeys.push(config.ZHIPU_API_KEY);
    seenKeys.add(config.ZHIPU_API_KEY);
  }
  // DB/config Ollama key (Zhipu id.secret format)
  if (config.OLLAMA_API_KEY && config.OLLAMA_API_KEY.includes('.') && !seenKeys.has(config.OLLAMA_API_KEY)) {
    zhipuKeys.push(config.OLLAMA_API_KEY);
    seenKeys.add(config.OLLAMA_API_KEY);
  }
  // ENV fallback keys (may differ from DB-configured keys)
  const envZhipuKey = process.env.ZHIPU_API_KEY;
  if (envZhipuKey && envZhipuKey.includes('.') && !seenKeys.has(envZhipuKey)) {
    zhipuKeys.push(envZhipuKey);
    seenKeys.add(envZhipuKey);
  }
  const envOllamaKey = process.env.OLLAMA_API_KEY;
  if (envOllamaKey && envOllamaKey.includes('.') && !seenKeys.has(envOllamaKey)) {
    zhipuKeys.push(envOllamaKey);
    seenKeys.add(envOllamaKey);
  }

  if (zhipuKeys.length === 0) {
    throw new Error('No Zhipu API key configured for CogView-3 fallback');
  }

  const models = ['cogview-3-plus', 'cogview-3'];
  const timeout = config.AI_TIMEOUT || 120000;

  let lastError = '';

  for (const model of models) {
    for (const zhipuKey of zhipuKeys) {
      // Try two auth methods: JWT token first, then direct Bearer
      const authMethods = [
        { label: 'JWT', header: () => `Bearer ${generateZhipuToken(zhipuKey)}` },
        { label: 'Direct', header: () => `Bearer ${zhipuKey}` },
      ];

      for (const auth of authMethods) {
        console.log(`[ImageGen-CogView] Trying model=${model}, key=${zhipuKey.substring(0, 8)}..., auth=${auth.label}, size=${cogviewSize}`);

        try {
          const controller = new AbortController();
          const timeoutId = setTimeout(() => controller.abort(), timeout);

          const response = await fetch(cogviewUrl, {
            method: 'POST',
            headers: {
              'Content-Type': 'application/json',
              'Authorization': auth.header(),
            },
            body: JSON.stringify({
              model,
              prompt,
              size: cogviewSize,
            }),
            signal: controller.signal,
          });

          clearTimeout(timeoutId);

          if (!response.ok) {
            const errorText = await response.text();
            lastError = `CogView API error (${response.status}) with model=${model}, auth=${auth.label}: ${errorText}`;
            console.warn(`[ImageGen-CogView] ${lastError}`);

            // If auth fails (401), try next auth method
            if (response.status === 401 || response.status === 403) {
              continue;
            }

            // If model not found, try next model
            if (response.status === 404 || errorText.includes('model_not_found') || errorText.includes('Model not found')) {
              break; // Break auth loop, try next model
            }

            // Rate limit — wait and retry
            if (response.status === 429) {
              console.log('[ImageGen-CogView] Rate limited, waiting 5s...');
              await new Promise(resolve => setTimeout(resolve, 5000));
              continue;
            }

            // Other errors — try next key
            continue;
          }

          const data: any = await response.json();

          // CogView returns: { data: [{ url: "https://..." }] }
          const imageUrl = data.data?.[0]?.url;
          if (!imageUrl) {
            console.error('[ImageGen-CogView] No image URL in response:', JSON.stringify(data).substring(0, 500));
            lastError = 'No image URL in CogView-3 response';
            continue;
          }

          // Download the image and convert to base64
          console.log(`[ImageGen-CogView] Downloading image from CogView URL...`);
          const imageDownloadResponse = await fetch(imageUrl, { signal: AbortSignal.timeout(30000) });
          if (!imageDownloadResponse.ok) {
            lastError = `Failed to download CogView image: ${imageDownloadResponse.status}`;
            console.warn(`[ImageGen-CogView] ${lastError}`);
            continue;
          }

          const arrayBuffer = await imageDownloadResponse.arrayBuffer();
          const base64Data = Buffer.from(arrayBuffer).toString('base64');

          const latencyMs = Date.now() - callStartTime;
          console.log(`[ImageGen-CogView] Image generated successfully with ${model} in ${latencyMs}ms | base64 length=${base64Data.length}`);

          // CogView doesn't provide token counts; estimate from prompt length
          const estimatedTokens = Math.ceil(prompt.length / 4); // rough character-to-token estimate
          const tokenUsage = { inputTokens: estimatedTokens, outputTokens: 0, totalTokens: estimatedTokens };

          return {
            base64Data,
            revisedPrompt: '', // CogView does not return a revised prompt
            model,
            provider: 'zhipu-cogview',
            tokenUsage,
            latencyMs,
          };
        } catch (error: any) {
          console.warn(`[ImageGen-CogView] Error with model=${model}, auth=${auth.label}: ${error.message}`);
          lastError = error.message;
          // Network errors — try next auth/key
          continue;
        }
      }
    }
  }

  throw new Error(`CogView-3 fallback failed: ${lastError || 'All attempts exhausted'}`);
}

// ============================================
// FLUX Image Generation
// ============================================

/**
 * Generate an image using FLUX (Black Forest Labs / fal.ai API).
 * Uses OpenAI-compatible /v1/images/generations endpoint format.
 */
export async function generateImageWithFlux(
  prompt: string,
  size: string,
  model: string = 'flux-pro',
  keyOverride?: any,
  userId?: string,
): Promise<{ base64Data: string; revisedPrompt: string; model: string; provider: string; tokenUsage: { inputTokens: number; outputTokens: number; totalTokens: number }; latencyMs: number }> {
  const config = await getAIConfig(userId);
  const callStartTime = Date.now();

  // Collect available FLUX keys
  const fluxKeys: Array<{ key: string; url?: string; model?: string }> = [];
  const seenKeys = new Set<string>();

  // DB-configured keys (priority order: active+healthy first)
  if (config.FLUX_KEY_LIST && config.FLUX_KEY_LIST.length > 0) {
    for (const keyEntry of config.FLUX_KEY_LIST) {
      if (keyEntry.key && !seenKeys.has(keyEntry.key)) {
        fluxKeys.push({ key: keyEntry.key, url: keyEntry.url, model: keyEntry.model });
        seenKeys.add(keyEntry.key);
      }
    }
  }
  // Global/config key
  if (config.FLUX_API_KEY && !seenKeys.has(config.FLUX_API_KEY)) {
    fluxKeys.push({ key: config.FLUX_API_KEY, url: config.FLUX_API_URL, model: config.FLUX_MODEL });
    seenKeys.add(config.FLUX_API_KEY);
  }

  // Use key override if provided
  if (keyOverride?.key) {
    fluxKeys.unshift({ key: keyOverride.key, url: keyOverride.url, model: keyOverride.model });
  }

  if (fluxKeys.length === 0) {
    throw new Error(`No active API key configured for FLUX.\nPlease configure one in Super Admin Settings.`);
  }

  // Map aspect ratio sizes for FLUX (supports common sizes)
  const FLUX_SIZE_MAP: Record<string, string> = {
    '1024x1024': '1024x1024',
    '1536x1024': '1536x1024',
    '1024x1536': '1024x1536',
    '1024x768': '1024x768',
    '768x1024': '768x1024',
    '1:1': '1024x1024',
    '16:9': '1536x1024',
    '9:16': '1024x1536',
    '4:3': '1024x768',
    '3:4': '768x1024',
  };
  const fluxSize = FLUX_SIZE_MAP[size] || size || '1024x1024';

  const timeout = config.AI_TIMEOUT || 180000;
  let lastError = '';

  for (const fluxKey of fluxKeys) {
    const fluxBaseUrl = (fluxKey.url || config.FLUX_API_URL || 'https://api.bfl.ml/v1')
      .replace(/\/chat\/completions$/, '')
      .replace(/\/$/, '');
    const imagesUrl = `${fluxBaseUrl}/images/generations`;
    const fluxModel = fluxKey.model || model || config.FLUX_MODEL || 'flux-pro';

    console.log(`[ImageGen-FLUX] Trying model=${fluxModel}, size=${fluxSize}, key=****${fluxKey.key.slice(-4)}, url=${fluxBaseUrl}`);

    try {
      const controller = new AbortController();
      const timeoutId = setTimeout(() => controller.abort(), timeout);

      const body: any = {
        model: fluxModel,
        prompt,
        image_size: fluxSize,
        num_images: 1,
      };

      const response = await fetch(imagesUrl, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'Authorization': `Bearer ${fluxKey.key}`,
        },
        body: JSON.stringify(body),
        signal: controller.signal,
      });

      clearTimeout(timeoutId);

      if (!response.ok) {
        const errorText = await response.text();
        lastError = `FLUX API error (${response.status}) with model=${fluxModel}: ${errorText}`;
        console.warn(`[ImageGen-FLUX] ${lastError}`);

        // Auth errors — try next key
        if (response.status === 401 || response.status === 403) {
          continue;
        }
        // Server errors — retry won't help, but try next key
        if (response.status >= 500) {
          continue;
        }
        // Other errors (billing, rate limit, etc.) — throw
        throw new Error(lastError);
      }

      const data: any = await response.json();

      // FLUX API can return data in multiple formats:
      // Format 1: { data: [{ b64_json }] } — base64 image data
      // Format 2: { data: [{ url }] } — URL to download
      // Format 3: { images: [{ url }] } — alternative URL format
      let base64 = '';
      let revisedPrompt = '';

      if (data.data?.[0]) {
        const image = data.data[0];
        base64 = image.b64_json || image.base64_data || '';
        revisedPrompt = image.revised_prompt || '';
        if (!base64 && image.url) {
          // Download from URL
          try {
            const urlResponse = await fetch(image.url, { signal: AbortSignal.timeout(30000) });
            if (urlResponse.ok) {
              const arrayBuffer = await urlResponse.arrayBuffer();
              base64 = Buffer.from(arrayBuffer).toString('base64');
              console.log(`[ImageGen-FLUX] Downloaded image from URL, base64 length: ${base64.length}`);
            }
          } catch (downloadErr: any) {
            console.error(`[ImageGen-FLUX] Error downloading image URL: ${downloadErr.message}`);
          }
        }
      } else if (data.images?.[0]) {
        const image = data.images[0];
        base64 = image.b64_json || image.base64_data || '';
        revisedPrompt = image.revised_prompt || '';
        if (!base64 && image.url) {
          try {
            const urlResponse = await fetch(image.url, { signal: AbortSignal.timeout(30000) });
            if (urlResponse.ok) {
              const arrayBuffer = await urlResponse.arrayBuffer();
              base64 = Buffer.from(arrayBuffer).toString('base64');
              console.log(`[ImageGen-FLUX] Downloaded image from URL, base64 length: ${base64.length}`);
            }
          } catch (downloadErr: any) {
            console.error(`[ImageGen-FLUX] Error downloading image URL: ${downloadErr.message}`);
          }
        }
      } else if (data.image) {
        // Some FLUX APIs return { image: "base64..." } or { image: "url..." }
        base64 = data.image.b64_json || data.image.base64_data || '';
        if (!base64 && data.image.url) {
          try {
            const urlResponse = await fetch(data.image.url, { signal: AbortSignal.timeout(30000) });
            if (urlResponse.ok) {
              const arrayBuffer = await urlResponse.arrayBuffer();
              base64 = Buffer.from(arrayBuffer).toString('base64');
            }
          } catch { /* fallback to error below */ }
        }
      }

      if (!base64) {
        console.error('[ImageGen-FLUX] No base64 data found. Response:', JSON.stringify(data).substring(0, 500));
        throw new Error('No image data found in FLUX response');
      }

      const latencyMs = Date.now() - callStartTime;
      const tokenUsage = {
        inputTokens: data.usage?.input_tokens || data.usage?.prompt_tokens || 0,
        outputTokens: data.usage?.output_tokens || data.usage?.completion_tokens || 0,
        totalTokens: (data.usage?.input_tokens || data.usage?.prompt_tokens || 0) + (data.usage?.output_tokens || data.usage?.completion_tokens || 0),
      };

      console.log(`[ImageGen-FLUX] Image generated successfully with ${fluxModel} in ${latencyMs}ms | base64 length=${base64.length}`);

      return {
        base64Data: base64,
        revisedPrompt,
        model: fluxModel,
        provider: 'flux',
        tokenUsage,
        latencyMs,
      };
    } catch (error: any) {
      if (error.name === 'AbortError') {
        throw new Error(`FLUX request timed out after ${timeout / 1000}s`);
      }
      lastError = error.message;
      console.warn(`[ImageGen-FLUX] Error with model=${fluxModel}: ${error.message}`);
      // Continue to next key
      continue;
    }
  }

  throw new Error(lastError || `FLUX image generation failed: No keys available`);
}

// ============================================
// Ideogram Image Generation
// ============================================

/**
 * Generate an image using Ideogram API.
 * Uses Ideogram v3 API format: POST {baseUrl}/ideograms with Api-Key header.
 */
export async function generateImageWithIdeogram(
  prompt: string,
  size: string,
  model: string = 'ideogram-v3',
  keyOverride?: any,
  userId?: string,
): Promise<{ base64Data: string; revisedPrompt: string; model: string; provider: string; tokenUsage: { inputTokens: number; outputTokens: number; totalTokens: number }; latencyMs: number }> {
  const config = await getAIConfig(userId);
  const callStartTime = Date.now();

  const ideogramKeys: Array<{ key: string; url?: string; model?: string }> = [];
  const seenKeys = new Set<string>();

  if (config.IDEOGRAM_KEY_LIST && config.IDEOGRAM_KEY_LIST.length > 0) {
    for (const keyEntry of config.IDEOGRAM_KEY_LIST) {
      if (keyEntry.key && !seenKeys.has(keyEntry.key)) {
        ideogramKeys.push({ key: keyEntry.key, url: keyEntry.url, model: keyEntry.model });
        seenKeys.add(keyEntry.key);
      }
    }
  }
  if (config.IDEOGRAM_API_KEY && !seenKeys.has(config.IDEOGRAM_API_KEY)) {
    ideogramKeys.push({ key: config.IDEOGRAM_API_KEY, url: config.IDEOGRAM_API_URL, model: config.IDEOGRAM_MODEL });
    seenKeys.add(config.IDEOGRAM_API_KEY);
  }
  if (keyOverride?.key) {
    ideogramKeys.unshift({ key: keyOverride.key, url: keyOverride.url, model: keyOverride.model });
  }

  if (ideogramKeys.length === 0) {
    throw new Error(`No active API key configured for Ideogram.\nPlease configure one in Super Admin Settings.`);
  }

  // Map sizes for Ideogram
  const ASPECT_RATIO_MAP: Record<string, string> = {
    '1024x1024': 'ASPECT_1_1',
    '1536x1024': 'ASPECT_16_9',
    '1024x1536': 'ASPECT_9_16',
    '1024x768': 'ASPECT_4_3',
    '768x1024': 'ASPECT_3_4',
    '1:1': 'ASPECT_1_1',
    '16:9': 'ASPECT_16_9',
    '9:16': 'ASPECT_9_16',
    '4:3': 'ASPECT_4_3',
    '3:4': 'ASPECT_3_4',
  };
  const aspectRatio = ASPECT_RATIO_MAP[size] || 'ASPECT_1_1';

  const timeout = config.AI_TIMEOUT || 180000;
  let lastError = '';

  for (const ideogramKey of ideogramKeys) {
    const ideogramBaseUrl = (ideogramKey.url || config.IDEOGRAM_API_URL || 'https://api.ideogram.ai/api/ideogram/v3')
      .replace(/\/$/, '');
    const ideogramUrl = `${ideogramBaseUrl}/ideograms`;
    const ideogramModel = ideogramKey.model || model || config.IDEOGRAM_MODEL || 'ideogram-v3';

    console.log(`[ImageGen-Ideogram] Trying model=${ideogramModel}, aspect_ratio=${aspectRatio}, key=****${ideogramKey.key.slice(-4)}, url=${ideogramBaseUrl}`);

    try {
      const controller = new AbortController();
      const timeoutId = setTimeout(() => controller.abort(), timeout);

      const body: any = {
        model: ideogramModel,
        prompt,
        aspect_ratio: aspectRatio,
        output_format: 'png',
      };

      const response = await fetch(ideogramUrl, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'Api-Key': ideogramKey.key,
        },
        body: JSON.stringify(body),
        signal: controller.signal,
      });

      clearTimeout(timeoutId);

      if (!response.ok) {
        const errorText = await response.text();
        lastError = `Ideogram API error (${response.status}) with model=${ideogramModel}: ${errorText}`;
        console.warn(`[ImageGen-Ideogram] ${lastError}`);

        if (response.status === 401 || response.status === 403) {
          continue;
        }
        if (response.status >= 500) {
          continue;
        }
        throw new Error(lastError);
      }

      const data: any = await response.json();

      // Ideogram returns: { data: [{ url, ... }] } or { data: [{ b64_json, ... }] }
      let base64 = '';
      let revisedPrompt = '';

      if (data.data?.[0]) {
        const image = data.data[0];
        revisedPrompt = image.revised_prompt || image.prompt || '';
        if (image.b64_json || image.base64_data) {
          base64 = image.b64_json || image.base64_data;
        } else if (image.url) {
          try {
            const urlResponse = await fetch(image.url, { signal: AbortSignal.timeout(30000) });
            if (urlResponse.ok) {
              const arrayBuffer = await urlResponse.arrayBuffer();
              base64 = Buffer.from(arrayBuffer).toString('base64');
              console.log(`[ImageGen-Ideogram] Downloaded image from URL, base64 length: ${base64.length}`);
            }
          } catch (downloadErr: any) {
            console.error(`[ImageGen-Ideogram] Error downloading image URL: ${downloadErr.message}`);
          }
        }
      }

      if (!base64) {
        console.error('[ImageGen-Ideogram] No base64 data found. Response:', JSON.stringify(data).substring(0, 500));
        throw new Error('No image data found in Ideogram response');
      }

      const latencyMs = Date.now() - callStartTime;
      const tokenUsage = {
        inputTokens: data.usage?.input_tokens || data.usage?.prompt_tokens || 0,
        outputTokens: data.usage?.output_tokens || 0,
        totalTokens: (data.usage?.input_tokens || data.usage?.prompt_tokens || 0) + (data.usage?.output_tokens || 0),
      };

      console.log(`[ImageGen-Ideogram] Image generated successfully with ${ideogramModel} in ${latencyMs}ms | base64 length=${base64.length}`);

      return {
        base64Data: base64,
        revisedPrompt,
        model: ideogramModel,
        provider: 'ideogram',
        tokenUsage,
        latencyMs,
      };
    } catch (error: any) {
      if (error.name === 'AbortError') {
        throw new Error(`Ideogram request timed out after ${timeout / 1000}s`);
      }
      lastError = error.message;
      console.warn(`[ImageGen-Ideogram] Error with model=${model}: ${error.message}`);
      continue;
    }
  }

  throw new Error(lastError || `Ideogram image generation failed: No keys available`);
}

// ============================================
// Midjourney Image Generation (via configurable proxy API)
// ============================================

/**
 * Generate an image using Midjourney via a configurable proxy API.
 * The proxy service handles the actual Midjourney interaction.
 * Uses Bearer token auth and OpenAI-compatible request format.
 */
export async function generateImageWithMidjourney(
  prompt: string,
  size: string,
  model: string = 'midjourney-v6',
  keyOverride?: any,
  userId?: string,
): Promise<{ base64Data: string; revisedPrompt: string; model: string; provider: string; tokenUsage: { inputTokens: number; outputTokens: number; totalTokens: number }; latencyMs: number }> {
  const config = await getAIConfig(userId);
  const callStartTime = Date.now();

  const midjourneyKeys: Array<{ key: string; url?: string; model?: string }> = [];
  const seenKeys = new Set<string>();

  if (config.MIDJOURNEY_KEY_LIST && config.MIDJOURNEY_KEY_LIST.length > 0) {
    for (const keyEntry of config.MIDJOURNEY_KEY_LIST) {
      if (keyEntry.key && !seenKeys.has(keyEntry.key)) {
        midjourneyKeys.push({ key: keyEntry.key, url: keyEntry.url, model: keyEntry.model });
        seenKeys.add(keyEntry.key);
      }
    }
  }
  if (config.MIDJOURNEY_API_KEY && !seenKeys.has(config.MIDJOURNEY_API_KEY)) {
    midjourneyKeys.push({ key: config.MIDJOURNEY_API_KEY, url: config.MIDJOURNEY_API_URL, model: config.MIDJOURNEY_MODEL });
    seenKeys.add(config.MIDJOURNEY_API_KEY);
  }
  if (keyOverride?.key) {
    midjourneyKeys.unshift({ key: keyOverride.key, url: keyOverride.url, model: keyOverride.model });
  }

  if (midjourneyKeys.length === 0) {
    throw new Error(`No active API key configured for Midjourney.\nPlease configure one in Super Admin Settings.`);
  }

  const timeout = config.AI_TIMEOUT || 180000;
  let lastError = '';

  for (const mjKey of midjourneyKeys) {
    const mjBaseUrl = (mjKey.url || config.MIDJOURNEY_API_URL || 'https://api.midjourney.com/v1')
      .replace(/\/chat\/completions$/, '')
      .replace(/\/$/, '');
    const imagesUrl = `${mjBaseUrl}/images/generations`;
    const mjModel = mjKey.model || model || config.MIDJOURNEY_MODEL || 'midjourney-v6';

    console.log(`[ImageGen-Midjourney] Trying model=${mjModel}, size=${size}, key=****${mjKey.key.slice(-4)}, url=${mjBaseUrl}`);

    try {
      const controller = new AbortController();
      const timeoutId = setTimeout(() => controller.abort(), timeout);

      const body: any = {
        model: mjModel,
        prompt,
        size: size || '1024x1024',
        n: 1,
      };

      const response = await fetch(imagesUrl, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'Authorization': `Bearer ${mjKey.key}`,
        },
        body: JSON.stringify(body),
        signal: controller.signal,
      });

      clearTimeout(timeoutId);

      if (!response.ok) {
        const errorText = await response.text();
        lastError = `Midjourney API error (${response.status}) with model=${mjModel}: ${errorText}`;
        console.warn(`[ImageGen-Midjourney] ${lastError}`);

        if (response.status === 401 || response.status === 403) {
          continue;
        }
        if (response.status >= 500) {
          continue;
        }
        throw new Error(lastError);
      }

      const data: any = await response.json();

      // Try multiple response formats (OpenAI-compatible, proxy-specific)
      let base64 = '';
      let revisedPrompt = '';

      if (data.data?.[0]) {
        const image = data.data[0];
        base64 = image.b64_json || image.base64_data || '';
        revisedPrompt = image.revised_prompt || '';
        if (!base64 && image.url) {
          try {
            const urlResponse = await fetch(image.url, { signal: AbortSignal.timeout(30000) });
            if (urlResponse.ok) {
              const arrayBuffer = await urlResponse.arrayBuffer();
              base64 = Buffer.from(arrayBuffer).toString('base64');
              console.log(`[ImageGen-Midjourney] Downloaded image from URL, base64 length: ${base64.length}`);
            }
          } catch (downloadErr: any) {
            console.error(`[ImageGen-Midjourney] Error downloading image URL: ${downloadErr.message}`);
          }
        }
      } else if (data.image) {
        base64 = data.image.b64_json || data.image.base64_data || data.image.data || '';
        revisedPrompt = data.image.revised_prompt || data.revised_prompt || '';
        if (!base64 && data.image.url) {
          try {
            const urlResponse = await fetch(data.image.url, { signal: AbortSignal.timeout(30000) });
            if (urlResponse.ok) {
              const arrayBuffer = await urlResponse.arrayBuffer();
              base64 = Buffer.from(arrayBuffer).toString('base64');
            }
          } catch { /* fallback to error */ }
        }
      } else if (data.url) {
        // Direct URL response
        try {
          const urlResponse = await fetch(data.url, { signal: AbortSignal.timeout(30000) });
          if (urlResponse.ok) {
            const arrayBuffer = await urlResponse.arrayBuffer();
            base64 = Buffer.from(arrayBuffer).toString('base64');
          }
        } catch { /* fallback */ }
      } else if (data.images?.[0]) {
        const image = data.images[0];
        base64 = image.b64_json || image.base64_data || '';
        if (!base64 && image.url) {
          try {
            const urlResponse = await fetch(image.url, { signal: AbortSignal.timeout(30000) });
            if (urlResponse.ok) {
              const arrayBuffer = await urlResponse.arrayBuffer();
              base64 = Buffer.from(arrayBuffer).toString('base64');
            }
          } catch { /* fallback */ }
        }
      }

      if (!base64) {
        console.error('[ImageGen-Midjourney] No base64 data found. Response:', JSON.stringify(data).substring(0, 500));
        throw new Error('No image data found in Midjourney response');
      }

      const latencyMs = Date.now() - callStartTime;
      const tokenUsage = {
        inputTokens: data.usage?.input_tokens || data.usage?.prompt_tokens || 0,
        outputTokens: data.usage?.output_tokens || 0,
        totalTokens: (data.usage?.input_tokens || data.usage?.prompt_tokens || 0) + (data.usage?.output_tokens || 0),
      };

      console.log(`[ImageGen-Midjourney] Image generated successfully with ${mjModel} in ${latencyMs}ms | base64 length=${base64.length}`);

      return {
        base64Data: base64,
        revisedPrompt,
        model: mjModel,
        provider: 'midjourney',
        tokenUsage,
        latencyMs,
      };
    } catch (error: any) {
      if (error.name === 'AbortError') {
        throw new Error(`Midjourney request timed out after ${timeout / 1000}s`);
      }
      lastError = error.message;
      console.warn(`[ImageGen-Midjourney] Error with model=${model}: ${error.message}`);
      continue;
    }
  }

  throw new Error(lastError || `Midjourney image generation failed: No keys available`);
}

// ============================================
// Nano Image Generation (OpenAI-compatible format)
// ============================================

/**
 * Generate an image using Nano (lightweight image generation API).
 * Uses OpenAI-compatible /v1/images/generations endpoint format.
 */
export async function generateImageWithNano(
  prompt: string,
  size: string,
  model: string = 'nano-v1',
  keyOverride?: any,
  userId?: string,
): Promise<{ base64Data: string; revisedPrompt: string; model: string; provider: string; tokenUsage: { inputTokens: number; outputTokens: number; totalTokens: number }; latencyMs: number }> {
  const config = await getAIConfig(userId);
  const callStartTime = Date.now();

  const nanoKeys: Array<{ key: string; url?: string; model?: string }> = [];
  const seenKeys = new Set<string>();

  if (config.NANO_KEY_LIST && config.NANO_KEY_LIST.length > 0) {
    for (const keyEntry of config.NANO_KEY_LIST) {
      if (keyEntry.key && !seenKeys.has(keyEntry.key)) {
        nanoKeys.push({ key: keyEntry.key, url: keyEntry.url, model: keyEntry.model });
        seenKeys.add(keyEntry.key);
      }
    }
  }
  if (config.NANO_API_KEY && !seenKeys.has(config.NANO_API_KEY)) {
    nanoKeys.push({ key: config.NANO_API_KEY, url: config.NANO_API_URL, model: config.NANO_MODEL });
    seenKeys.add(config.NANO_API_KEY);
  }
  if (keyOverride?.key) {
    nanoKeys.unshift({ key: keyOverride.key, url: keyOverride.url, model: keyOverride.model });
  }

  if (nanoKeys.length === 0) {
    throw new Error(`No active API key configured for Nano.\nPlease configure one in Super Admin Settings.`);
  }

  // Map sizes for Nano
  const NANO_SIZE_MAP: Record<string, string> = {
    '1024x1024': '1024x1024',
    '1536x1024': '1536x1024',
    '1024x1536': '1024x1536',
    '1024x768': '1024x768',
    '768x1024': '768x1024',
    '1:1': '1024x1024',
    '16:9': '1536x1024',
    '9:16': '1024x1536',
  };
  const nanoSize = NANO_SIZE_MAP[size] || size || '1024x1024';

  const timeout = config.AI_TIMEOUT || 180000;
  let lastError = '';

  for (const nanoKey of nanoKeys) {
    const nanoBaseUrl = (nanoKey.url || config.NANO_API_URL || 'https://api.nano.ai/v1')
      .replace(/\/chat\/completions$/, '')
      .replace(/\/$/, '');
    const imagesUrl = `${nanoBaseUrl}/images/generations`;
    const nanoModel = nanoKey.model || model || config.NANO_MODEL || 'nano-v1';

    console.log(`[ImageGen-Nano] Trying model=${nanoModel}, size=${nanoSize}, key=****${nanoKey.key.slice(-4)}, url=${nanoBaseUrl}`);

    try {
      const controller = new AbortController();
      const timeoutId = setTimeout(() => controller.abort(), timeout);

      const body: any = {
        model: nanoModel,
        prompt,
        size: nanoSize,
        n: 1,
      };

      const response = await fetch(imagesUrl, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'Authorization': `Bearer ${nanoKey.key}`,
        },
        body: JSON.stringify(body),
        signal: controller.signal,
      });

      clearTimeout(timeoutId);

      if (!response.ok) {
        const errorText = await response.text();
        lastError = `Nano API error (${response.status}) with model=${nanoModel}: ${errorText}`;
        console.warn(`[ImageGen-Nano] ${lastError}`);

        if (response.status === 401 || response.status === 403) {
          continue;
        }
        if (response.status >= 500) {
          continue;
        }
        throw new Error(lastError);
      }

      const data: any = await response.json();

      // Try multiple response formats (OpenAI-compatible)
      let base64 = '';
      let revisedPrompt = '';

      if (data.data?.[0]) {
        const image = data.data[0];
        base64 = image.b64_json || image.base64_data || '';
        revisedPrompt = image.revised_prompt || '';
        if (!base64 && image.url) {
          try {
            const urlResponse = await fetch(image.url, { signal: AbortSignal.timeout(30000) });
            if (urlResponse.ok) {
              const arrayBuffer = await urlResponse.arrayBuffer();
              base64 = Buffer.from(arrayBuffer).toString('base64');
              console.log(`[ImageGen-Nano] Downloaded image from URL, base64 length: ${base64.length}`);
            }
          } catch (downloadErr: any) {
            console.error(`[ImageGen-Nano] Error downloading image URL: ${downloadErr.message}`);
          }
        }
      } else if (data.output) {
        // gpt-image-1 style response
        for (const outputItem of data.output) {
          if (outputItem.content) {
            for (const contentItem of outputItem.content) {
              if (contentItem.type === 'image_url' && contentItem.url) {
                if (contentItem.url.startsWith('data:')) {
                  const base64Match = contentItem.url.match(/^data:image\/\w+;base64,(.+)$/);
                  if (base64Match) {
                    base64 = base64Match[1];
                  }
                } else {
                  try {
                    const urlResponse = await fetch(contentItem.url, { signal: AbortSignal.timeout(30000) });
                    if (urlResponse.ok) {
                      const arrayBuffer = await urlResponse.arrayBuffer();
                      base64 = Buffer.from(arrayBuffer).toString('base64');
                    }
                  } catch { /* fallback */ }
                }
              }
            }
          }
        }
      }

      if (!base64) {
        console.error('[ImageGen-Nano] No base64 data found. Response:', JSON.stringify(data).substring(0, 500));
        throw new Error('No image data found in Nano response');
      }

      const latencyMs = Date.now() - callStartTime;
      const tokenUsage = {
        inputTokens: data.usage?.input_tokens || data.usage?.prompt_tokens || 0,
        outputTokens: data.usage?.output_tokens || data.usage?.completion_tokens || 0,
        totalTokens: (data.usage?.input_tokens || data.usage?.prompt_tokens || 0) + (data.usage?.output_tokens || data.usage?.completion_tokens || 0),
      };

      console.log(`[ImageGen-Nano] Image generated successfully with ${nanoModel} in ${latencyMs}ms | base64 length=${base64.length}`);

      return {
        base64Data: base64,
        revisedPrompt,
        model: nanoModel,
        provider: 'nano',
        tokenUsage,
        latencyMs,
      };
    } catch (error: any) {
      if (error.name === 'AbortError') {
        throw new Error(`Nano request timed out after ${timeout / 1000}s`);
      }
      lastError = error.message;
      console.warn(`[ImageGen-Nano] Error with model=${model}: ${error.message}`);
      continue;
    }
  }

  throw new Error(lastError || `Nano image generation failed: No keys available`);
}

// ============================================
// IMAGE GENERATION PROVIDER ROUTER
// ============================================

/**
 * Unified dispatcher: route to the correct image generation function based on provider.
 */
export async function generateImageForProvider(
  provider: ImageProvider,
  model: string,
  prompt: string,
  size: string,
  quality: string,
  style: string,
  keyOverride?: any,
  userId?: string,
): Promise<{ base64Data: string; revisedPrompt: string; model: string; provider: string; tokenUsage: { inputTokens: number; outputTokens: number; totalTokens: number }; latencyMs: number }> {
  switch (provider) {
    case 'openai':
      return generateImageWithOpenAI(prompt, size, quality, style, model, keyOverride, userId);
    case 'zhipu':
      return generateImageWithZhipuCogView(prompt, size, userId);
    case 'flux':
      return generateImageWithFlux(prompt, size, model, keyOverride, userId);
    case 'midjourney':
      return generateImageWithMidjourney(prompt, size, model, keyOverride, userId);
    case 'ideogram':
      return generateImageWithIdeogram(prompt, size, model, keyOverride, userId);
    case 'nano':
      return generateImageWithNano(prompt, size, model, keyOverride, userId);
    default:
      throw new Error(`Unknown image generation provider: ${provider}`);
  }
}

// ============================================
// ROUTES
// ============================================
// IMPORTANT: Specific path routes MUST come before the /:companyId param route
// to avoid Express matching "detail", "enhance-prompt", etc. as companyId values.

/**
 * GET /prompts/:assetCategory — Get available prompt variants for an asset category.
 * Used by the frontend prompt selection UI when generating images.
 * Returns all active variants for both asset_category_guidance and asset_content_elements.
 */
router.get('/prompts/:assetCategory', async (req: Request, res: Response) => {
  try {
    const PromptConfig = getModels().PromptConfig;
    const { assetCategory } = req.params;

    // Fetch both guidance and content elements variants for this category
    const [guidanceVariants, contentVariants] = await Promise.all([
      PromptConfig.find({
        type: 'asset_category_guidance',
        key: assetCategory,
        isActive: true,
      }).select('_id type key name label prompt isDefaultVariant updatedAt').sort({ isDefaultVariant: -1, name: 1 }).lean(),
      PromptConfig.find({
        type: 'asset_content_elements',
        key: assetCategory,
        isActive: true,
      }).select('_id type key name label prompt isDefaultVariant updatedAt').sort({ isDefaultVariant: -1, name: 1 }).lean(),
    ]);

    res.json({
      category: assetCategory,
      guidancePrompts: guidanceVariants,
      contentElementPrompts: contentVariants,
    });
  } catch (error: any) {
    console.error('[ImageGen] Get prompts by category error:', error.message);
    res.status(500).json({ error: 'Failed to fetch prompts' });
  }
});

/**
 * GET /detail/:id — Get a single image generation with all versions
 */
router.get('/detail/:id', async (req: Request, res: Response) => {
  try {
    const { id } = req.params;
    const { ImageGeneration } = getModels();

    const generation = await ImageGeneration.findById(id).lean();
    if (!generation) {
      res.status(404).json({ error: 'Image generation not found' });
      return;
    }

    if (!authorizeCompany(req, generation.companyId)) {
      res.status(403).json({ error: 'Access denied' });
      return;
    }

    res.json(generation);
  } catch (error: any) {
    handleError(res, error);
  }
});

/**
 * GET /:companyId — List all image generations for a company
 * Supports: ?search=, ?style=, ?platform=, ?status=
 * NOTE: This MUST come after all other specific GET routes.
 */
router.get('/:companyId', async (req: Request, res: Response) => {
  try {
    const { companyId } = req.params;
    if (!authorizeCompany(req, companyId)) {
      res.status(403).json({ error: 'Access denied' });
      return;
    }

    const { ImageGeneration } = getModels();

    // Build filter
    const filter: any = { companyId };
    const { search, style, platform, status } = req.query;
    if (style) filter.style = style;
    if (platform) filter.platform = platform;
    if (status) filter.status = status;
    // Trim surrounding whitespace and escape regex metacharacters so that
    // arbitrary user input (e.g. "Logo (v2)", "*", "a+b") is matched as a
    // literal substring instead of being compiled as a raw — and often
    // invalid — regular expression, which previously crashed with a 500.
    const searchTerm = typeof search === 'string' ? search.trim() : '';
    if (searchTerm) {
      const safe = escapeRegex(searchTerm);
      filter.$or = [
        { name: { $regex: safe, $options: 'i' } },
        { description: { $regex: safe, $options: 'i' } },
      ];
    }

    const generations = await ImageGeneration.find(filter)
      .select('-versions.base64Data') // Exclude base64 from list view for performance
      .sort({ createdAt: -1 })
      .lean();

    res.json(generations);
  } catch (error: any) {
    handleError(res, error);
  }
});

/**
 * POST / — Create a new image generation record (metadata only, no generation yet)
 */
router.post(
  '/',
  requirePermission('ai-processing', 'ai-generate'),
  [
    body('companyId').notEmpty().withMessage('Company ID is required'),
    body('description').notEmpty().withMessage('Description is required'),
    body('style').notEmpty().withMessage('Style is required'),
    body('platform').notEmpty().withMessage('Platform is required'),
    body('aspectRatio').notEmpty().withMessage('Aspect ratio is required'),
  ],
  async (req: Request, res: Response) => {
    try {
      const errors = validationResult(req);
      if (!errors.isEmpty()) {
        res.status(400).json({ errors: errors.array() });
        return;
      }

      if (!authorizeCompany(req, req.body.companyId)) {
        res.status(403).json({ error: 'Access denied' });
        return;
      }

      const { ImageGeneration } = getModels();

      // Compose original prompt from form fields
      const { description, objective, style, targetAudience, platform, aspectRatio, additionalInstructions } = req.body;
      // Normalize platform to title case for consistency (e.g. 'instagram' → 'Instagram')
      const normalizedPlatform = platform ? platform.charAt(0).toUpperCase() + platform.slice(1).toLowerCase() : platform;
      const originalPromptParts = [description];
      if (objective) originalPromptParts.push(`Purpose: ${objective}`);
      originalPromptParts.push(`Style: ${style}`);
      if (targetAudience) originalPromptParts.push(`Audience: ${targetAudience}`);
      originalPromptParts.push(`Platform: ${platform}`);
      originalPromptParts.push(`Aspect ratio: ${aspectRatio}`);
      if (additionalInstructions) originalPromptParts.push(additionalInstructions);
      const originalPrompt = originalPromptParts.join('. ') + '.';

      const generation = new (ImageGeneration as any)({
        companyId: req.body.companyId,
        name: req.body.name,
        description,
        objective,
        style,
        targetAudience,
        platform: normalizedPlatform,
        aspectRatio,
        additionalInstructions,
        originalPrompt,
        status: 'pending',
        tags: req.body.tags || [],
        createdBy: req.user!.id,
      });

      await generation.save();
      res.status(201).json(generation);
    } catch (error: any) {
      handleError(res, error);
    }
  }
);

/**
 * POST /enhance-prompt — Ollama prompt enhancement
 * Returns the enhanced prompt without saving (client decides whether to use it).
 */
router.post(
  '/enhance-prompt',
  requirePermission('ai-processing', 'ai-generate'),
  async (req: Request, res: Response) => {
    try {
      const { companyId, description, objective, style, targetAudience, platform, aspectRatio, additionalInstructions, userInstructions, assetTypeCategory, assetCategory, assetRequirements, promptConfigId } = req.body;

      if (!companyId || !description) {
        res.status(400).json({ error: 'companyId and description are required' });
        return;
      }

      if (!authorizeCompany(req, companyId)) {
        res.status(403).json({ error: 'Access denied' });
        return;
      }

      // Extract linked data source IDs from request body
      const linkedData: Record<string, string[] | string | undefined> = {};
      const linkedKeys = ['businessProfileId', 'brandStrategyIds', 'visualIdentityIds'];
      for (const key of linkedKeys) {
        if (req.body[key] !== undefined) {
          linkedData[key] = req.body[key];
        }
      }
      // Also include any other linked data keys that were spread into the body
      for (const [key, value] of Object.entries(req.body)) {
        if ((key.endsWith('Id') || key.endsWith('Ids')) && value !== undefined && !linkedKeys.includes(key)) {
          linkedData[key] = value as string[] | string | undefined;
        }
      }

      // Fetch brand context enriched with linked data
      const brandContext = await fetchBrandContext(companyId, linkedData);

      // Build prompts
      const inputs: ImageEnhancementInputs = {
        description,
        objective,
        style,
        targetAudience,
        platform,
        aspectRatio,
        additionalInstructions,
        userInstructions,
        assetTypeCategory,
        assetCategory,
        assetRequirements,
        ...brandContext,
      };

      // Load prompt overrides from MongoDB (admin-editable prompts)
      const [styleGuidance, platformGuidance, assetCategoryGuidance, assetContentElements] = await Promise.all([
        getStyleGuidance().catch(() => null),
        getPlatformGuidance().catch(() => null),
        getAssetCategoryGuidance().catch(() => null),
        getAssetContentElements().catch(() => null),
      ]);
      const promptOverrides = {
        ...(styleGuidance && { styleGuidance }),
        ...(platformGuidance && { platformGuidance }),
        ...(assetCategoryGuidance && { assetCategoryGuidance }),
        ...(assetContentElements && { assetContentElements }),
      };

      // If a specific promptConfigId was provided, load that variant and override
      // the relevant prompt maps with the selected variant's content
      if (promptConfigId) {
        try {
          const { getPromptVariantById } = await import('../services/aiContext/promptConfigLoader');
          const selectedVariant = await getPromptVariantById(promptConfigId);
          if (selectedVariant) {
            // Override the appropriate map based on the variant's type
            if (selectedVariant.type === 'asset_category_guidance' && assetCategoryGuidance) {
              // Replace the guidance for this specific key with the selected variant
              const updatedGuidance = { ...assetCategoryGuidance };
              updatedGuidance[selectedVariant.key] = selectedVariant.prompt;
              promptOverrides.assetCategoryGuidance = updatedGuidance;
            } else if (selectedVariant.type === 'asset_content_elements' && assetContentElements) {
              const updatedElements = { ...assetContentElements };
              updatedElements[selectedVariant.key] = selectedVariant.prompt;
              promptOverrides.assetContentElements = updatedElements;
            }
          }
        } catch (err: any) {
          console.warn('[ImageGen] Failed to load selected prompt variant, using defaults:', err.message);
        }
      }

      const { systemPrompt, userPrompt, maxTokens } = buildEnhancementPrompts(inputs, promptOverrides);

      try {
        // Call Ollama (preferred) for prompt enhancement
        const result = await generateWithAI(userPrompt, systemPrompt, maxTokens, 0.7, 'text', 'ollama');

        res.json({
          enhancedPrompt: stripReasoning(result.content.trim()),
          originalPrompt: buildFallbackPrompt(inputs, promptOverrides),
          model: result.model,
          provider: result.provider,
          tokenUsage: result.tokenUsage,
          latencyMs: result.latencyMs,
        });
      } catch (ollamaError: any) {
        // Graceful fallback: if Ollama is unavailable, try other providers, then use fallback prompt
        console.warn(`[ImageGen] Ollama prompt enhancement failed: ${ollamaError.message}`);

        try {
          // Try with auto provider selection (might pick Claude, Zhipu, or OpenAI)
          const result = await generateWithAI(userPrompt, systemPrompt, maxTokens, 0.7, 'text');

          res.json({
            enhancedPrompt: stripReasoning(result.content.trim()),
            originalPrompt: buildFallbackPrompt(inputs, promptOverrides),
            model: result.model,
            provider: result.provider,
            tokenUsage: result.tokenUsage,
            latencyMs: result.latencyMs,
          });
        } catch (aiError: any) {
          // All AI providers failed — return fallback prompt with warning
          console.warn(`[ImageGen] All AI providers failed for prompt enhancement: ${aiError.message}`);
          const fallbackPrompt = buildFallbackPrompt(inputs, promptOverrides);
          res.json({
            enhancedPrompt: fallbackPrompt,
            originalPrompt: fallbackPrompt,
            model: 'fallback',
            provider: 'none',
            warning: 'AI prompt enhancement unavailable. Using structured prompt instead. Configure an AI provider for better results.',
            tokenUsage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
          });
        }
      }
    } catch (error: any) {
      handleError(res, error);
    }
  }
);

/**
 * Composite a brand logo directly onto a generated image using sharp.
 * The logo is overlaid in the top-left header region (a common logo position for
 * stationery) at a proportionate size, with the logo's own aspect ratio preserved
 * and any surrounding transparent/empty space trimmed. The logo is NEVER sent to
 * the AI — this is a pure post-generation overlay. Returns the new base64 PNG, or
 * the original base64 unchanged if anything fails.
 */
async function compositeLogoOntoImage(baseBase64: string, logoBuffer: Buffer): Promise<string> {
  const cleanBase64 = baseBase64.replace(/^data:[^;]+;base64,/, '');
  const baseBuffer = Buffer.from(cleanBase64, 'base64');
  const meta = await sharp(baseBuffer).metadata();
  const width = meta.width || 1024;
  const height = meta.height || 1024;

  // Logo ~20% of the page width, padded in from the top-left corner.
  const targetWidth = Math.max(64, Math.round(width * 0.20));
  const padding = Math.round(Math.min(width, height) * 0.045);

  // Trim empty borders around the logo (best-effort), then resize to the target width.
  let trimmedLogo = logoBuffer;
  try {
    trimmedLogo = await sharp(logoBuffer).trim().toBuffer();
  } catch {
    // Logo has no uniform border to trim (or trim unsupported) — use it as-is.
  }
  const resizedLogo = await sharp(trimmedLogo)
    .resize({ width: targetWidth, withoutEnlargement: false })
    .png()
    .toBuffer();

  const composited = await sharp(baseBuffer)
    .composite([{ input: resizedLogo, top: padding, left: padding }])
    .png()
    .toBuffer();

  return composited.toString('base64');
}

/**
 * Overlay the referenced brand asset (logo) onto a freshly generated image.
 * Loads the asset for the company, reads its bytes, and composites it. Returns the
 * original base64 unchanged when no referenceAssetId is given or anything fails, so
 * generation never breaks because of the overlay step.
 */
async function maybeOverlayBrandLogo(
  base64Data: string,
  referenceAssetId: string | undefined,
  companyId: string,
): Promise<string> {
  if (!referenceAssetId || !base64Data) return base64Data;
  try {
    const { BrandAsset } = getModels();
    const refAsset = await (BrandAsset as any).findOne({ _id: referenceAssetId, companyId });
    const logo = refAsset ? await readBrandAssetImageBytes(refAsset) : null;
    if (!logo) {
      console.warn(`[ImageGen] referenceAssetId ${referenceAssetId} has no usable image — skipping logo overlay`);
      return base64Data;
    }
    const out = await compositeLogoOntoImage(base64Data, logo.buffer);
    console.log(`[ImageGen] Overlaid brand logo "${refAsset.name || refAsset.type}" onto generated image`);
    return out;
  } catch (e: any) {
    console.warn(`[ImageGen] Logo overlay failed (${e?.message}); using generated image without overlay`);
    return base64Data;
  }
}

/**
 * POST /generate/:id — Generate image for an existing generation record
 * Calls OpenAI DALL-E 3 API with the enhanced prompt. When req.body.referenceAssetId
 * is provided, the referenced brand asset (e.g. the logo) is composited directly
 * onto the finished image (post-generation overlay) — the logo is NOT sent to the AI.
 */
router.post(
  '/generate/:id',
  requirePermission('ai-processing', 'ai-generate'),
  async (req: Request, res: Response) => {
    try {
      const { id } = req.params;
      const { ImageGeneration } = getModels();

      const generation = await ImageGeneration.findById(id);
      if (!generation) {
        res.status(404).json({ error: 'Image generation not found' });
        return;
      }

      if (!authorizeCompany(req, generation.companyId)) {
        res.status(403).json({ error: 'Access denied' });
        return;
      }

      // Update status to generating
      generation.status = 'generating';
      await generation.save();

      // Determine which prompt to use
      const prompt = req.body.prompt || generation.ollamaEnhancedPrompt || generation.originalPrompt;
      const quality = req.body.quality || 'standard';
      const style = req.body.style || 'vivid';
      const model = req.body.model || 'gpt-image-1';

      // Determine the image generation provider based on the model
      const provider = getImageProviderForModel(model);

      // Map aspect ratio to DALL-E size
      const size = ASPECT_RATIO_SIZE_MAP[generation.aspectRatio] || '1024x1024';

      try {
        let imageResult = await generateImageForProvider(provider, model, prompt, size, quality, style);
        let fallbackUsed = false;
        let fallbackProvider: string | null = null;

        // Insert the real brand logo directly into the generated image (post-generation
        // overlay). The logo is NOT included in the AI prompt — it is composited here
        // from the actual brand asset. No-op when no referenceAssetId is supplied.
        imageResult.base64Data = await maybeOverlayBrandLogo(imageResult.base64Data, req.body.referenceAssetId, generation.companyId);

        // Mark all previous versions as not current
        for (const v of generation.versions) {
          v.isCurrent = false;
        }

        // Add new version
        const versionNumber = generation.versions.length + 1;
        generation.versions.push({
          version: versionNumber,
          prompt: generation.originalPrompt,
          enhancedPrompt: stripReasoning(prompt),
          imageUrl: '', // DALL-E URLs expire, we store base64
          base64Data: imageResult.base64Data,
          generationProvider: imageResult.provider,
          generationModel: imageResult.model,
          aspectRatio: generation.aspectRatio,
          size,
          quality,
          style,
          revisedPrompt: imageResult.revisedPrompt,
          tokenUsage: (imageResult as any).tokenUsage || { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
          latencyMs: (imageResult as any).latencyMs || 0,
          isCurrent: true,
        } as any);

        generation.currentVersion = versionNumber;
        generation.totalVersions = versionNumber;
        generation.generationProvider = imageResult.provider;
        generation.generationModel = imageResult.model;
        if (!generation.ollamaEnhancedPrompt && prompt !== generation.originalPrompt) {
          generation.ollamaEnhancedPrompt = stripReasoning(prompt);
        }
        generation.status = 'completed';
        await generation.save();

        notifyImageGenerationCompleted(req, {
          generationId: String(generation._id),
          version: versionNumber,
          companyId: generation.companyId,
          name: generation.name,
        });

        const responseObj = generation.toObject ? generation.toObject() : generation;
        if (fallbackUsed) {
          (responseObj as any).fallbackUsed = true;
          (responseObj as any).fallbackProvider = fallbackProvider;
        }
        res.json(responseObj);
      } catch (primaryError: any) {
        // Primary provider failed
        console.warn(`[ImageGen] ${IMAGE_PROVIDER_NAMES[provider] || provider} image generation failed: ${primaryError.message}`);

        // Only attempt CogView-3 fallback when the primary provider is OpenAI
        // (preserving existing fallback behavior). Other providers should NOT
        // auto-fallback — the user chose that provider specifically.
        if (provider !== 'openai') {
          generation.status = 'failed';
          await generation.save();
          const friendlyError = parseOpenAIError(primaryError.message);
          res.status(500).json({ error: friendlyError });
          return;
        }

        // OpenAI failed — try Zhipu CogView-3 fallback
        let imageResult: { base64Data: string; revisedPrompt: string; model: string; provider: string } | null = null;
        let fallbackUsed = false;
        let fallbackProvider: string | null = null;

        try {
          const fallbackConfig = await getAIConfig();
          if (fallbackConfig.ZHIPU_API_KEY || (fallbackConfig.ZHIPU_KEY_LIST && fallbackConfig.ZHIPU_KEY_LIST.length > 0) || (fallbackConfig.OLLAMA_API_KEY && fallbackConfig.OLLAMA_API_KEY.includes('.'))) {
            console.log('[ImageGen] Attempting Zhipu CogView-3 fallback...');
            imageResult = await generateImageWithZhipuCogView(prompt, size);
            fallbackUsed = true;
            fallbackProvider = imageResult.provider;
            console.log(`[ImageGen] Zhipu CogView-3 fallback succeeded (model: ${imageResult.model})`);
          } else {
            console.warn('[ImageGen] No Zhipu key configured for CogView-3 fallback');
          }
        } catch (cogviewError: any) {
          console.error(`[ImageGen] Zhipu CogView-3 fallback also failed: ${cogviewError.message}`);
        }

        if (!imageResult) {
          // Both providers failed — return user-friendly error
          generation.status = 'failed';
          await generation.save();
          const friendlyError = parseOpenAIError(primaryError.message);
          res.status(500).json({ error: friendlyError });
          return;
        }

        // CogView fallback succeeded — overlay the brand logo, then save the result
        imageResult.base64Data = await maybeOverlayBrandLogo(imageResult.base64Data, req.body.referenceAssetId, generation.companyId);

        for (const v of generation.versions) {
          v.isCurrent = false;
        }

        const versionNumber = generation.versions.length + 1;
        generation.versions.push({
          version: versionNumber,
          prompt: generation.originalPrompt,
          enhancedPrompt: stripReasoning(prompt),
          imageUrl: '',
          base64Data: imageResult.base64Data,
          generationProvider: imageResult.provider,
          generationModel: imageResult.model,
          aspectRatio: generation.aspectRatio,
          size,
          quality: 'auto',
          style: 'vivid',
          revisedPrompt: imageResult.revisedPrompt,
          tokenUsage: (imageResult as any).tokenUsage || { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
          latencyMs: (imageResult as any).latencyMs || 0,
          isCurrent: true,
        } as any);

        generation.currentVersion = versionNumber;
        generation.totalVersions = versionNumber;
        generation.generationProvider = imageResult.provider;
        generation.generationModel = imageResult.model;
        if (!generation.ollamaEnhancedPrompt && prompt !== generation.originalPrompt) {
          generation.ollamaEnhancedPrompt = stripReasoning(prompt);
        }
        generation.status = 'completed';
        await generation.save();

        notifyImageGenerationCompleted(req, {
          generationId: String(generation._id),
          version: versionNumber,
          companyId: generation.companyId,
          name: generation.name,
        });

        const responseObj = generation.toObject ? generation.toObject() : generation;
        (responseObj as any).fallbackUsed = true;
        (responseObj as any).fallbackProvider = fallbackProvider;
        res.json(responseObj);
      }
    } catch (error: any) {
      handleError(res, error);
    }
  }
);

/**
 * POST /regenerate/:id — Create a new version and regenerate the image
 */
router.post(
  '/regenerate/:id',
  requirePermission('ai-processing', 'ai-generate'),
  async (req: Request, res: Response) => {
    try {
      const { id } = req.params;
      const { ImageGeneration } = getModels();

      const generation = await ImageGeneration.findById(id);
      if (!generation) {
        res.status(404).json({ error: 'Image generation not found' });
        return;
      }

      if (!authorizeCompany(req, generation.companyId)) {
        res.status(403).json({ error: 'Access denied' });
        return;
      }

      // Use provided prompt, or the last enhanced prompt, or original prompt
      const prompt = req.body.prompt || generation.ollamaEnhancedPrompt || generation.originalPrompt;
      const quality = req.body.quality || 'standard';
      const style = req.body.style || 'vivid';
      const model = req.body.model || 'gpt-image-1';

      // Determine the image generation provider based on the model
      const provider = getImageProviderForModel(model);

      const size = ASPECT_RATIO_SIZE_MAP[generation.aspectRatio] || '1024x1024';

      generation.status = 'generating';
      await generation.save();

      try {
        let imageResult = await generateImageForProvider(provider, model, prompt, size, quality, style);
        let fallbackUsed = false;
        let fallbackProvider: string | null = null;

        // Mark all previous versions as not current
        for (const v of generation.versions) {
          v.isCurrent = false;
        }

        const versionNumber = generation.totalVersions + 1;
        generation.versions.push({
          version: versionNumber,
          prompt: generation.originalPrompt,
          enhancedPrompt: stripReasoning(prompt),
          imageUrl: '',
          base64Data: imageResult.base64Data,
          generationProvider: imageResult.provider,
          generationModel: imageResult.model,
          aspectRatio: generation.aspectRatio,
          size,
          quality,
          style,
          revisedPrompt: imageResult.revisedPrompt,
          tokenUsage: (imageResult as any).tokenUsage || { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
          latencyMs: (imageResult as any).latencyMs || 0,
          isCurrent: true,
        } as any);

        generation.currentVersion = versionNumber;
        generation.totalVersions = versionNumber;
        generation.generationProvider = imageResult.provider;
        generation.generationModel = imageResult.model;
        generation.status = 'completed';
        await generation.save();

        notifyImageGenerationCompleted(req, {
          generationId: String(generation._id),
          version: versionNumber,
          companyId: generation.companyId,
          name: generation.name,
        });

        const responseObj = generation.toObject ? generation.toObject() : generation;
        if (fallbackUsed) {
          (responseObj as any).fallbackUsed = true;
          (responseObj as any).fallbackProvider = fallbackProvider;
        }
        res.json(responseObj);
      } catch (primaryError: any) {
        // Primary provider failed
        console.warn(`[ImageGen] ${IMAGE_PROVIDER_NAMES[provider] || provider} regeneration failed: ${primaryError.message}`);

        // Only attempt CogView-3 fallback when the primary provider is OpenAI
        if (provider !== 'openai') {
          generation.status = 'failed';
          await generation.save();
          const friendlyError = parseOpenAIError(primaryError.message);
          res.status(500).json({ error: friendlyError });
          return;
        }

        // OpenAI failed — try Zhipu CogView-3 fallback
        let imageResult: { base64Data: string; revisedPrompt: string; model: string; provider: string } | null = null;
        let fallbackUsed = false;
        let fallbackProvider: string | null = null;

        try {
          const fallbackConfig = await getAIConfig();
          if (fallbackConfig.ZHIPU_API_KEY || (fallbackConfig.ZHIPU_KEY_LIST && fallbackConfig.ZHIPU_KEY_LIST.length > 0) || (fallbackConfig.OLLAMA_API_KEY && fallbackConfig.OLLAMA_API_KEY.includes('.'))) {
            console.log('[ImageGen] Attempting Zhipu CogView-3 fallback for regeneration...');
            imageResult = await generateImageWithZhipuCogView(prompt, size);
            fallbackUsed = true;
            fallbackProvider = imageResult.provider;
            console.log(`[ImageGen] Zhipu CogView-3 fallback succeeded for regeneration (model: ${imageResult.model})`);
          } else {
            console.warn('[ImageGen] No Zhipu key configured for CogView-3 fallback');
          }
        } catch (cogviewError: any) {
          console.error(`[ImageGen] Zhipu CogView-3 fallback also failed for regeneration: ${cogviewError.message}`);
        }

        if (!imageResult) {
          // Both providers failed — return user-friendly error
          generation.status = 'failed';
          await generation.save();
          const friendlyError = parseOpenAIError(primaryError.message);
          res.status(500).json({ error: friendlyError });
          return;
        }

        // CogView fallback succeeded — save the result
        for (const v of generation.versions) {
          v.isCurrent = false;
        }

        const versionNumber = generation.totalVersions + 1;
        generation.versions.push({
          version: versionNumber,
          prompt: generation.originalPrompt,
          enhancedPrompt: stripReasoning(prompt),
          imageUrl: '',
          base64Data: imageResult.base64Data,
          generationProvider: imageResult.provider,
          generationModel: imageResult.model,
          aspectRatio: generation.aspectRatio,
          size,
          quality: 'auto',
          style: 'vivid',
          revisedPrompt: imageResult.revisedPrompt,
          tokenUsage: (imageResult as any).tokenUsage || { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
          latencyMs: (imageResult as any).latencyMs || 0,
          isCurrent: true,
        } as any);

        generation.currentVersion = versionNumber;
        generation.totalVersions = versionNumber;
        generation.generationProvider = imageResult.provider;
        generation.generationModel = imageResult.model;
        generation.status = 'completed';
        await generation.save();

        notifyImageGenerationCompleted(req, {
          generationId: String(generation._id),
          version: versionNumber,
          companyId: generation.companyId,
          name: generation.name,
        });

        const responseObj = generation.toObject ? generation.toObject() : generation;
        (responseObj as any).fallbackUsed = true;
        (responseObj as any).fallbackProvider = fallbackProvider;
        res.json(responseObj);
      }
    } catch (error: any) {
      handleError(res, error);
    }
  }
);

/**
 * PUT /:id — Update generation metadata (name, tags, etc.)
 */
router.put('/:id', requirePermission('ai-processing', 'edit'), async (req: Request, res: Response) => {
  try {
    const { id } = req.params;
    const { ImageGeneration } = getModels();

    const generation = await ImageGeneration.findById(id);
    if (!generation) {
      res.status(404).json({ error: 'Image generation not found' });
      return;
    }

    if (!authorizeCompany(req, generation.companyId)) {
      res.status(403).json({ error: 'Access denied' });
      return;
    }

    // Only allow updating safe fields
    const allowedUpdates = ['name', 'tags', 'ollamaEnhancedPrompt'];
    for (const field of allowedUpdates) {
      if (req.body[field] !== undefined) {
        (generation as any)[field] = req.body[field];
      }
    }

    generation.updatedAt = new Date();
    await generation.save();
    res.json(generation);
  } catch (error: any) {
    handleError(res, error);
  }
});

/**
 * DELETE /:id — Delete an entire image generation
 */
router.delete('/:id', requirePermission('ai-processing', 'delete'), async (req: Request, res: Response) => {
  try {
    const { id } = req.params;
    const { ImageGeneration } = getModels();

    const generation = await ImageGeneration.findById(id);
    if (!generation) {
      res.status(404).json({ error: 'Image generation not found' });
      return;
    }

    if (!authorizeCompany(req, generation.companyId)) {
      res.status(403).json({ error: 'Access denied' });
      return;
    }

    await ImageGeneration.deleteOne({ _id: id });
    res.json({ message: 'Image generation deleted successfully' });
  } catch (error: any) {
    handleError(res, error);
  }
});

/**
 * POST /save-to-brand-assets/:id — Copy current version image to Brand Assets
 */
router.post('/save-to-brand-assets/:id', requirePermission('ai-processing', 'ai-generate'), async (req: Request, res: Response) => {
  try {
    const { id } = req.params;
    const { ImageGeneration, BrandAsset } = getModels();

    const generation = await ImageGeneration.findById(id);
    if (!generation) {
      res.status(404).json({ error: 'Image generation not found' });
      return;
    }

    if (!authorizeCompany(req, generation.companyId)) {
      res.status(403).json({ error: 'Access denied' });
      return;
    }

    // Get current version
    const currentVersion = generation.versions.find((v: any) => v.isCurrent);
    if (!currentVersion || !currentVersion.base64Data) {
      res.status(400).json({ error: 'No image available. Generate an image first.' });
      return;
    }

    // Map platform to asset type
    const platformToAssetType: Record<string, string> = {
      'Instagram': 'web-banner',
      'Facebook': 'web-banner',
      'LinkedIn': 'web-banner',
      'Twitter/X': 'web-banner',
      'Website': 'web-banner',
      'Blog': 'web-banner',
      'Email': 'email-header',
      'Presentation': 'presentation',
      'Print': 'print-ready',
      'YouTube': 'web-banner',
      'TikTok': 'web-banner',
      'Pinterest': 'web-banner',
      'Other': 'other',
    };

    // Save the AI-generated image to disk instead of MongoDB
    const base64WithPrefix = `data:image/png;base64,${currentVersion.base64Data}`;
    const { buffer, mimeType } = base64ToBuffer(base64WithPrefix, 'image/png');
    const ext = getExtensionFromMime(mimeType, 'png');
    const { url: assetUrl, fileSize } = await saveBrandAssetFile(buffer, `ai-generated.${ext}`, mimeType);

    // Create Brand Asset with filesystem URL (no base64Data)
    const assetData = {
      companyId: generation.companyId,
      name: req.body.name || generation.name || `AI Generated — ${generation.description.substring(0, 50)}`,
      // Use explicit type from request body (e.g. from enhanced Add Asset form) or fall back to platform mapping
      type: req.body.type || platformToAssetType[generation.platform] || 'other',
      description: `AI-generated image. Original prompt: ${generation.originalPrompt.substring(0, 200)}`,
      format: ext === 'jpeg' ? 'jpg' : ext,
      url: assetUrl,
      source: 'ai-generation',
      designBrief: currentVersion.enhancedPrompt || generation.ollamaEnhancedPrompt || generation.originalPrompt,
      tags: ['ai-generated', generation.style.toLowerCase(), generation.platform.toLowerCase()],
      isPrimary: false,
      fileSize,
      fileType: mimeType,
    };

    const brandAsset = new (BrandAsset as any)(assetData);
    await brandAsset.save();

    // Update generation with brand asset link
    generation.brandAssetId = (brandAsset as any)._id || (brandAsset as any).id;
    generation.updatedAt = new Date();
    await generation.save();

    res.status(201).json({
      message: 'Image saved to Brand Assets successfully',
      brandAsset,
      generationId: generation._id || generation.id,
    });
  } catch (error: any) {
    handleError(res, error);
  }
});

/**
 * POST /save-to-hr-assets/:id — Copy current version image to HR Assets
 */
router.post('/save-to-hr-assets/:id', requirePermission('ai-processing', 'ai-generate'), async (req: Request, res: Response) => {
  try {
    const { id } = req.params;
    const { ImageGeneration, HRAsset } = getModels();

    const generation = await ImageGeneration.findById(id);
    if (!generation) {
      res.status(404).json({ error: 'Image generation not found' });
      return;
    }

    if (!authorizeCompany(req, generation.companyId)) {
      res.status(403).json({ error: 'Access denied' });
      return;
    }

    // Get current version
    const currentVersion = generation.versions.find((v: any) => v.isCurrent);
    if (!currentVersion || !currentVersion.base64Data) {
      res.status(400).json({ error: 'No image available. Generate an image first.' });
      return;
    }

    // Map HR asset type to category
    const TYPE_TO_CATEGORY: Record<string, string> = {
      // Internal Branding
      'id-card-front': 'internal-branding',
      'id-card-back': 'internal-branding',
      'lanyard-design': 'internal-branding',
      'employee-badge': 'internal-branding',
      'visiting-card': 'internal-branding',
      'attendance-sheet': 'internal-branding',
      'internal-memo': 'internal-branding',
      // Desk & Office
      'notepad': 'desk-office',
      'diary-planner': 'desk-office',
      'file-folder': 'desk-office',
      'document-folder': 'desk-office',
      'pen-branding': 'desk-office',
      'desk-name-plate': 'desk-office',
      // Letters
      'offer-letter': 'letters',
      'relieving-letter': 'letters',
      'increment-letter': 'letters',
      'termination-letter': 'letters',
      'experience-letter': 'letters',
      'appointment-letter': 'letters',
      'promotion-letter': 'letters',
      'warning-letter': 'letters',
      // Certifications
      'experience-certificate': 'certifications',
      'training-certificate': 'certifications',
      'appreciation-certificate': 'certifications',
      'completion-certificate': 'certifications',
      'internship-certificate': 'certifications',
      // Onboarding
      'welcome-kit': 'onboarding',
      'onboarding-checklist': 'onboarding',
      'orientation-presentation': 'onboarding',
      'handbook': 'onboarding',
      'code-of-conduct': 'onboarding',
      // Other
      'other': 'other',
    };

    // Save the AI-generated image to disk
    const base64WithPrefix = `data:image/png;base64,${currentVersion.base64Data}`;
    const { buffer, mimeType } = base64ToBuffer(base64WithPrefix, 'image/png');
    const ext = getExtensionFromMime(mimeType, 'png');
    const { url: assetUrl, fileSize } = await saveHrAssetFile(buffer, `ai-generated.${ext}`, mimeType);

    // Create HR Asset document
    const assetType = req.body.type || 'other';
    const assetCategory = req.body.category || TYPE_TO_CATEGORY[assetType] || 'other';

    const assetData = {
      companyId: generation.companyId,
      name: req.body.name || generation.name || `AI Generated — ${generation.description.substring(0, 50)}`,
      type: assetType,
      category: assetCategory,
      description: `AI-generated image. Original prompt: ${generation.originalPrompt.substring(0, 200)}`,
      templateUrl: assetUrl,
      previewImageUrl: assetUrl,
      sourceUrl: 'ai-generation',
      base64Data: currentVersion.base64Data, // Keep base64 for inline previews
      fileName: `hr-asset-${Date.now()}.${ext}`,
      fileSize,
      fileType: mimeType,
      status: 'draft',
      tags: ['ai-generated', generation.style.toLowerCase(), generation.platform.toLowerCase()],
    };

    const hrAsset = new (HRAsset as any)(assetData);
    await hrAsset.save();

    // Update generation with HR asset link
    generation.hrAssetId = (hrAsset as any)._id || (hrAsset as any).id;
    generation.updatedAt = new Date();
    await generation.save();

    res.status(201).json({
      message: 'Image saved to HR Assets successfully',
      hrAsset,
      generationId: generation._id || generation.id,
    });
  } catch (error: any) {
    handleError(res, error);
  }
});

/**
 * POST /save-to-stationery/:id — Copy current version image to Stationery
 */
router.post('/save-to-stationery/:id', requirePermission('ai-processing', 'ai-generate'), async (req: Request, res: Response) => {
  try {
    const { id } = req.params;
    const { ImageGeneration, Stationery } = getModels();

    const generation = await ImageGeneration.findById(id);
    if (!generation) {
      res.status(404).json({ error: 'Image generation not found' });
      return;
    }

    if (!authorizeCompany(req, generation.companyId)) {
      res.status(403).json({ error: 'Access denied' });
      return;
    }

    // Get current version
    const currentVersion = generation.versions.find((v: any) => v.isCurrent);
    if (!currentVersion || !currentVersion.base64Data) {
      res.status(400).json({ error: 'No image available. Generate an image first.' });
      return;
    }

    // Map Stationery type to category
    const TYPE_TO_CATEGORY: Record<string, string> = {
      // Core Stationery
      'business-card': 'core-stationery',
      'letterhead': 'core-stationery',
      'envelope-a4': 'core-stationery',
      'envelope-dl': 'core-stationery',
      'email-signature': 'core-stationery',
      'presentation-template': 'core-stationery',
      // Office Use Assets
      'invoice-template': 'office-assets',
      'quotation-template': 'office-assets',
      'receipt-design': 'office-assets',
      'purchase-order': 'office-assets',
      'billing-format': 'office-assets',
      'proposal-template': 'office-assets',
      // Packaging Stationery
      'thank-you-card': 'packaging-stationery',
      'warranty-card': 'packaging-stationery',
      'instruction-manual': 'packaging-stationery',
      'product-insert-card': 'packaging-stationery',
      'branded-stickers': 'packaging-stationery',
      'packaging-tape': 'packaging-stationery',
      // Print Stationery
      'stamps': 'print-stationery',
      'branding-print': 'print-stationery',
      'standees-print': 'print-stationery',
      'booth-designs': 'print-stationery',
      't-shirts': 'print-stationery',
      'notebook': 'print-stationery',
      'coffee-mug': 'print-stationery',
      'tote-bag': 'print-stationery',
      // Marketing Assets
      'newsletter-template': 'marketing-assets',
      'brochure-pdf': 'marketing-assets',
      'pitch-deck': 'marketing-assets',
      'tagline': 'marketing-assets',
      'hook-style': 'marketing-assets',
      'standees-marketing': 'marketing-assets',
      'marketing-collateral': 'marketing-assets',
      // Other
      'memo-pad': 'other',
      'folder': 'other',
      'compliment-slip': 'other',
      'envelope': 'other',
      'other': 'other',
    };

    // Save the AI-generated image to disk
    const base64WithPrefix = `data:image/png;base64,${currentVersion.base64Data}`;
    const { buffer, mimeType } = base64ToBuffer(base64WithPrefix, 'image/png');
    const ext = getExtensionFromMime(mimeType, 'png');
    const { url: assetUrl, fileSize } = await saveStationeryFile(buffer, `ai-generated.${ext}`, mimeType);

    // Create Stationery document
    const assetType = req.body.type || 'other';
    const assetCategory = req.body.category || TYPE_TO_CATEGORY[assetType] || 'other';

    const assetData = {
      companyId: generation.companyId,
      name: req.body.name || generation.name || `AI Generated — ${generation.description.substring(0, 50)}`,
      type: assetType,
      description: `AI-generated image. Original prompt: ${generation.originalPrompt.substring(0, 200)}`,
      category: assetCategory,
      templateUrl: assetUrl,
      previewImageUrl: assetUrl,
      sourceUrl: 'ai-generation',
      base64Data: currentVersion.base64Data,
      fileName: `stationery-${Date.now()}.${ext}`,
      fileSize,
      fileType: mimeType,
      status: 'draft',
      tags: ['ai-generated', generation.style.toLowerCase(), generation.platform.toLowerCase()],
    };

    const stationeryItem = new (Stationery as any)(assetData);
    await stationeryItem.save();

    // Update generation with stationery link
    generation.stationeryAssetId = (stationeryItem as any)._id || (stationeryItem as any).id;
    generation.updatedAt = new Date();
    await generation.save();

    res.status(201).json({
      message: 'Image saved to Stationery successfully',
      stationery: stationeryItem,
      generationId: generation._id || generation.id,
    });
  } catch (error: any) {
    handleError(res, error);
  }
});

/**
 * DELETE /version/:id/:version — Delete a specific version
 */
router.delete('/version/:id/:version', requirePermission('ai-processing', 'delete'), async (req: Request, res: Response) => {
  try {
    const { id, version } = req.params;
    const versionNumber = parseInt(version, 10);
    const { ImageGeneration } = getModels();

    const generation = await ImageGeneration.findById(id);
    if (!generation) {
      res.status(404).json({ error: 'Image generation not found' });
      return;
    }

    if (!authorizeCompany(req, generation.companyId)) {
      res.status(403).json({ error: 'Access denied' });
      return;
    }

    const versionIndex = generation.versions.findIndex((v: any) => v.version === versionNumber);
    if (versionIndex === -1) {
      res.status(404).json({ error: 'Version not found' });
      return;
    }

    // If deleting the current version, set the most recent remaining version as current
    const wasCurrent = generation.versions[versionIndex].isCurrent;
    generation.versions.splice(versionIndex, 1);

    if (wasCurrent && generation.versions.length > 0) {
      generation.versions[generation.versions.length - 1].isCurrent = true;
      generation.currentVersion = generation.versions[generation.versions.length - 1].version;
    } else if (generation.versions.length === 0) {
      generation.currentVersion = 0;
    }

    generation.totalVersions = generation.versions.length;
    generation.updatedAt = new Date();
    await generation.save();

    res.json(generation);
  } catch (error: any) {
    handleError(res, error);
  }
});

// ============================================
// STATIONERY — NON-BLOCKING BACKGROUND GENERATION
// ============================================

/**
 * Output format of a stationery asset.
 *
 *  - 'html' : EDITABLE TEMPLATE assets (business card, letterhead, envelopes,
 *             email signature, the office documents, the packaging cards and the
 *             newsletter). These are text-bearing layouts whose whole point is
 *             that the wording stays editable after generation, so they are
 *             produced as clean, standalone HTML — never flattened to an image.
 *  - 'png'  : GRAPHIC assets (stickers, tape, stamps, print branding, standees,
 *             booth panels, merchandise, the presentation slide template). These
 *             are artwork, not documents, so they keep the existing image
 *             generation path unchanged.
 */
type StationeryOutputFormat = 'html' | 'png';

interface StationeryGenConfig {
  label: string;
  promptHint: string;
  style: string;
  platform: string;
  aspectRatio: string;
  assetTypeCategory: string;
  /** Which format this asset type is produced in — see StationeryOutputFormat. */
  outputFormat: StationeryOutputFormat;
}

/**
 * Per-category generation config for stationery. Drives the standard size
 * (via aspectRatio → ASPECT_RATIO_SIZE_MAP), art style, platform, which
 * asset-category guidance the prompt builder applies, and the OUTPUT FORMAT.
 * Documents/core stationery are print-ready flat layouts; merchandise stays
 * flat/print-ready too.
 *
 * This table is the single source of truth for "is this asset an editable
 * template or a graphic?" — nothing classifies stationery by name anywhere else.
 */
const STATIONERY_GEN_CONFIG: Record<string, StationeryGenConfig> = {
  // Core Stationery (official print sizes)
  'business-card': { label: 'Business Card', promptHint: 'professional business card design with logo placement and contact details', style: 'Minimalist', platform: 'Print', aspectRatio: '16:9', assetTypeCategory: 'core-stationery', outputFormat: 'html' },
  'letterhead': { label: 'Letterhead', promptHint: 'corporate letterhead with branded header, logo and footer', style: 'Minimalist', platform: 'Print', aspectRatio: '3:4', assetTypeCategory: 'core-stationery', outputFormat: 'html' },
  'envelope-a4': { label: 'Envelope (A4)', promptHint: 'branded A4 envelope with logo and return address', style: 'Minimalist', platform: 'Print', aspectRatio: '3:4', assetTypeCategory: 'core-stationery', outputFormat: 'html' },
  'envelope-dl': { label: 'Envelope (DL)', promptHint: 'branded DL envelope with logo and address layout', style: 'Minimalist', platform: 'Print', aspectRatio: '16:9', assetTypeCategory: 'core-stationery', outputFormat: 'html' },
  'email-signature': { label: 'Email Signature', promptHint: 'professional email signature with logo, name, title and contact info', style: 'Minimalist', platform: 'Email', aspectRatio: '16:9', assetTypeCategory: 'core-stationery', outputFormat: 'html' },
  'presentation-template': { label: 'Presentation Template', promptHint: 'branded presentation slide template with title and content layouts', style: 'Minimalist', platform: 'Presentation', aspectRatio: '16:9', assetTypeCategory: 'core-stationery', outputFormat: 'png' },
  // Document templates (A4)
  'invoice-template': { label: 'Invoice Template', promptHint: 'professional invoice with branded header, item table and totals', style: 'Minimalist', platform: 'Print', aspectRatio: '3:4', assetTypeCategory: 'office-assets', outputFormat: 'html' },
  'quotation-template': { label: 'Quotation Template', promptHint: 'professional quotation with company branding and itemised layout', style: 'Minimalist', platform: 'Print', aspectRatio: '3:4', assetTypeCategory: 'office-assets', outputFormat: 'html' },
  'receipt-design': { label: 'Receipt Design', promptHint: 'branded receipt with logo and transaction layout', style: 'Minimalist', platform: 'Print', aspectRatio: '3:4', assetTypeCategory: 'office-assets', outputFormat: 'html' },
  'purchase-order': { label: 'Purchase Order Template', promptHint: 'professional purchase order with branded header and order table', style: 'Minimalist', platform: 'Print', aspectRatio: '3:4', assetTypeCategory: 'office-assets', outputFormat: 'html' },
  'billing-format': { label: 'Billing Format', promptHint: 'branded billing template with company header and payment details', style: 'Minimalist', platform: 'Print', aspectRatio: '3:4', assetTypeCategory: 'office-assets', outputFormat: 'html' },
  'proposal-template': { label: 'Proposal Template', promptHint: 'professional proposal with cover page, sections and branded footer', style: 'Realistic', platform: 'Print', aspectRatio: '3:4', assetTypeCategory: 'office-assets', outputFormat: 'html' },
  'thank-you-card': { label: 'Thank You Card', promptHint: 'elegant branded thank you card with logo and message area', style: 'Realistic', platform: 'Print', aspectRatio: '4:3', assetTypeCategory: 'packaging-stationery', outputFormat: 'html' },
  'warranty-card': { label: 'Warranty Card', promptHint: 'professional warranty card with branding and terms layout', style: 'Minimalist', platform: 'Print', aspectRatio: '4:3', assetTypeCategory: 'packaging-stationery', outputFormat: 'html' },
  'instruction-manual': { label: 'Instruction Manual', promptHint: 'branded instruction manual cover with logo and product imagery area', style: 'Minimalist', platform: 'Print', aspectRatio: '3:4', assetTypeCategory: 'packaging-stationery', outputFormat: 'html' },
  'product-insert-card': { label: 'Product Insert Card', promptHint: 'branded product insert card for packaging inclusion', style: 'Realistic', platform: 'Print', aspectRatio: '4:3', assetTypeCategory: 'packaging-stationery', outputFormat: 'html' },
  // Branding & merchandise (flat print-ready, no 3D mockups)
  'branded-stickers': { label: 'Branded Stickers', promptHint: 'branded sticker sheet with logo variations and decorative elements', style: 'Realistic', platform: 'Print', aspectRatio: '1:1', assetTypeCategory: 'print-marketing', outputFormat: 'png' },
  'packaging-tape': { label: 'Packaging Tape', promptHint: 'branded packaging tape with repeating logo pattern and brand colours', style: 'Minimalist', platform: 'Print', aspectRatio: '16:9', assetTypeCategory: 'print-marketing', outputFormat: 'png' },
  'stamps': { label: 'Stamps', promptHint: 'branded rubber stamp with logo and company address', style: 'Minimalist', platform: 'Print', aspectRatio: '1:1', assetTypeCategory: 'print-marketing', outputFormat: 'png' },
  'branding-print': { label: 'Branding (Print)', promptHint: 'brand identity print sheet showcasing logo, colours and typography', style: 'Realistic', platform: 'Print', aspectRatio: '4:3', assetTypeCategory: 'print-marketing', outputFormat: 'png' },
  'standees-print': { label: 'Standees', promptHint: 'branded standee with messaging, logo and visual areas', style: 'Realistic', platform: 'Print', aspectRatio: '9:16', assetTypeCategory: 'print-marketing', outputFormat: 'png' },
  'booth-designs': { label: 'Booth Designs', promptHint: 'branded exhibition booth panel with logo and product showcase areas', style: 'Realistic', platform: 'Print', aspectRatio: '16:9', assetTypeCategory: 'print-marketing', outputFormat: 'png' },
  't-shirts': { label: 'T-shirts', promptHint: 'branded t-shirt design with logo, tagline area and visual elements', style: 'Realistic', platform: 'Print', aspectRatio: '3:4', assetTypeCategory: 'print-marketing', outputFormat: 'png' },
  'notebook': { label: 'Corporate Notebook / Diary', promptHint: 'branded notebook cover with logo and elegant brand pattern', style: 'Minimalist', platform: 'Print', aspectRatio: '3:4', assetTypeCategory: 'print-marketing', outputFormat: 'png' },
  'coffee-mug': { label: 'Coffee Mug / Tumbler', promptHint: 'branded mug wrap design with logo and tagline area', style: 'Realistic', platform: 'Print', aspectRatio: '16:9', assetTypeCategory: 'print-marketing', outputFormat: 'png' },
  'tote-bag': { label: 'Tote Bag', promptHint: 'branded tote bag design with logo, tagline and visual elements', style: 'Realistic', platform: 'Print', aspectRatio: '3:4', assetTypeCategory: 'print-marketing', outputFormat: 'png' },
  'newsletter-template': { label: 'Newsletter Template', promptHint: 'branded newsletter with header, article sections and footer', style: 'Minimalist', platform: 'Print', aspectRatio: '3:4', assetTypeCategory: 'print-marketing', outputFormat: 'html' },
};

/**
 * Resolve the output format for a stationery type.
 *
 * STATIONERY_GEN_CONFIG is the source of truth. Legacy types that predate the
 * config table ('memo-pad', 'compliment-slip', 'folder', 'other', …) fall back
 * to the existing full-page-document classification, so they still land on the
 * correct side of the split without a second hand-maintained list.
 */
function resolveStationeryOutputFormat(stationeryType: string): StationeryOutputFormat {
  const cfg = STATIONERY_GEN_CONFIG[stationeryType];
  if (cfg) return cfg.outputFormat;
  return isFullPageDocumentCategory(stationeryType) ? 'html' : 'png';
}

/**
 * Standard print size per HTML stationery type, so the generated template is
 * print-correct rather than an arbitrary web page. `css` drives the `@page`
 * rule of the offline fallback below; `label` is what the model is told.
 * Types absent from the map are laid out on A4.
 */
const STATIONERY_HTML_PAGE: Record<string, { css: string; label: string }> = {
  'business-card': { css: '90mm 54mm', label: '90mm wide x 54mm tall (standard business card)' },
  'letterhead': { css: '210mm 297mm', label: '210mm wide x 297mm tall (A4 portrait)' },
  'envelope-a4': { css: '324mm 229mm', label: '324mm wide x 229mm tall (C4 envelope, landscape)' },
  'envelope-dl': { css: '220mm 110mm', label: '220mm wide x 110mm tall (DL envelope, landscape)' },
  'email-signature': { css: '160mm 60mm', label: '600px wide, height driven by content (email signature block)' },
  'invoice-template': { css: '210mm 297mm', label: '210mm wide x 297mm tall (A4 portrait)' },
  'quotation-template': { css: '210mm 297mm', label: '210mm wide x 297mm tall (A4 portrait)' },
  'receipt-design': { css: '80mm 200mm', label: '80mm wide x 200mm tall (thermal receipt)' },
  'purchase-order': { css: '210mm 297mm', label: '210mm wide x 297mm tall (A4 portrait)' },
  'billing-format': { css: '210mm 297mm', label: '210mm wide x 297mm tall (A4 portrait)' },
  'proposal-template': { css: '210mm 297mm', label: '210mm wide x 297mm tall (A4 portrait)' },
  'thank-you-card': { css: '148mm 105mm', label: '148mm wide x 105mm tall (A6 landscape)' },
  'warranty-card': { css: '148mm 105mm', label: '148mm wide x 105mm tall (A6 landscape)' },
  'instruction-manual': { css: '148mm 210mm', label: '148mm wide x 210mm tall (A5 portrait)' },
  'product-insert-card': { css: '105mm 148mm', label: '105mm wide x 148mm tall (A6 portrait)' },
  'newsletter-template': { css: '210mm 297mm', label: '210mm wide x 297mm tall (A4 portrait)' },
};

const STATIONERY_HTML_PAGE_DEFAULT = { css: '210mm 297mm', label: '210mm wide x 297mm tall (A4 portrait)' };

/** Token the AI is told to use for the logo `src`; replaced with the real logo afterwards. */
const STATIONERY_LOGO_TOKEN = '{{LOGO_URL}}';

/**
 * Prompt selection resolved from the Super Admin prompt configuration.
 *
 * `overrides` feeds the existing image prompt builder (PNG path); `guidanceText`
 * / `contentText` are the same prompts as flat text for the HTML path. Both come
 * from the same PromptConfig documents, so the two formats can never drift onto
 * different prompt sources.
 */
interface StationeryPromptSelection {
  overrides: Record<string, any>;
  guidanceText: string;
  contentText: string;
}

/**
 * Load the admin-editable prompts for a stationery type and apply the user's
 * choice from the prompt step: a selected Super Admin variant (`promptConfigId`)
 * and/or a prompt the user typed themselves (`customPrompt`, which wins).
 */
async function resolveStationeryPrompts(
  stationeryType: string,
  promptConfigId?: string,
  customPrompt?: string,
): Promise<StationeryPromptSelection> {
  const [styleGuidance, platformGuidance, assetCategoryGuidance, assetContentElements] = await Promise.all([
    getStyleGuidance().catch(() => null),
    getPlatformGuidance().catch(() => null),
    getAssetCategoryGuidance().catch(() => null),
    getAssetContentElements().catch(() => null),
  ]);

  const overrides: Record<string, any> = {
    ...(styleGuidance && { styleGuidance }),
    ...(platformGuidance && { platformGuidance }),
    ...(assetCategoryGuidance && { assetCategoryGuidance }),
    ...(assetContentElements && { assetContentElements }),
  };

  let guidanceText = assetCategoryGuidance?.[stationeryType] || '';
  let contentText = assetContentElements?.[stationeryType] || '';

  if (promptConfigId) {
    try {
      const { getPromptVariantById } = await import('../services/aiContext/promptConfigLoader');
      const selected = await getPromptVariantById(promptConfigId);
      if (selected) {
        if (selected.type === 'asset_content_elements') {
          overrides.assetContentElements = { ...(assetContentElements || {}), [selected.key]: selected.prompt };
          contentText = selected.prompt;
        } else {
          overrides.assetCategoryGuidance = { ...(assetCategoryGuidance || {}), [selected.key]: selected.prompt };
          guidanceText = selected.prompt;
        }
      }
    } catch (err: any) {
      console.warn('[Stationery] Failed to load selected prompt variant, using defaults:', err?.message);
    }
  }

  const typed = (customPrompt || '').trim();
  if (typed) {
    overrides.assetCategoryGuidance = {
      ...(overrides.assetCategoryGuidance || assetCategoryGuidance || {}),
      [stationeryType]: typed,
    };
    guidanceText = typed;
  }

  return { overrides, guidanceText, contentText };
}

/**
 * Output contract for HTML stationery. The template has to stay EDITABLE after
 * generation, so the model is told to emit plain HTML with the copy in ordinary
 * elements — no images of text, no scripts, no markdown wrapper.
 */
const STATIONERY_HTML_SYSTEM_PROMPT = `You are a senior brand designer who writes production-quality, print-ready HTML/CSS stationery templates.

OUTPUT CONTRACT — follow every rule exactly:
- Return ONE complete standalone HTML5 document: it must start with <!DOCTYPE html> and end with </html>.
- Return HTML CODE ONLY. No markdown code fences, no backticks, no commentary, no explanation, no notes before or after the document.
- Put ALL styling in a single <style> block inside <head>. No external stylesheets, no webfont links, no external images.
- Use NO JavaScript: no <script> tags, no inline on* event handlers.
- Every piece of copy (company name, address, phone, email, website, labels, table headings) must sit in its own plain HTML element with real text, so a person can open the file and edit the wording directly. Never render text as an image or as CSS content.
- Include an @page rule and an @media print block so the document prints at the specified size with no extra margins.
- Use ONLY the brand colours and typography supplied. Never invent a different palette. Use web-safe font stacks.
- Where the brand logo belongs, output exactly <img src="${STATIONERY_LOGO_TOKEN}" alt="Company logo"> and size it with CSS. Never draw, letter or describe a logo any other way, and never use a placeholder image service.`;

/** Brand context condensed into the few lines the HTML prompt actually needs. */
function buildStationeryBrandBrief(brandContext: any): string {
  const lines: string[] = [];
  if (brandContext?.brandName) lines.push(`Brand name: ${brandContext.brandName}`);
  if (Array.isArray(brandContext?.brandColors) && brandContext.brandColors.length) {
    lines.push(`Brand colours (use these exactly): ${brandContext.brandColors.slice(0, 6).join(', ')}`);
  }
  const typography = brandContext?.visualIdentity?.typography;
  if (typography) {
    lines.push(`Typography: ${typeof typography === 'string' ? typography : JSON.stringify(typography).slice(0, 400)}`);
  }
  if (brandContext?.businessDescription) lines.push(`Business: ${String(brandContext.businessDescription).slice(0, 400)}`);
  if (brandContext?.businessIndustry) lines.push(`Industry: ${brandContext.businessIndustry}`);
  if (brandContext?.brandTone) lines.push(`Tone of voice: ${brandContext.brandTone}`);
  if (brandContext?.visualIdentity?.visualStyle) lines.push(`Visual style: ${brandContext.visualIdentity.visualStyle}`);
  return lines.join('\n');
}

/**
 * User prompt for an HTML stationery template.
 *
 * When `existingHtml` is supplied this becomes a REGENERATION: the current
 * template is handed back to the model and it is told to edit that document in
 * place, so a user asking to "change the address and add a website" gets their
 * own template updated rather than a brand new design.
 */
function buildStationeryHtmlUserPrompt(params: {
  cfg: StationeryGenConfig;
  stationeryType: string;
  assetName: string;
  description?: string;
  brandContext: any;
  guidanceText: string;
  contentText: string;
  existingHtml?: string;
  feedback?: string;
}): string {
  const { cfg, stationeryType, assetName, description, brandContext, guidanceText, contentText, existingHtml, feedback } = params;
  const pageSize = (STATIONERY_HTML_PAGE[stationeryType] || STATIONERY_HTML_PAGE_DEFAULT).label;
  const parts: string[] = [];

  if (existingHtml) {
    parts.push(
      `Below is the EXISTING HTML template for the ${cfg.label} "${assetName}". Update it according to the change request.`,
      '',
      'RULES FOR THIS UPDATE:',
      '- Start from the existing document. Keep its layout, structure, CSS, typography, spacing and colours exactly as they are.',
      '- Change ONLY what the request asks for. If the request adds a field, insert it in the place that matches the existing style.',
      '- Do not redesign, do not restyle, and do not drop sections that were not mentioned.',
      '- Return the COMPLETE updated document, not a diff and not a fragment.',
      '',
      'EXISTING TEMPLATE:',
      existingHtml,
      '',
      `CHANGE REQUEST: ${feedback || description || 'Refresh the content.'}`,
    );
  } else {
    parts.push(
      `Create an editable ${cfg.label} template: ${cfg.promptHint}.`,
      `Asset name: ${assetName}`,
      `Page size: ${pageSize}`,
    );
    if (description) parts.push(`User requirements: ${description}`);
    if (feedback) parts.push(`Additional instructions: ${feedback}`);
  }

  const brief = buildStationeryBrandBrief(brandContext);
  if (brief) parts.push('', 'BRAND CONTEXT:', brief);
  if (guidanceText) parts.push('', 'DESIGN GUIDANCE:', guidanceText);
  if (contentText) parts.push('', 'CONTENT ELEMENTS TO INCLUDE:', contentText);

  parts.push('', 'Respond with the HTML document only.');
  return parts.join('\n');
}

/**
 * Turn a model response into a clean, standalone HTML document.
 *
 * Strips reasoning traces and markdown fences, drops any prose the model put
 * around the document, removes scripts/handlers, and wraps a bare fragment in a
 * minimal document so the file always opens on its own.
 */
function extractCleanStationeryHtml(raw: string, title: string): string {
  let html = stripReasoning(String(raw || '')).trim();

  // Markdown fences — the contract forbids them, but models still add them.
  const fenced = html.match(/```(?:html)?\s*\n?([\s\S]*?)```/i);
  if (fenced) html = fenced[1].trim();
  html = html.replace(/^```(?:html)?\s*\n?/i, '').replace(/\n?```\s*$/, '').trim();

  // Drop any commentary before/after the document itself.
  const docStart = html.search(/<!DOCTYPE html/i);
  if (docStart > 0) html = html.slice(docStart);
  const closeIdx = html.toLowerCase().lastIndexOf('</html>');
  if (closeIdx !== -1) html = html.slice(0, closeIdx + '</html>'.length);

  // No scripts, no inline handlers, no embedded frames — the same rules the HR
  // asset template generator applies to raw AI HTML.
  html = html
    .replace(/<script\b[^<]*(?:(?!<\/script>)<[^<]*)*<\/script>/gi, '')
    .replace(/\son\w+\s*=\s*["'][^"']*["']/gi, '')
    .replace(/<iframe\b[^>]*>[\s\S]*?<\/iframe>/gi, '')
    .replace(/<object\b[^>]*>[\s\S]*?<\/object>/gi, '')
    .replace(/<embed\b[^>]*>/gi, '')
    .trim();

  if (!html) return '';

  // A fragment still has to open on its own.
  if (!/<html[\s>]/i.test(html)) {
    html = [
      '<!DOCTYPE html>',
      '<html lang="en">',
      '<head>',
      '<meta charset="UTF-8">',
      '<meta name="viewport" content="width=device-width, initial-scale=1.0">',
      `<title>${title.replace(/[<>&]/g, '')}</title>`,
      '</head>',
      '<body>',
      html,
      '</body>',
      '</html>',
    ].join('\n');
  }
  return html;
}

/** Escape a value for safe interpolation into the fallback template. */
function escapeStationeryHtml(value: string): string {
  return String(value ?? '')
    .replace(/&/g, '&amp;')
    .replace(/</g, '&lt;')
    .replace(/>/g, '&gt;')
    .replace(/"/g, '&quot;');
}

/**
 * Deterministic, brand-styled HTML template used when no AI provider can be
 * reached.
 *
 * The HR asset template generator does the same thing: a generation that cannot
 * reach a model still returns a usable, editable document rather than failing
 * the whole job. The copy is deliberately placeholder text in ordinary elements
 * so the user can edit it, and the page size and brand colours are the real
 * ones for this asset type.
 */
function buildFallbackStationeryHtml(params: {
  cfg: StationeryGenConfig;
  stationeryType: string;
  displayName: string;
  brandContext: any;
  description?: string;
}): string {
  const { cfg, stationeryType, displayName, brandContext, description } = params;
  const page = STATIONERY_HTML_PAGE[stationeryType] || STATIONERY_HTML_PAGE_DEFAULT;
  const colors: string[] = Array.isArray(brandContext?.brandColors) ? brandContext.brandColors.filter(Boolean) : [];
  const primary = escapeStationeryHtml(colors[0] || '#1F2937');
  const accent = escapeStationeryHtml(colors[1] || colors[0] || '#4B5563');
  const brandName = escapeStationeryHtml(brandContext?.brandName || 'Company Name');
  const title = escapeStationeryHtml(displayName);
  const label = escapeStationeryHtml(cfg.label);
  const note = description ? `<p class="note">${escapeStationeryHtml(description)}</p>` : '';

  return [
    '<!DOCTYPE html>',
    '<html lang="en">',
    '<head>',
    '<meta charset="UTF-8">',
    '<meta name="viewport" content="width=device-width, initial-scale=1.0">',
    `<title>${title}</title>`,
    '<style>',
    `@page { size: ${page.css}; margin: 0; }`,
    '* { box-sizing: border-box; }',
    'body { margin: 0; font-family: "Helvetica Neue", Arial, sans-serif; color: #111827; background: #f3f4f6; }',
    `.sheet { width: 100%; max-width: 100%; min-height: 100vh; padding: 8mm; background: #ffffff; border-top: 4mm solid ${primary}; display: flex; flex-direction: column; gap: 6mm; }`,
    '.header { display: flex; align-items: center; gap: 4mm; }',
    '.logo { max-height: 16mm; max-width: 45mm; object-fit: contain; }',
    `.brand { font-size: 14pt; font-weight: 700; letter-spacing: .02em; color: ${primary}; }`,
    `.kind { font-size: 8pt; text-transform: uppercase; letter-spacing: .12em; color: ${accent}; }`,
    '.body { flex: 1; }',
    '.body p { margin: 0 0 2mm; font-size: 10pt; line-height: 1.5; }',
    `.note { font-size: 9pt; color: ${accent}; }`,
    `.footer { border-top: 0.4mm solid ${accent}; padding-top: 3mm; font-size: 9pt; line-height: 1.6; color: #374151; }`,
    '@media print { body { background: #ffffff; } .sheet { min-height: auto; } }',
    '</style>',
    '</head>',
    '<body>',
    '<div class="sheet">',
    '<header class="header">',
    `<img class="logo" src="${STATIONERY_LOGO_TOKEN}" alt="Company logo">`,
    '<div>',
    `<div class="brand">${brandName}</div>`,
    `<div class="kind">${label}</div>`,
    '</div>',
    '</header>',
    '<main class="body">',
    `<p>${title}</p>`,
    '<p>Edit this text to add your content.</p>',
    note,
    '</main>',
    '<footer class="footer">',
    '<p>Address line 1, Address line 2</p>',
    '<p>Phone: +00 0000 000000</p>',
    '<p>Email: hello@example.com</p>',
    '</footer>',
    '</div>',
    '</body>',
    '</html>',
  ].join('\n');
}

/**
 * Generate (or update) an HTML stationery template.
 *
 * Regeneration path: when the linked Stationery item already holds HTML, that
 * HTML is passed back to the model with the user's change request so the
 * existing template is edited in place instead of being replaced by a new
 * design. The item is then updated in place, exactly as the image path does.
 */
async function generateOneStationeryTemplate(
  params: StationeryGenParams,
  onProgress?: (pct: number, step: string) => void,
): Promise<StationeryItemResult> {
  const { companyId, stationeryType, assetName, description, feedback, promptConfigId, customPrompt, stationeryId: existingStationeryId } = params;
  const cfg: StationeryGenConfig = STATIONERY_GEN_CONFIG[stationeryType] || {
    label: stationeryType, promptHint: stationeryType, style: 'Minimalist', platform: 'Print',
    aspectRatio: '3:4', assetTypeCategory: 'core-stationery', outputFormat: 'html',
  };
  const { Stationery } = getModels();
  const displayName = `${assetName} - ${cfg.label}`;

  onProgress?.(10, 'Gathering brand assets…');
  const brandContext = await fetchBrandContext(companyId, {});

  // Existing template (regeneration) — retrieved so the model edits it rather
  // than inventing a replacement.
  let existingItem: any = null;
  if (existingStationeryId) {
    existingItem = await (Stationery as any).findOne({ _id: existingStationeryId, companyId });
  }
  let existingHtml = '';
  if (existingItem) {
    existingHtml = existingItem.populatedHtml || '';
    if (!existingHtml && existingItem.base64Data && existingItem.fileType === 'text/html') {
      try { existingHtml = Buffer.from(String(existingItem.base64Data), 'base64').toString('utf-8'); } catch { existingHtml = ''; }
    }
  }

  onProgress?.(35, existingHtml ? `Updating your ${cfg.label.toLowerCase()}…` : `Designing your ${cfg.label.toLowerCase()}…`);
  const { guidanceText, contentText } = await resolveStationeryPrompts(stationeryType, promptConfigId, customPrompt);
  const userPrompt = buildStationeryHtmlUserPrompt({
    cfg, stationeryType, assetName: displayName, description, brandContext,
    guidanceText, contentText,
    existingHtml: existingHtml || undefined,
    feedback,
  });

  onProgress?.(55, 'Writing the template…');
  // generateWithAI already walks the whole provider chain (Ollama → Zhipu →
  // Claude → OpenAI), so one call is the whole attempt — a second call would
  // only repeat the same failures and double the wait.
  let html = '';
  try {
    const result = await generateWithAI(userPrompt, STATIONERY_HTML_SYSTEM_PROMPT, 12000, 0.5, 'text', 'ollama');
    html = extractCleanStationeryHtml(result.content, displayName);
    if (!html) console.warn(`[Stationery] "${stationeryType}": the model returned no usable HTML (${result.content?.length || 0} chars).`);
  } catch (aiErr: any) {
    console.warn(`[Stationery] "${stationeryType}": HTML generation failed on every provider — ${aiErr?.message}`);
  }
  if (!html) {
    // Never fail the item and never overwrite a good template with nothing:
    // keep what exists, otherwise hand back a usable branded starting point the
    // user can edit and regenerate — the HR asset generator behaves the same way.
    html = existingHtml || buildFallbackStationeryHtml({ cfg, stationeryType, displayName, brandContext, description });
    console.warn(`[Stationery] "${stationeryType}": using the ${existingHtml ? 'existing' : 'offline fallback'} template.`);
  }

  // Inline the real brand logo so the file stands alone.
  onProgress?.(85, 'Adding your brand logo…');
  const logoDataUri = await resolveBrandLogoDataUri(companyId);
  html = logoDataUri
    ? html.split(STATIONERY_LOGO_TOKEN).join(logoDataUri)
    // No usable logo — drop the placeholder images rather than leaving broken ones.
    : html.replace(/<img[^>]*src=["']\{\{LOGO_URL\}\}["'][^>]*>/gi, '');

  onProgress?.(95, 'Saving your template…');
  const base64Data = Buffer.from(html, 'utf-8').toString('base64');
  const fileName = `${displayName.replace(/[^a-zA-Z0-9_-]/g, '_')}.html`;

  let savedId: string;
  if (existingItem) {
    existingItem.populatedHtml = html;
    existingItem.base64Data = base64Data;
    existingItem.fileName = fileName;
    existingItem.fileSize = Buffer.byteLength(html, 'utf-8');
    existingItem.fileType = 'text/html';
    existingItem.updatedAt = new Date();
    await existingItem.save();
    savedId = String(existingItem._id);
  } else {
    const item = new (Stationery as any)({
      companyId,
      name: displayName,
      type: stationeryType,
      description: `AI-generated editable ${cfg.label} template.`,
      category: STATIONERY_TYPE_TO_CATEGORY[stationeryType] || cfg.assetTypeCategory || 'other',
      templateUrl: `/stationery/template/${stationeryType}`,
      sourceUrl: 'ai-generation',
      base64Data,
      populatedHtml: html,
      fileName,
      fileSize: Buffer.byteLength(html, 'utf-8'),
      fileType: 'text/html',
      kind: 'ai',
      exportFormats: ['html', 'pdf'],
      status: 'draft',
      tags: ['ai-generated', stationeryType, 'html-template'],
    });
    await item.save();
    savedId = String(item._id);
  }

  return {
    imageGenerationId: '',
    stationeryType,
    stationeryLabel: cfg.label,
    assetName,
    stationeryId: savedId,
    outputFormat: 'html',
  };
}

/**
 * Hard instruction appended to every stationery image prompt: the AI must NOT draw
 * any logo (it kept inventing one, e.g. a "U" mark). The real logo is composited on
 * top afterwards, in the reserved top-left area.
 */
const NO_LOGO_DIRECTIVE = `

ABSOLUTELY CRITICAL — NO LOGO GENERATION: Do NOT draw, generate, invent, recreate, simplify, or include ANY logo, brand mark, monogram, lettermark, company initial, emblem, badge, symbol, or icon anywhere in this design. The brand's real logo will be added separately as an image afterwards. Reserve the TOP-LEFT area as clean, EMPTY space (plain background only) — keep it completely free of text, letters, graphics, shapes, or any design element so the real logo can be placed there without overlapping anything. Do NOT put a placeholder logo or the company's initial in that space.`;

interface StationeryGenParams {
  companyId: string;
  stationeryType: string;
  assetName: string;
  description?: string;
  userId: string;
  feedback?: string;            // regeneration instructions
  existingGenerationId?: string; // reuse the record and append a new version
  /** Super Admin prompt variant chosen on the prompt step (PromptConfig _id). */
  promptConfigId?: string;
  /** Prompt the user typed/edited on the prompt step — wins over the variant. */
  customPrompt?: string;
  /** Existing HTML stationery item to update in place (HTML regeneration). */
  stationeryId?: string;
}

/**
 * Background orchestration for a single stationery design: gather brand context →
 * enhance the prompt → generate the image at the standard size → overlay the real
 * brand logo → persist as an ImageGeneration version. Reports progress to the job.
 */
type StationeryItemResult = { imageGenerationId: string; stationeryType: string; stationeryLabel: string; assetName: string; stationeryId?: string; outputFormat: StationeryOutputFormat };

/** Category mapping mirroring the Stationery module's grouping (used for the saved record). */
const STATIONERY_TYPE_TO_CATEGORY: Record<string, string> = {
  'business-card': 'core-stationery', 'letterhead': 'core-stationery', 'envelope-a4': 'core-stationery',
  'envelope-dl': 'core-stationery', 'email-signature': 'core-stationery', 'presentation-template': 'core-stationery',
  'invoice-template': 'office-assets', 'quotation-template': 'office-assets', 'receipt-design': 'office-assets',
  'purchase-order': 'office-assets', 'billing-format': 'office-assets', 'proposal-template': 'office-assets',
  'thank-you-card': 'packaging-stationery', 'warranty-card': 'packaging-stationery', 'instruction-manual': 'packaging-stationery',
  'product-insert-card': 'packaging-stationery', 'branded-stickers': 'packaging-stationery', 'packaging-tape': 'packaging-stationery',
  'stamps': 'print-stationery', 'branding-print': 'print-stationery', 'standees-print': 'print-stationery',
  'booth-designs': 'print-stationery', 't-shirts': 'print-stationery', 'notebook': 'print-stationery',
  'coffee-mug': 'print-stationery', 'tote-bag': 'print-stationery', 'newsletter-template': 'marketing-assets',
};

async function generateOneStationeryDesign(params: StationeryGenParams, onProgress?: (pct: number, step: string) => void): Promise<StationeryItemResult> {
  // Editable-template assets are produced as HTML, graphic assets as PNG. The
  // asset type decides — see STATIONERY_GEN_CONFIG.outputFormat.
  if (resolveStationeryOutputFormat(params.stationeryType) === 'html') {
    return generateOneStationeryTemplate(params, onProgress);
  }

  const { companyId, stationeryType, assetName, description, userId, feedback, existingGenerationId, promptConfigId, customPrompt } = params;
  const cfg: StationeryGenConfig = STATIONERY_GEN_CONFIG[stationeryType] || { label: stationeryType, promptHint: stationeryType, style: 'Minimalist', platform: 'Print', aspectRatio: '3:4', assetTypeCategory: 'core-stationery', outputFormat: 'png' };
  const { ImageGeneration } = getModels();

  onProgress?.(10, 'Gathering brand assets…');
  const brandContext = await fetchBrandContext(companyId, {});
  const referenceAssetId = await resolvePrimaryLogoAssetId(companyId);

  onProgress?.(30, `Designing your ${cfg.label.toLowerCase()}…`);
  const baseDescription = `Create ${cfg.promptHint} for "${assetName}".${description ? ' ' + description : ''}`;
  const inputs: ImageEnhancementInputs = {
    description: baseDescription,
    objective: `Professional, print-ready ${cfg.label} design with brand consistency`,
    style: cfg.style,
    platform: cfg.platform,
    aspectRatio: cfg.aspectRatio,
    assetTypeCategory: cfg.assetTypeCategory,
    assetCategory: stationeryType,
    assetRequirements: description,
    userInstructions: feedback,
    ...brandContext,
  };

  // Admin-editable prompts, with the variant the user picked (or the prompt they
  // typed) on the prompt step applied.
  const { overrides: promptOverrides } = await resolveStationeryPrompts(stationeryType, promptConfigId, customPrompt);

  const { systemPrompt, userPrompt, maxTokens } = buildEnhancementPrompts(inputs, promptOverrides);
  let enhancedPrompt = '';
  try {
    const result = await generateWithAI(userPrompt, systemPrompt, maxTokens, 0.7, 'text', 'ollama');
    enhancedPrompt = stripReasoning(result.content.trim());
  } catch {
    try {
      const result = await generateWithAI(userPrompt, systemPrompt, maxTokens, 0.7, 'text');
      enhancedPrompt = stripReasoning(result.content.trim());
    } catch {
      enhancedPrompt = buildFallbackPrompt(inputs, promptOverrides);
    }
  }
  if (feedback) {
    enhancedPrompt = `${enhancedPrompt}\n\nREVISION REQUESTED BY USER — apply this change: ${feedback}`;
  }
  // Never let the AI invent a logo — the real one is overlaid afterwards.
  enhancedPrompt = `${enhancedPrompt}${NO_LOGO_DIRECTIVE}`;

  onProgress?.(55, 'Generating the design…');
  const size = ASPECT_RATIO_SIZE_MAP[cfg.aspectRatio] || '1024x1536';

  // Create or reuse the ImageGeneration record (regenerate appends a new version)
  let generation: any = null;
  if (existingGenerationId) {
    generation = await (ImageGeneration as any).findOne({ _id: existingGenerationId, companyId });
  }
  if (!generation) {
    generation = new (ImageGeneration as any)({
      companyId,
      name: `${assetName} - ${cfg.label}`,
      description: baseDescription,
      objective: inputs.objective,
      style: cfg.style,
      platform: cfg.platform,
      aspectRatio: cfg.aspectRatio,
      originalPrompt: baseDescription,
      status: 'generating',
      tags: [stationeryType, 'stationery'],
      createdBy: userId,
    });
    await generation.save();
  } else {
    generation.status = 'generating';
    await generation.save();
  }

  // Generate the image (OpenAI, then Zhipu CogView fallback), then overlay the logo.
  let imageResult: { base64Data: string; revisedPrompt: string; model: string; provider: string; tokenUsage?: any; latencyMs?: number };
  try {
    imageResult = await generateImageWithOpenAI(enhancedPrompt, size, 'hd', 'vivid', 'gpt-image-1');
  } catch (openaiErr: any) {
    console.warn(`[Stationery] OpenAI generation failed (${openaiErr?.message}); trying Zhipu CogView fallback`);
    imageResult = await generateImageWithZhipuCogView(enhancedPrompt, size);
  }
  imageResult.base64Data = await maybeOverlayBrandLogo(imageResult.base64Data, referenceAssetId, companyId);

  onProgress?.(90, 'Finalising your design…');
  for (const v of generation.versions) v.isCurrent = false;
  const versionNumber = generation.versions.length + 1;
  generation.versions.push({
    version: versionNumber,
    prompt: generation.originalPrompt,
    enhancedPrompt: stripReasoning(enhancedPrompt),
    imageUrl: '',
    base64Data: imageResult.base64Data,
    generationProvider: imageResult.provider,
    generationModel: imageResult.model,
    aspectRatio: cfg.aspectRatio,
    size,
    quality: 'hd',
    style: 'vivid',
    revisedPrompt: imageResult.revisedPrompt,
    tokenUsage: (imageResult as any).tokenUsage || { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
    latencyMs: (imageResult as any).latencyMs || 0,
    isCurrent: true,
  } as any);
  generation.currentVersion = versionNumber;
  generation.totalVersions = versionNumber;
  generation.generationProvider = imageResult.provider;
  generation.generationModel = imageResult.model;
  generation.status = 'completed';
  await generation.save();

  // Auto-save the design into the Stationery library so it appears in the Stationery
  // tab immediately (as a draft). Regeneration updates the SAME item in place.
  let stationeryId: string | undefined;
  try {
    const { Stationery } = getModels();
    const base64WithPrefix = `data:image/png;base64,${imageResult.base64Data}`;
    const { buffer, mimeType } = base64ToBuffer(base64WithPrefix, 'image/png');
    const ext = getExtensionFromMime(mimeType, 'png');
    const { url: assetUrl, fileSize } = await saveStationeryFile(buffer, `ai-generated.${ext}`, mimeType);

    const linkedId = (generation as any).stationeryAssetId;
    const existing = linkedId ? await (Stationery as any).findOne({ _id: linkedId, companyId }) : null;
    if (existing) {
      // Regenerated → replace the image on the existing stationery item
      existing.templateUrl = assetUrl;
      existing.previewImageUrl = assetUrl;
      existing.base64Data = imageResult.base64Data;
      existing.fileName = `stationery-${Date.now()}.${ext}`;
      existing.fileSize = fileSize;
      existing.fileType = mimeType;
      existing.updatedAt = new Date();
      await existing.save();
      stationeryId = String(existing._id);
    } else {
      const stationeryItem = new (Stationery as any)({
        companyId,
        name: `${assetName} - ${cfg.label}`,
        type: stationeryType,
        description: `AI-generated ${cfg.label}.`,
        category: STATIONERY_TYPE_TO_CATEGORY[stationeryType] || cfg.assetTypeCategory || 'other',
        templateUrl: assetUrl,
        previewImageUrl: assetUrl,
        sourceUrl: 'ai-generation',
        base64Data: imageResult.base64Data,
        fileName: `stationery-${Date.now()}.${ext}`,
        fileSize,
        fileType: mimeType,
        status: 'draft',
        tags: ['ai-generated', stationeryType],
      });
      await stationeryItem.save();
      stationeryId = String(stationeryItem._id);
      (generation as any).stationeryAssetId = stationeryId;
      await generation.save();
    }
  } catch (saveErr: any) {
    console.warn(`[Stationery] Auto-save to Stationery failed: ${saveErr?.message}`);
  }

  return { imageGenerationId: String(generation._id), stationeryType, stationeryLabel: cfg.label, assetName, stationeryId, outputFormat: 'png' };
}

/**
 * Process a batch of selected stationery items ONE AT A TIME (sequentially). Each
 * item produces a single design; once one finishes, the next begins. Items that fail
 * are skipped so the rest still complete. The job result carries the list of designs
 * (the Jobs screen fetches each image by id when the user clicks View).
 */
async function runStationeryGenerationBatch(
  jobId: string,
  params: { companyId: string; stationeryTypes: string[]; assetName: string; description?: string; userId: string; promptConfigId?: string; customPrompt?: string },
): Promise<void> {
  const { companyId, stationeryTypes, assetName, description, userId, promptConfigId, customPrompt } = params;
  const total = stationeryTypes.length;
  const items: StationeryItemResult[] = [];
  // Per-item reasons, so a job that fails outright can tell the user WHY rather
  // than only that it failed.
  const failures: string[] = [];

  for (let i = 0; i < total; i++) {
    const stationeryType = stationeryTypes[i];
    const base = Math.round((i / total) * 100);
    const span = 100 / total;
    updateJobProgress(jobId, Math.min(99, base), `Generating item ${i + 1} of ${total}…`);
    try {
      const result = await generateOneStationeryDesign(
        { companyId, stationeryType, assetName, description, userId, promptConfigId, customPrompt },
        (pct, step) => updateJobProgress(jobId, Math.min(99, base + Math.round((pct / 100) * span)), `(${i + 1}/${total}) ${step}`),
      );
      items.push(result);
    } catch (err: any) {
      const reason = err?.message || 'unknown error';
      console.warn(`[Stationery] Item "${stationeryType}" failed: ${reason}`);
      failures.push(`${stationeryType}: ${reason}`);
    }
  }

  if (items.length === 0) {
    // Carry the underlying reasons through — the job error is what the user sees.
    const detail = failures.length ? ` ${failures.slice(0, 3).join(' | ')}` : '';
    throw new Error(`All stationery designs failed to generate.${detail}`);
  }
  completeJob(
    jobId,
    { autoFillData: { items, assetName, total, failed: total - items.length }, source: 'openai' },
    'openai',
  );
}

/**
 * POST /stationery-generate — start a non-blocking job that generates the selected
 * stationery items ONE AT A TIME. Accepts `stationeryTypes` (array) or a single
 * `stationeryType` for backward compatibility.
 */
router.post('/stationery-generate', requirePermission('ai-processing', 'ai-generate'), async (req: Request, res: Response) => {
  const { companyId, stationeryTypes, stationeryType, assetName, description, promptConfigId, customPrompt } = req.body;
  const types: string[] = Array.isArray(stationeryTypes)
    ? stationeryTypes.filter((t: any) => typeof t === 'string' && t.trim())
    : (stationeryType ? [stationeryType] : []);
  if (!companyId || types.length === 0 || !assetName) {
    res.status(400).json({ error: 'companyId, at least one stationeryType, and assetName are required' });
    return;
  }
  if (!authorizeCompany(req, companyId)) {
    res.status(403).json({ error: 'Access denied' });
    return;
  }
  const job = createJob('stationery-generation', companyId, 'stationery-generation');
  res.status(202).json({ jobId: job.jobId, status: 'processing' });
  const userId = req.user!.id;
  setImmediate(async () => {
    try {
      await runStationeryGenerationBatch(job.jobId, { companyId, stationeryTypes: types, assetName, description, userId, promptConfigId, customPrompt });
    } catch (err: any) {
      failJob(job.jobId, err?.message || 'Stationery generation failed');
    }
  });
});

/** POST /stationery-regenerate — regenerate a single design from user feedback (new version). */
router.post('/stationery-regenerate', requirePermission('ai-processing', 'ai-generate'), async (req: Request, res: Response) => {
  const { companyId, stationeryType, assetName, description, feedback, generationId, stationeryId, promptConfigId, customPrompt } = req.body;
  if (!companyId || !stationeryType || !assetName) {
    res.status(400).json({ error: 'companyId, stationeryType and assetName are required' });
    return;
  }
  if (!authorizeCompany(req, companyId)) {
    res.status(403).json({ error: 'Access denied' });
    return;
  }
  const job = createJob('stationery-generation', companyId, 'stationery-generation');
  res.status(202).json({ jobId: job.jobId, status: 'processing' });
  const userId = req.user!.id;
  setImmediate(async () => {
    try {
      const result = await generateOneStationeryDesign(
        { companyId, stationeryType, assetName, description, userId, feedback, existingGenerationId: generationId, stationeryId, promptConfigId, customPrompt },
        (pct, step) => updateJobProgress(job.jobId, pct, step),
      );
      completeJob(job.jobId, { autoFillData: { items: [result], assetName, total: 1, failed: 0 }, source: 'openai' }, 'openai');
    } catch (err: any) {
      failJob(job.jobId, err?.message || 'Stationery regeneration failed');
    }
  });
});

/** GET /status/:jobId — poll the status of a stationery generation job. */
router.get('/status/:jobId', async (req: Request, res: Response) => {
  const job = getJob(req.params.jobId);
  if (!job) {
    res.status(404).json({ error: 'Job not found' });
    return;
  }
  res.json({
    jobId: job.jobId,
    status: job.status,
    progress: job.progress,
    step: job.step,
    result: job.result,
    error: job.error,
  });
});

export default router;