aigen/utils/globalLocks.ts

export const PERMANENCE_REF_LABELS = ['Subject', 'Vehicle', 'Interior', 'Object'] as const

export type PermanenceRefLabel = (typeof PERMANENCE_REF_LABELS)[number]

export interface PermanenceRef {
  stillId: string
  label: string
}

export const DEFAULT_GLOBAL_LOCKS = `[GLOBAL LOCKS]
Subject:
Vehicle:
Interior:
Object:
Extras:
Camera:
Speech:
Never:
State:`

const LOCK_FIELD = /^(Subject|Vehicle|Interior|Object|Extras|Camera|Speech|Never|State)\s*:?\s*/i

export function normalizePermanenceRefs(raw: unknown, limit = 8): PermanenceRef[] {
  if (!Array.isArray(raw)) return []
  const seen = new Set<string>()
  const refs: PermanenceRef[] = []
  for (const item of raw) {
    if (refs.length >= limit) break
    const stillId = String((item && (item as { stillId?: unknown }).stillId) || '').trim()
    if (!stillId || seen.has(stillId)) continue
    const label = sanitizeRefLabel((item as { label?: unknown })?.label)
    seen.add(stillId)
    refs.push({ stillId, label })
  }
  return refs
}

export function sanitizeRefLabel(raw: unknown) {
  const label = String(raw || '').trim().slice(0, 24)
  if (!label) return 'Object'
  const match = PERMANENCE_REF_LABELS.find(item => item.toLowerCase() === label.toLowerCase())
  return match || label
}

export function nextPermanenceLabel(existing: PermanenceRef[]) {
  const used = new Set(existing.map(item => item.label))
  return PERMANENCE_REF_LABELS.find(label => !used.has(label)) || 'Object'
}

export function mergePermanenceRefs(family?: PermanenceRef[] | null, shot?: PermanenceRef[] | null) {
  return normalizePermanenceRefs([...(family || []), ...(shot || [])])
}

export function isBlankGlobalLocks(text: string) {
  const stripped = String(text || '')
    .replace(/\[GLOBAL LOCKS\]/gi, '')
    .replace(/\[INHERITS\]/gi, '')
    .replace(LOCK_FIELD, '')
    .replace(/\r\n/g, '\n')
    .split('\n')
    .map(line => line.replace(LOCK_FIELD, '').trim())
    .filter(Boolean)
    .join('')
    .trim()
  return !stripped
}

export function looksAlreadyInherited(prompt: string) {
  return /^\s*\[(GLOBAL LOCKS|INHERITS)\]/i.test(String(prompt || ''))
}

export function extractGlobalLocks(text: string): { locks: string; script: string } {
  const src = String(text || '').replace(/\r\n/g, '\n')
  const marker = /\[GLOBAL LOCKS\]/i.exec(src)
  if (!marker || marker.index === undefined) return { locks: '', script: src.trim() }
  const after = src.slice(marker.index)
  const shot = /^[ \t]*shot\s+\d+\b/im.exec(after)
  if (!shot || shot.index === undefined) return { locks: '', script: src.trim() }
  const locks = after.slice(0, shot.index).replace(/\n\[INHERITS\]\s*$/i, '').trim()
  const script = `${src.slice(0, marker.index)}${after.slice(shot.index)}`.trim()
  return isBlankGlobalLocks(locks) ? { locks: '', script: src.trim() } : { locks, script }
}

export function resolveGlobalLocks(locks: string, firstPrompt: string) {
  if (!isBlankGlobalLocks(locks)) return { locks: String(locks || '').trim(), prompt: firstPrompt }
  const extracted = extractGlobalLocks(firstPrompt)
  return extracted.locks
    ? { locks: extracted.locks, prompt: extracted.script }
    : { locks: '', prompt: firstPrompt }
}

function extractState(locks: string) {
  const match = String(locks || '').match(/(?:^|\n)\s*State:\s*(.+?)(?=\n\s*(?:Subject|Vehicle|Interior|Object|Extras|Camera|Speech|Never)\s*:|\n\s*\[|$)/is)
  return String(match?.[1] || '').replace(/\s+/g, ' ').trim()
}

function permanenceLine(refs: PermanenceRef[]) {
  const labels = [...new Set(refs.map(item => item.label).filter(Boolean))]
  if (!labels.length) return ''
  return `Permanence: ${labels.join(', ')} stay the same.`
}

function inheritLine(locks: string, refs: PermanenceRef[]) {
  const parts = ['[INHERITS] Global locks.']
  const labels = [...new Set(refs.map(item => item.label).filter(Boolean))]
  if (labels.length) parts.push(`Same ${labels.join(', ')} as locked.`)
  const state = extractState(locks)
  if (state) parts.push(`State: ${state}`)
  return parts.join(' ')
}

export function applyGlobalLocks(params: {
  locks: string
  prompt: string
  shotIndex: number
  refs?: PermanenceRef[]
}) {
  const beat = String(params.prompt || '').trim()
  if (!beat) return beat
  if (looksAlreadyInherited(beat)) return beat
  const refs = normalizePermanenceRefs(params.refs)
  const locks = String(params.locks || '').trim()
  const blank = isBlankGlobalLocks(locks)
  if (blank && !refs.length) return beat
  if (params.shotIndex <= 0) {
    const chunks = [
      blank ? '' : locks,
      permanenceLine(refs),
      beat
    ].filter(Boolean)
    return chunks.join('\n\n')
  }
  return `${inheritLine(blank ? '' : locks, refs)}\n\n${beat}`
}

export function composeShotPrompt(params: {
  globalLocks?: string
  prompt: string
  shotIndex: number
  familyRefs?: PermanenceRef[] | null
  shotRefs?: PermanenceRef[] | null
}) {
  return applyGlobalLocks({
    locks: params.globalLocks || '',
    prompt: params.prompt,
    shotIndex: params.shotIndex,
    refs: mergePermanenceRefs(params.familyRefs, params.shotRefs)
  })
}