aigen/utils/promptParts.ts

export type PromptParts = {
  pre: string
  prompt: string
  post: string
}

export function sanitizePromptPart(raw: unknown) {
  return String(raw || '').replace(/\r\n/g, '\n').replace(/\r/g, '')
}

export function composePromptParts(pre: string, prompt: string, post: string) {
  return [pre, prompt, post]
    .map(part => sanitizePromptPart(part).trim())
    .filter(Boolean)
    .join('\n\n')
}

export type PromptWrapperFields = {
  promptPre?: string
  promptPost?: string
}

/** Combine wrappers for Comfy. If mid already includes pre/post, do not wrap twice. */
export function joinPromptParts(pre?: string, mid?: string, post?: string) {
  const parts = restorePromptParts(mid || '', {
    pre: pre || '',
    prompt: mid || '',
    post: post || ''
  })
  return composePromptParts(parts.pre, parts.prompt, parts.post)
}

/** Prefer stored mid when wrappers exist. If wrappers are missing, use the already-joined fallback. */
export function wrappedPromptForComfy(
  wrappers?: PromptWrapperFields | null,
  mid?: string,
  fallback?: string
) {
  const pre = sanitizePromptPart(wrappers?.promptPre || '').trim()
  const post = sanitizePromptPart(wrappers?.promptPost || '').trim()
  const fallbackText = sanitizePromptPart(fallback || '').trim()
  const midText = sanitizePromptPart(mid || '').trim()
  if (pre || post) {
    return composePromptParts(pre, peelPromptWrappers(midText || fallbackText, pre, post), post)
  }
  return fallbackText || midText
}

export function ingestPromptParts(fields: {
  prompt?: unknown
  promptMid?: unknown
  promptPre?: unknown
  promptPost?: unknown
}) {
  const promptPre = sanitizePromptPart(fields.promptPre).trim()
  const promptPost = sanitizePromptPart(fields.promptPost).trim()
  const postedPrompt = sanitizePromptPart(fields.prompt).trim()
  const postedMid = sanitizePromptPart(fields.promptMid).trim()
  const promptMid = postedMid || restorePromptParts(postedPrompt, {
    pre: promptPre,
    prompt: postedPrompt,
    post: promptPost
  }).prompt
  const fromParts = composePromptParts(promptPre, promptMid, promptPost)
  return {
    promptPre,
    promptPost,
    promptMid,
    prompt: (promptPre || promptPost) ? fromParts : (postedPrompt || fromParts)
  }
}

export function applyPromptParts(parts: PromptParts, middle: string) {
  return composePromptParts(parts.pre, middle, parts.post)
}

function peelPromptWrappers(mid: string, pre: string, post: string) {
  let next = sanitizePromptPart(mid)
  if (pre && next.startsWith(pre)) next = next.slice(pre.length).replace(/^\n+/, '')
  if (post && next.endsWith(post)) next = next.slice(0, Math.max(0, next.length - post.length)).replace(/\n+$/, '')
  return next.trim()
}

export function restorePromptParts(saved: string, stored?: Partial<PromptParts> | null): PromptParts {
  const pre = sanitizePromptPart(stored?.pre ?? '').trim()
  const post = sanitizePromptPart(stored?.post ?? '').trim()
  return {
    pre,
    prompt: peelPromptWrappers(stored?.prompt ?? saved, pre, post),
    post
  }
}

function hasWrapperKey(item: object | null | undefined, key: 'promptPre' | 'promptPost') {
  return Boolean(item && Object.prototype.hasOwnProperty.call(item, key))
}

/** First non-empty wrapper wins. Blank strings fall through so a shot can inherit family pre/post. */
export function resolvePromptWrappers(
  ...sources: Array<PromptWrapperFields | null | undefined>
) {
  let promptPre = ''
  let promptPost = ''
  let sawPre = false
  let sawPost = false
  for (const source of sources) {
    if (!source) continue
    if (!sawPre && hasWrapperKey(source, 'promptPre') && source.promptPre != null) {
      const value = sanitizePromptPart(source.promptPre)
      if (value.trim()) {
        promptPre = value
        sawPre = true
      }
    }
    if (!sawPost && hasWrapperKey(source, 'promptPost') && source.promptPost != null) {
      const value = sanitizePromptPart(source.promptPost)
      if (value.trim()) {
        promptPost = value
        sawPost = true
      }
    }
    if (sawPre && sawPost) break
  }
  return { promptPre, promptPost }
}

export function persistPromptWrappers(item?: PromptWrapperFields | null): PromptWrapperFields {
  if (!item) return {}
  const next: PromptWrapperFields = {}
  if (hasWrapperKey(item, 'promptPre') && item.promptPre != null) next.promptPre = sanitizePromptPart(item.promptPre)
  if (hasWrapperKey(item, 'promptPost') && item.promptPost != null) next.promptPost = sanitizePromptPart(item.promptPost)
  return next
}

export function formatPromptParts(
  stored?: { prompt?: string; promptPre?: string; promptPost?: string } | null,
  midOverride?: string
) {
  const pre = sanitizePromptPart(stored?.promptPre || '').trim()
  const post = sanitizePromptPart(stored?.promptPost || '').trim()
  const prompt = peelPromptWrappers(midOverride ?? stored?.prompt ?? '', pre, post)
  const blocks: string[] = []
  if (pre) blocks.push(`Pre-text\n${pre}`)
  if (prompt) blocks.push(pre || post ? `Prompt\n${prompt}` : prompt)
  if (post) blocks.push(`Post-text\n${post}`)
  return blocks.join('\n\n')
}