aigen/utils/promptSequences.ts

export type PromptSequenceOutfit = {
  id: string
  name: string
  description: string
}

export type PromptSequenceCharacter = {
  id: string
  name: string
  description: string
  outfits: PromptSequenceOutfit[]
}

export type PromptSequenceCharacterState = {
  characterId: string
  on: boolean
  /** Active outfit id for this scene, or null/'' for none. */
  outfitId: string | null
}

/** Suggested labels only — not a hard enum. */
export const PROMPT_SEQUENCE_FRAMING_SUGGESTIONS = ['wide', 'medium', 'narrow'] as const

export type PromptSequenceFraming = {
  id: string
  label: string
  text: string
}

export type PromptSequenceSet = {
  id: string
  name: string
  framings: PromptSequenceFraming[]
}

export type PromptSequenceSetState = {
  setId: string
  on: boolean
  /** Chosen framing for this set in this scene (like outfitId). Empty = copy nothing. */
  framingId: string | null
}

export type PromptSequenceScene = {
  id: string
  title: string
  style: string
  settings: string
  setStates: PromptSequenceSetState[]
  characterStates: PromptSequenceCharacterState[]
  prompt: string
}

/** Scene prompt sections. Default copy order. */
export const PROMPT_SEQUENCE_SECTION_KEYS = ['style', 'settings', 'set', 'characters', 'prompt'] as const
export type PromptSequenceSectionKey = (typeof PROMPT_SEQUENCE_SECTION_KEYS)[number]
export const DEFAULT_PROMPT_SEQUENCE_SECTION_ORDER: PromptSequenceSectionKey[] = [
  ...PROMPT_SEQUENCE_SECTION_KEYS
]

export type PromptSequence = {
  id: string
  name: string
  /** Order of scene sections when assembling / in the editor. */
  sectionOrder: PromptSequenceSectionKey[]
  characters: PromptSequenceCharacter[]
  sets: PromptSequenceSet[]
  scenes: PromptSequenceScene[]
  createdAt: number
  updatedAt: number
}

/** Valid unique keys in given order; missing keys append in default order. */
export function normalizeSectionOrder(raw: unknown): PromptSequenceSectionKey[] {
  const allowed = new Set<string>(PROMPT_SEQUENCE_SECTION_KEYS)
  const seen = new Set<PromptSequenceSectionKey>()
  const ordered: PromptSequenceSectionKey[] = []
  if (Array.isArray(raw)) {
    for (const item of raw) {
      const key = String(item || '').trim()
      if (!allowed.has(key) || seen.has(key as PromptSequenceSectionKey)) continue
      const typed = key as PromptSequenceSectionKey
      seen.add(typed)
      ordered.push(typed)
    }
  }
  for (const key of DEFAULT_PROMPT_SEQUENCE_SECTION_ORDER) {
    if (!seen.has(key)) ordered.push(key)
  }
  return ordered
}

export const PROMPT_SEQUENCE_NAME_MAX = 80
export const PROMPT_SEQUENCE_FIELD_MAX = 50_000
export const PROMPT_SEQUENCE_TITLE_MAX = 60

export function newId() {
  if (typeof crypto !== 'undefined' && typeof crypto.randomUUID === 'function') {
    return crypto.randomUUID()
  }
  return `ps-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 10)}`
}

export function emptyOutfit(): PromptSequenceOutfit {
  return { id: newId(), name: '', description: '' }
}

export function emptyCharacter(): PromptSequenceCharacter {
  return { id: newId(), name: '', description: '', outfits: [] }
}

export function emptyFraming(label = ''): PromptSequenceFraming {
  return { id: newId(), label, text: '' }
}

export function emptySet(): PromptSequenceSet {
  return { id: newId(), name: '', framings: [emptyFraming()] }
}

export function defaultCharacterState(characterId: string, on = true): PromptSequenceCharacterState {
  return { characterId, on, outfitId: null }
}

export function defaultSetState(setId: string, on = true, framingId: string | null = null): PromptSequenceSetState {
  return { setId, on, framingId }
}

