aigen/shared/studio2/photos.mjs

/**
 * Inspector Photos: role cards (plain English). The runner maps roles → graph sockets.
 * UI never teaches <imageN> / Hero / Edit source.
 */

export const PHOTO_ROLE = {
  person: 'person',
  change: 'change',
  outfit: 'outfit',
  look: 'look',
  first: 'first',
  extra: 'extra'
}

export const PHOTO_ROLE_LABEL = {
  person: 'Person to keep',
  change: 'Photo to change',
  outfit: 'Outfit / object',
  look: 'Look / style',
  first: 'First frame',
  extra: 'Extra photo'
}

/** Human phrase inserted by Mention in prompt (not a Comfy token). */
export function mentionPhrase(role, index = 0) {
  const r = String(role || '')
  if (r === PHOTO_ROLE.change) return 'this photo'
  if (r === PHOTO_ROLE.outfit) return 'the outfit photo'
  if (r === PHOTO_ROLE.person) return 'the face photo'
  if (r === PHOTO_ROLE.look) return 'the look photo'
  if (r === PHOTO_ROLE.first) return 'the first frame'
  return `Photo ${Number(index) + 1}`
}

export function insertMention(text, phrase, selectionStart = 0, selectionEnd = 0) {
  const value = String(text || '')
  const chip = String(phrase || '')
  if (!chip) return { text: value, selectionStart: 0, selectionEnd: 0 }
  const existing = value.indexOf(chip)
  if (existing >= 0) {
    return { text: value, selectionStart: existing, selectionEnd: existing + chip.length }
  }
  const start = Math.max(0, Math.min(Number(selectionStart) || 0, value.length))
  const end = Math.max(start, Math.min(Number(selectionEnd) || start, value.length))
  const next = `${value.slice(0, start)}${chip}${value.slice(end)}`
  return { text: next, selectionStart: start, selectionEnd: start + chip.length }
}

/**
 * @returns {{
 *   show: boolean,
 *   max: number,
 *   roles: string[],
 *   defaultRole: string,
 *   emptyHint: string|null,
 *   handoffReadonly: boolean
 * }}
 */
export function photosPlan({ mode, engine } = {}) {
  const task = String(mode || 'generate')
  const eng = String(engine || 'flux')

  if (task === 'music') {
    return { show: false, max: 0, roles: [], defaultRole: '', emptyHint: null, handoffReadonly: false }
  }

  if (eng === 'qwen21') {
    if (task === 'generate') {
      return { show: false, max: 0, roles: [], defaultRole: '', emptyHint: null, handoffReadonly: false }
    }
    if (task === 'edit') {
      return {
        show: true,
        max: 3,
        roles: [PHOTO_ROLE.change, PHOTO_ROLE.outfit, PHOTO_ROLE.person, PHOTO_ROLE.extra],
        defaultRole: PHOTO_ROLE.change,
        emptyHint: 'Add a photo, then say what it is.',
        handoffReadonly: false
      }
    }
    return { show: false, max: 0, roles: [], defaultRole: '', emptyHint: null, handoffReadonly: false }
  }

  if (eng === 'krea') {
    if (task === 'generate') {
      return { show: false, max: 0, roles: [], defaultRole: '', emptyHint: null, handoffReadonly: false }
    }
    if (task === 'refine') {
      return {
        show: true,
        max: 1,
        roles: [PHOTO_ROLE.change],
        defaultRole: PHOTO_ROLE.change,
        emptyHint: 'Add a photo, then say what it is.',
        handoffReadonly: false
      }
    }
    return { show: false, max: 0, roles: [], defaultRole: '', emptyHint: null, handoffReadonly: false }
  }

  if (task === 'extend') {
    return {
      show: true,
      max: 0,
      roles: [],
      defaultRole: '',
      emptyHint: null,
      handoffReadonly: true
    }
  }

  if (task === 'video') {
    if (eng === 'minimax') {
      return {
        show: true,
        max: 2,
        roles: [PHOTO_ROLE.first, PHOTO_ROLE.person, PHOTO_ROLE.look],
        defaultRole: PHOTO_ROLE.first,
        emptyHint: 'Add a photo, then say what it is.',
        handoffReadonly: false
      }
    }
    // LTX: first frame only — existing start binding
    return {
      show: true,
      max: 1,
      roles: [PHOTO_ROLE.first],
      defaultRole: PHOTO_ROLE.first,
      emptyHint: 'Add a photo, then say what it is.',
      handoffReadonly: false
    }
  }

  // Klein (flux)
  if (task === 'generate') {
    return {
      show: true,
      max: 1,
      roles: [PHOTO_ROLE.person],
      defaultRole: PHOTO_ROLE.person,
      emptyHint: 'Add a photo, then say what it is.',
      handoffReadonly: false
    }
  }
  if (task === 'edit' || task === 'iterate') {
    return {
      show: true,
      max: 2,
      roles: [PHOTO_ROLE.change, PHOTO_ROLE.person],
      defaultRole: PHOTO_ROLE.change,
      emptyHint: 'Add a photo, then say what it is.',
      handoffReadonly: false
    }
  }
  if (task === 'compose') {
    return {
      show: true,
      max: 3,
      roles: [PHOTO_ROLE.change, PHOTO_ROLE.person, PHOTO_ROLE.extra],
      defaultRole: PHOTO_ROLE.change,
      emptyHint: 'Add a photo, then say what it is.',
      handoffReadonly: false
    }
  }
  if (task === 'refine') {
    return {
      show: true,
      max: 1,
      roles: [PHOTO_ROLE.change],
      defaultRole: PHOTO_ROLE.change,
      emptyHint: 'Add a photo, then say what it is.',
      handoffReadonly: false
    }
  }

  return { show: false, max: 0, roles: [], defaultRole: '', emptyHint: null, handoffReadonly: false }
}

