aigen/utils/globalLocks.ts

export interface PermanenceRef {
  stillId: string
  label: string
}

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

const KNOWN_LOCK_KEYS = 'Subject|Vehicle|Interior|Object|Extras|Camera|Speech|Never|State'
const LOCK_FIELD = new RegExp(`^(?:(?:${KNOWN_LOCK_KEYS})\\s*:?|[A-Za-z][A-Za-z0-9 _-]{0,22}\\s*:)\\s*`, 'i')
const LOCK_KEY_LINE = /^([A-Za-z][A-Za-z0-9 _-]{0,22})\s*:/
const SKIP_BINDING_KEYS = /^(speech|never|state)$/i

export function parseLockBindingKeys(text: string): string[] {
  const keys: string[] = []
  const seen = new Set<string>()
  for (const line of String(text || '').replace(/\r\n/g, '\n').split('\n')) {
    const trimmed = line.trim()
    if (!trimmed || trimmed.startsWith('[')) continue
    const match = trimmed.match(LOCK_KEY_LINE)
    if (!match) continue
    const key = match[1].replace(/\s+/g, ' ').trim()
    if (!key || SKIP_BINDING_KEYS.test(key)) continue
    const id = key.toLowerCase()
    if (seen.has(id)) continue
    seen.add(id)
    keys.push(key)
  }
  return keys
}

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) {
  return String(raw || '').replace(/\s+/g, ' ').trim().slice(0, 24)
}

export function nextPermanenceLabel(existing: PermanenceRef[], lockKeys: string[] = []) {
  const used = new Set(existing.map(item => item.label.trim().toLowerCase()).filter(Boolean))
  return lockKeys.find(key => !used.has(key.trim().toLowerCase())) || ''
}

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 }
}

export function permanenceLabels(refs: PermanenceRef[] | null | undefined) {
  return [...new Set(normalizePermanenceRefs(refs).map(item => item.label).filter(Boolean))]
}

const PERSON_WORD = /\b(man|men|woman|women|person|people|boy|girl|guy|lady|male|female|human|child|kid|figure|character|actor|actress)\b/i
const MAIN_SUBJECT_ALIAS = /^(subject|woman|women|she|her|female)$/i

export function extraPersonLabels(refs: PermanenceRef[] | null | undefined) {
  return permanenceLabels(refs).filter(label => {
    if (MAIN_SUBJECT_ALIAS.test(label)) return false
    return PERSON_WORD.test(label) || /^extras?$/i.test(label)
  })
}

export function shouldForceLastFrameI2V(refs?: PermanenceRef[] | null, _locks = '') {
  return extraPersonLabels(refs).length > 0
}

export function allowIdentityRefs(
  requested: boolean,
  familyRefs?: PermanenceRef[] | null,
  locks = '',
  shotRefs?: PermanenceRef[][] | null
) {
  if (!requested) return false
  return !shouldForceLastFrameI2V(mergePermanenceRefs(familyRefs, (shotRefs || []).flat()), locks)
}

function permanenceLine(refs: PermanenceRef[]) {
  const labels = permanenceLabels(refs)
  if (!labels.length) return ''
  return `Permanence: ${labels.join(', ')} stay the same.`
}

function inheritLine(refs: PermanenceRef[]) {
  const labels = permanenceLabels(refs)
  if (labels.length) return `[INHERITS] Same ${labels.join(', ')} as now`
  return '[INHERITS] Same as now'
}

export function stripLockHeader(prompt: string) {
  let src = String(prompt || '').replace(/\r\n/g, '\n').trim()
  while (/^\[(GLOBAL LOCKS|INHERITS)\]/i.test(src)) {
    if (/^\[GLOBAL LOCKS\]/i.test(src)) {
      const config = src.search(/\n\[SHOT CONFIGURATION\]/i)
      if (config >= 0) {
        src = src.slice(config).trim()
        break
      }
    }
    const split = src.search(/\n\s*\n/)
    src = (split >= 0 ? src.slice(split) : src.replace(/^\s*\[[^\]]+\][^\n]*/, '')).trim()
    if (/^\[GLOBAL LOCKS\]/i.test(src)) {
      src = src.replace(/^\s*\[GLOBAL LOCKS\]\s*\n?/, '').trim()
      break
    }
  }
  return src
}

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

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)
  })
}