export function emptyScene(
  index = 1,
  characters: PromptSequenceCharacter[] = [],
  sets: PromptSequenceSet[] = []
): PromptSequenceScene {
  return {
    id: newId(),
    title: `Scene ${Math.max(1, index)}`,
    style: '',
    settings: '',
    setStates: sets.map(set => defaultSetState(set.id, true, null)),
    characterStates: characters.map(character => defaultCharacterState(character.id, true)),
    prompt: ''
  }
}

export function cloneSceneForNext(previous: PromptSequenceScene, nextIndex: number): PromptSequenceScene {
  return {
    id: newId(),
    title: `Scene ${nextIndex}`,
    style: previous.style,
    settings: previous.settings,
    setStates: (previous.setStates || []).map(state => ({
      setId: state.setId,
      on: !!state.on,
      framingId: state.framingId || null
    })),
    characterStates: (previous.characterStates || []).map(state => ({
      characterId: state.characterId,
      on: !!state.on,
      outfitId: state.outfitId || null
    })),
    prompt: previous.prompt
  }
}

function outfitPromptText(outfit?: { name?: string; description?: string } | null) {
  if (!outfit) return ''
  return String(outfit.description || '').trim() || String(outfit.name || '').trim()
}

/**
 * Character line when on:
 * - with outfit: "Name: description. Outfit: outfit description"
 * - without: "Name: description" (or name-only / description-only)
 */
export function formatCharacterLine(
  character: { name?: string; description?: string },
  outfit?: { name?: string; description?: string } | null
) {
  const name = String(character?.name || '').trim()
  const description = String(character?.description || '').trim()
  let base = ''
  if (name && description) base = `${name}: ${description}`
  else base = name || description
  if (!base) return ''
  const outfitText = outfitPromptText(outfit)
  if (!outfitText) return base
  return `${base}. Outfit: ${outfitText}`
}

/** Selected framing text for each on set. No set name, no framing label. */
export function assembleSetBlock(
  scene: { setStates?: PromptSequenceSetState[] },
  sets: PromptSequenceSet[] = []
) {
  const states = scene.setStates || []
  const stateById = new Map(states.map(state => [state.setId, state]))
  const chunks: string[] = []
  for (const set of sets) {
    const state = stateById.get(set.id)
    if (state && !state.on) continue
    const framingId = state?.framingId || null
    if (!framingId) continue
    const framing = (set.framings || []).find(item => item.id === framingId)
    const text = String(framing?.text || '').trim()
    if (text) chunks.push(text)
  }
  return chunks.join('\n')
}

export function resolveSceneCharacterLines(
  scene: { characterStates?: PromptSequenceCharacterState[] },
  characters: PromptSequenceCharacter[] = []
) {
  const states = scene.characterStates || []
  const stateById = new Map(states.map(state => [state.characterId, state]))
  return characters
    .map((character) => {
      const state = stateById.get(character.id)
      if (state && !state.on) return ''
      const outfitId = state?.outfitId || null
      const outfit = outfitId
        ? (character.outfits || []).find(item => item.id === outfitId) || null
        : null
      return formatCharacterLine(character, outfit)
    })
    .filter(Boolean)
}

export function assembleSceneText(
  scene: {
    style?: string
    settings?: string
    prompt?: string
    setStates?: PromptSequenceSetState[]
    characterStates?: PromptSequenceCharacterState[]
    /** Legacy per-scene characters (migration / old tests). */
    characters?: Array<{ name?: string; description?: string }>
  },
  sequenceCharacters?: PromptSequenceCharacter[],
  sequenceSets?: PromptSequenceSet[],
  sectionOrder?: PromptSequenceSectionKey[] | unknown
) {
  const setBlock = sequenceSets?.length
    ? assembleSetBlock(scene, sequenceSets)
    : ''
  const characterBlock = sequenceCharacters?.length
    ? resolveSceneCharacterLines(scene, sequenceCharacters).join('\n')
    : (scene.characters || [])
        .map(character => formatCharacterLine(character))
        .filter(Boolean)
        .join('\n')
  const blocks: Record<PromptSequenceSectionKey, string> = {
    style: String(scene.style || '').trim(),
    settings: String(scene.settings || '').trim(),
    set: setBlock,
    characters: characterBlock,
    prompt: String(scene.prompt || '').trim()
  }
  const order = normalizeSectionOrder(sectionOrder)
  const parts = order.map(key => blocks[key]).filter(Boolean)
  return parts.join('\n\n')
}