export function normalizePhotos(raw, plan) {
  const list = Array.isArray(raw) ? raw : []
  const roles = new Set(plan?.roles || [])
  const max = Math.max(0, Number(plan?.max) || 0)
  const out = []
  for (const row of list) {
    if (out.length >= max) break
    const stillId = String(row?.stillId || row?.id || '').trim()
    if (!stillId) continue
    let role = String(row?.role || plan?.defaultRole || '').trim()
    if (!roles.has(role)) role = plan?.defaultRole || [...roles][0] || PHOTO_ROLE.extra
    out.push({ stillId, role })
  }
  return out
}

/** First card wins when two share a role. Returns duplicate role label or null. */
export function duplicatePhotoRole(photos) {
  const seen = new Set()
  for (const row of photos || []) {
    const role = String(row?.role || '')
    if (!role) continue
    if (seen.has(role)) return PHOTO_ROLE_LABEL[role] || role
    seen.add(role)
  }
  return null
}

/**
 * Map Photos cards → legacy request fields the graphs already consume.
 * Qwen Edit: change→imageAId, outfit|extra→imageBId, person→imageBId if empty.
 * Klein: person→identityStillId, change→imageAId.
 */
export function bindPhotosToRequest(photos, { mode, engine } = {}) {
  const task = String(mode || 'generate')
  const eng = String(engine || 'flux')
  const list = Array.isArray(photos) ? photos.filter(p => p?.stillId) : []
  const find = (...roles) => list.find(p => roles.includes(p.role))

  let imageAId = ''
  let imageBId = ''
  let identityStillId = ''

  if (eng === 'qwen21' && task === 'edit') {
    let change = find(PHOTO_ROLE.change)
    if (!change && list[0]) change = { ...list[0], role: PHOTO_ROLE.change }
    const outfit = find(PHOTO_ROLE.outfit, PHOTO_ROLE.extra)
    const person = find(PHOTO_ROLE.person)
    imageAId = change?.stillId || ''
    imageBId = outfit?.stillId || ''
    if (!imageBId && person) imageBId = person.stillId
    // Person with outfit already in image_2: leave as imageBId only (2-input graph).
    identityStillId = ''
    return { imageAId, imageBId, identityStillId, photos: list }
  }

  if (eng === 'krea' && task === 'refine') {
    imageAId = find(PHOTO_ROLE.change)?.stillId || list[0]?.stillId || ''
    return { imageAId, imageBId: '', identityStillId: '', photos: list }
  }

  if (task === 'video' || (task === 'extend' && eng === 'minimax')) {
    imageAId = find(PHOTO_ROLE.first)?.stillId || list[0]?.stillId || ''
    identityStillId = find(PHOTO_ROLE.person, PHOTO_ROLE.look)?.stillId || ''
    return { imageAId, imageBId: '', identityStillId, photos: list }
  }

  if (eng === 'ltx' && task === 'video') {
    imageAId = find(PHOTO_ROLE.first)?.stillId || list[0]?.stillId || ''
    return { imageAId, imageBId: '', identityStillId: '', photos: list }
  }

  // Klein / flux
  if (task === 'generate') {
    identityStillId = find(PHOTO_ROLE.person)?.stillId || list[0]?.stillId || ''
    return { imageAId: '', imageBId: '', identityStillId, photos: list }
  }
  if (task === 'edit' || task === 'iterate' || task === 'compose' || task === 'refine') {
    imageAId = find(PHOTO_ROLE.change)?.stillId || list[0]?.stillId || ''
    identityStillId = find(PHOTO_ROLE.person)?.stillId || ''
    imageBId = find(PHOTO_ROLE.extra)?.stillId || ''
    return { imageAId, imageBId, identityStillId, photos: list }
  }

  return { imageAId, imageBId, identityStillId, photos: list }
}