/**
 * Whole sequence as one copyable block.
 * Scenes join with a double blank line (\n\n\n) — same separator desktop Batch uses.
 * Sequence name is never included.
 */
export function assembleSequenceText(sequence: {
  sectionOrder?: PromptSequenceSectionKey[] | unknown
  characters?: PromptSequenceCharacter[]
  sets?: PromptSequenceSet[]
  scenes?: Array<{
    style?: string
    settings?: string
    prompt?: string
    setStates?: PromptSequenceSetState[]
    characterStates?: PromptSequenceCharacterState[]
    characters?: Array<{ name?: string; description?: string }>
  }>
}) {
  const characters = sequence.characters || []
  const sets = sequence.sets || []
  const sectionOrder = normalizeSectionOrder(sequence.sectionOrder)
  return (sequence.scenes || [])
    .map(scene => assembleSceneText(
      scene,
      characters.length ? characters : undefined,
      sets.length ? sets : undefined,
      sectionOrder
    ))
    .filter(Boolean)
    .join('\n\n\n')
}

export function normalizePromptSequenceName(raw: unknown) {
  return String(raw || '').trim().replace(/\s+/g, ' ').slice(0, PROMPT_SEQUENCE_NAME_MAX)
}

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

export function normalizePromptSequenceTitle(raw: unknown, fallback: string) {
  const title = String(raw || '').trim().replace(/\s+/g, ' ').slice(0, PROMPT_SEQUENCE_TITLE_MAX)
  return title || fallback
}

export function normalizeOutfit(raw: unknown): PromptSequenceOutfit {
  const rec = (raw && typeof raw === 'object' ? raw : {}) as Record<string, unknown>
  return {
    id: String(rec.id || '').trim() || newId(),
    name: String(rec.name || '').trim().slice(0, 120),
    description: normalizePromptSequenceField(rec.description).trim()
  }
}

export function normalizeCharacter(raw: unknown): PromptSequenceCharacter {
  const rec = (raw && typeof raw === 'object' ? raw : {}) as Record<string, unknown>
  const outfits = Array.isArray(rec.outfits) ? rec.outfits.map(normalizeOutfit) : []
  return {
    id: String(rec.id || '').trim() || newId(),
    name: String(rec.name || '').trim().slice(0, 120),
    description: normalizePromptSequenceField(rec.description).trim(),
    outfits
  }
}

export function normalizeFraming(raw: unknown): PromptSequenceFraming {
  const rec = (raw && typeof raw === 'object' ? raw : {}) as Record<string, unknown>
  return {
    id: String(rec.id || '').trim() || newId(),
    label: String(rec.label || '').trim().slice(0, 120),
    text: normalizePromptSequenceField(rec.text)
  }
}

/** Fold legacy beat-only sets into one framing labeled "wide". */
function framingsFromLegacyBeats(rec: Record<string, unknown>): PromptSequenceFraming[] {
  const parts: string[] = []
  const description = normalizePromptSequenceField(rec.description).trim()
  if (description) parts.push(description)
  const beats = Array.isArray(rec.beats) ? rec.beats : []
  for (const beat of beats) {
    if (!beat || typeof beat !== 'object') continue
    const text = normalizePromptSequenceField((beat as Record<string, unknown>).text).trim()
    if (text) parts.push(text)
  }
  if (!parts.length) return [emptyFraming()]
  return [{ id: newId(), label: 'wide', text: parts.join('\n') }]
}