/** Hydrate Photos cards from legacy hero/source fields when loading a job. */
export function photosFromRequest(request = {}) {
  if (Array.isArray(request.photos) && request.photos.length) {
    return request.photos.map(p => ({
      stillId: String(p.stillId || p.id || ''),
      role: String(p.role || PHOTO_ROLE.change)
    })).filter(p => p.stillId)
  }
  const mode = String(request.mode || 'generate')
  const engine = String(request.engine || 'flux')
  const out = []
  if (engine === 'qwen21' && mode === 'edit') {
    if (request.imageAId) out.push({ stillId: request.imageAId, role: PHOTO_ROLE.change })
    if (request.imageBId) out.push({ stillId: request.imageBId, role: PHOTO_ROLE.outfit })
    return out
  }
  if (mode === 'generate' && engine === 'flux' && request.identityStillId) {
    out.push({ stillId: request.identityStillId, role: PHOTO_ROLE.person })
    return out
  }
  if (['edit', 'iterate', 'compose', 'refine'].includes(mode)) {
    if (request.imageAId) out.push({ stillId: request.imageAId, role: PHOTO_ROLE.change })
    if (request.identityStillId && request.identityStillId !== request.imageAId) {
      out.push({ stillId: request.identityStillId, role: PHOTO_ROLE.person })
    }
    if (request.imageBId) out.push({ stillId: request.imageBId, role: PHOTO_ROLE.extra })
    return out
  }
  if (mode === 'video') {
    if (request.imageAId) out.push({ stillId: request.imageAId, role: PHOTO_ROLE.first })
    if (request.identityStillId) out.push({ stillId: request.identityStillId, role: PHOTO_ROLE.person })
  }
  return out
}

/**
 * Apply Photos → request fields. Call before validate / queue.
 * If Edit+Qwen has cards but no Photo to change, promote the first card.
 * When `photos` is an explicit array (even empty), it is source of truth —
 * do not resurrect cards from leftover imageAId / identityStillId.
 */
export function applyPhotosToRequest(raw = {}) {
  const next = { ...raw }
  const plan = photosPlan({ mode: next.mode, engine: next.engine })
  if (!plan.show || plan.handoffReadonly) {
    if (!Array.isArray(next.photos)) next.photos = []
    return next
  }
  const photosProvided = Array.isArray(raw.photos)
  let photos = normalizePhotos(next.photos, plan)
  if (!photosProvided && !photos.length && (next.imageAId || next.identityStillId || next.imageBId)) {
    photos = normalizePhotos(photosFromRequest(next), plan)
  }
  // Promote first card to Photo to change when Edit needs it.
  if (
    plan.roles.includes(PHOTO_ROLE.change)
    && photos.length
    && !photos.some(p => p.role === PHOTO_ROLE.change)
  ) {
    photos = [{ ...photos[0], role: PHOTO_ROLE.change }, ...photos.slice(1)]
  }
  const bound = bindPhotosToRequest(photos, { mode: next.mode, engine: next.engine })
  next.photos = bound.photos
  next.imageAId = bound.imageAId
  next.imageBId = bound.imageBId
  next.identityStillId = bound.identityStillId
  return next
}

/**
 * Qwen Edit: map mention phrases → <imageN> for TextEncode.
 * Photo to change → <image1>; outfit/extra/person-on-2 → <image2>.
 */
export function injectQwenEditMentionTags(prompt, photos = []) {
  let text = String(prompt || '')
  const list = Array.isArray(photos) ? photos.filter(p => p?.stillId) : []
  let change = list.find(p => p.role === PHOTO_ROLE.change)
  if (!change && list[0]) change = { ...list[0], role: PHOTO_ROLE.change }
  const outfit = list.find(p => p.role === PHOTO_ROLE.outfit || p.role === PHOTO_ROLE.extra)
  const person = list.find(p => p.role === PHOTO_ROLE.person)
  const image2 = outfit || (!outfit ? person : null)

  const pairs = []
  if (change) {
    const idx = list.indexOf(list.find(p => p.stillId === change.stillId && p.role === change.role) || change)
    pairs.push({ phrase: mentionPhrase(PHOTO_ROLE.change, Math.max(0, idx)), tag: '<image1>' })
    // Also map generic Photo 1 if they used chip style
    pairs.push({ phrase: 'Photo 1', tag: '<image1>' })
  }
  if (image2) {
    const idx = list.findIndex(p => p.stillId === image2.stillId)
    pairs.push({ phrase: mentionPhrase(image2.role, Math.max(0, idx)), tag: '<image2>' })
    pairs.push({ phrase: 'Photo 2', tag: '<image2>' })
  }

  // Longer phrases first so "the outfit photo" wins over fragments.
  pairs.sort((a, b) => b.phrase.length - a.phrase.length)
  for (const { phrase, tag } of pairs) {
    if (!phrase || text.includes(tag) && !text.includes(phrase)) continue
    if (text.includes(phrase)) text = text.split(phrase).join(tag)
  }
  return text
}