export function normalizeSet(raw: unknown): PromptSequenceSet {
  const rec = (raw && typeof raw === 'object' ? raw : {}) as Record<string, unknown>
  let framings: PromptSequenceFraming[]
  if (Array.isArray(rec.framings) && rec.framings.length) {
    framings = rec.framings.map(normalizeFraming)
  } else if (Array.isArray(rec.beats) && rec.beats.length) {
    framings = framingsFromLegacyBeats(rec)
  } else if (String(rec.description || '').trim()) {
    framings = framingsFromLegacyBeats(rec)
  } else {
    framings = [emptyFraming()]
  }
  return {
    id: String(rec.id || '').trim() || newId(),
    name: String(rec.name || '').trim().slice(0, 120),
    framings
  }
}

export function normalizeCharacterState(raw: unknown, fallbackCharacterId = ''): PromptSequenceCharacterState {
  const rec = (raw && typeof raw === 'object' ? raw : {}) as Record<string, unknown>
  const characterId = String(rec.characterId || fallbackCharacterId || '').trim()
  const outfitRaw = rec.outfitId
  const outfitId = outfitRaw == null || outfitRaw === '' ? null : String(outfitRaw).trim() || null
  return {
    characterId,
    on: rec.on === undefined ? true : !!rec.on,
    outfitId
  }
}

export function normalizeSetState(
  raw: unknown,
  fallbackSetId = '',
  sets: PromptSequenceSet[] = [],
  legacyShotSize = ''
): PromptSequenceSetState {
  const rec = (raw && typeof raw === 'object' ? raw : {}) as Record<string, unknown>
  const setId = String(rec.setId || fallbackSetId || '').trim()
  const framingRaw = rec.framingId
  let framingId = framingRaw == null || framingRaw === '' ? null : String(framingRaw).trim() || null
  const set = sets.find(item => item.id === setId)
  if (framingId && set && !set.framings.some(framing => framing.id === framingId)) {
    framingId = null
  }
  if (!framingId && legacyShotSize && set) {
    const match = set.framings.find(
      framing => framing.label.trim().toLowerCase() === legacyShotSize.trim().toLowerCase()
    )
    if (match) framingId = match.id
  }
  return {
    setId,
    on: rec.on === undefined ? true : !!rec.on,
    framingId
  }
}

function liftLegacyCharacters(scenesRaw: unknown[]): PromptSequenceCharacter[] {
  const byId = new Map<string, PromptSequenceCharacter>()
  for (const scene of scenesRaw) {
    if (!scene || typeof scene !== 'object') continue
    const chars = (scene as Record<string, unknown>).characters
    if (!Array.isArray(chars)) continue
    for (const character of chars) {
      const normalized = normalizeCharacter(character)
      if (!byId.has(normalized.id)) byId.set(normalized.id, normalized)
    }
  }
  return [...byId.values()]
}

function syncCharacterStates(
  statesRaw: unknown,
  characters: PromptSequenceCharacter[],
  defaultOn = true
): PromptSequenceCharacterState[] {
  const incoming = Array.isArray(statesRaw)
    ? statesRaw.map(state => normalizeCharacterState(state)).filter(state => state.characterId)
    : []
  const byId = new Map(incoming.map(state => [state.characterId, state]))
  return characters.map((character) => {
    const existing = byId.get(character.id)
    if (existing) {
      const outfitOk = existing.outfitId
        && character.outfits.some(outfit => outfit.id === existing.outfitId)
      return {
        characterId: character.id,
        on: existing.on,
        outfitId: outfitOk ? existing.outfitId : null
      }
    }
    return defaultCharacterState(character.id, defaultOn)
  })
}

function syncSetStates(
  statesRaw: unknown,
  sets: PromptSequenceSet[],
  defaultOn = true,
  legacyShotSize = ''
): PromptSequenceSetState[] {
  const incoming = Array.isArray(statesRaw)
    ? statesRaw.map(state => normalizeSetState(state, '', sets, legacyShotSize)).filter(state => state.setId)
    : []
  const byId = new Map(incoming.map(state => [state.setId, state]))
  return sets.map((set) => {
    const existing = byId.get(set.id)
    if (existing) {
      const framingOk = existing.framingId
        && set.framings.some(framing => framing.id === existing.framingId)
      return {
        setId: set.id,
        on: existing.on,
        framingId: framingOk ? existing.framingId : null
      }
    }
    let framingId: string | null = null
    if (legacyShotSize) {
      const match = set.framings.find(
        framing => framing.label.trim().toLowerCase() === legacyShotSize.trim().toLowerCase()
      )
      if (match) framingId = match.id
    }
    return defaultSetState(set.id, defaultOn, framingId)
  })
}

export function normalizeScene(
  raw: unknown,
  index: number,
  characters: PromptSequenceCharacter[] = [],
  sets: PromptSequenceSet[] = []
): PromptSequenceScene {
  const rec = (raw && typeof raw === 'object' ? raw : {}) as Record<string, unknown>
  const legacyShotSize = String(rec.shotSize || '').trim()
  return {
    id: String(rec.id || '').trim() || newId(),
    title: normalizePromptSequenceTitle(rec.title, `Scene ${index + 1}`),
    style: normalizePromptSequenceField(rec.style),
    settings: normalizePromptSequenceField(rec.settings),
    setStates: syncSetStates(rec.setStates, sets, true, legacyShotSize),
    characterStates: syncCharacterStates(rec.characterStates, characters, true),
    prompt: normalizePromptSequenceField(rec.prompt)
  }
}

export function normalizePromptSequence(raw: unknown): PromptSequence | null {
  if (!raw || typeof raw !== 'object') return null
  const rec = raw as Record<string, unknown>
  const name = normalizePromptSequenceName(rec.name)
  if (!name) return null
  const scenesRaw = Array.isArray(rec.scenes) ? rec.scenes : []
  let characters = Array.isArray(rec.characters)
    ? rec.characters.map(normalizeCharacter)
    : liftLegacyCharacters(scenesRaw)
  if (!characters.length) characters = [emptyCharacter()]
  const sets = Array.isArray(rec.sets) ? rec.sets.map(normalizeSet) : []
  const scenes = scenesRaw.length
    ? scenesRaw.map((scene, index) => normalizeScene(scene, index, characters, sets))
    : [emptyScene(1, characters, sets)]
  const createdAt = Number(rec.createdAt) || Date.now()
  return {
    id: String(rec.id || '').trim() || newId(),
    name,
    sectionOrder: normalizeSectionOrder(rec.sectionOrder),
    characters,
    sets,
    scenes,
    createdAt,
    updatedAt: Number(rec.updatedAt) || createdAt
  }
}

/** Ensure every scene has a state row for every sequence character. */
export function ensureSceneStatesForCharacters(
  scenes: PromptSequenceScene[],
  characters: PromptSequenceCharacter[],
  opts?: { newCharacterId?: string; turnOnSceneIndex?: number }
) {
  const newCharacterId = opts?.newCharacterId
  const turnOnIndex = opts?.turnOnSceneIndex
  for (let i = 0; i < scenes.length; i++) {
    const scene = scenes[i]
    const byId = new Map(scene.characterStates.map(state => [state.characterId, state]))
    scene.characterStates = characters.map((character) => {
      const existing = byId.get(character.id)
      if (existing) return existing
      const on = newCharacterId && character.id === newCharacterId
        ? i === turnOnIndex
        : true
      return defaultCharacterState(character.id, on)
    })
  }
}

/** Ensure every scene has a state row for every sequence set. */
export function ensureSceneStatesForSets(
  scenes: PromptSequenceScene[],
  sets: PromptSequenceSet[],
  opts?: { newSetId?: string; turnOnSceneIndex?: number }
) {
  const newSetId = opts?.newSetId
  const turnOnIndex = opts?.turnOnSceneIndex
  for (let i = 0; i < scenes.length; i++) {
    const scene = scenes[i]
    const byId = new Map((scene.setStates || []).map(state => [state.setId, state]))
    scene.setStates = sets.map((set) => {
      const existing = byId.get(set.id)
      if (existing) {
        const framingOk = existing.framingId
          && set.framings.some(framing => framing.id === existing.framingId)
        return {
          setId: set.id,
          on: existing.on,
          framingId: framingOk ? existing.framingId : null
        }
      }
      const on = newSetId && set.id === newSetId
        ? i === turnOnIndex
        : true
      return defaultSetState(set.id, on, null)
    })
  }
}