const REFERENCE_TOKEN_PATTERN = /@(?:图片|图)(\d+)/g;

export function referenceToken(index: number): string {
  return `@图片${index + 1}`;
}

export function promptReferencesImage(prompt: string, index: number): boolean {
  const expected = index + 1;
  for (const match of prompt.matchAll(REFERENCE_TOKEN_PATTERN)) {
    if (Number(match[1]) === expected) return true;
  }
  return false;
}

export function findUnboundReferenceNumbers(prompt: string, boundReferenceCount: number): number[] {
  const unbound = new Set<number>();
  for (const match of prompt.matchAll(REFERENCE_TOKEN_PATTERN)) {
    const number = Number(match[1]);
    if (number < 1 || number > boundReferenceCount) unbound.add(number);
  }
  return [...unbound].sort((left, right) => left - right);
}

export function findReferenceMentionRange(prompt: string, cursor: number): {
  start: number;
  end: number;
} | null {
  const prefix = prompt.slice(0, cursor);
  const match = /@(?:图(?:片)?)?\d*$/.exec(prefix);
  if (!match || match.index < 0) return null;
  return { start: match.index, end: cursor };
}

export function insertReferenceToken(
  prompt: string,
  index: number,
  range?: { start: number; end: number } | null,
): { value: string; cursor: number } {
  const token = referenceToken(index);
  if (range) {
    const trailingSpace = prompt.slice(range.end).startsWith(' ') ? '' : ' ';
    const value = `${prompt.slice(0, range.start)}${token}${trailingSpace}${prompt.slice(range.end)}`;
    return { value, cursor: range.start + token.length + trailingSpace.length };
  }

  const trimmedRight = prompt.trimEnd();
  const separator = trimmedRight ? (/[，。！？、：；,.!?;:]$/.test(trimmedRight) ? ' ' : '，参考 ') : '参考 ';
  const value = `${trimmedRight}${separator}${token} `;
  return { value, cursor: value.length };
}

export function removeReferenceTokenAndShift(prompt: string, removedIndex: number): string {
  const removedNumber = removedIndex + 1;
  const removedMarker = '__MAKEL0RE_REMOVED_REFERENCE__';
  return prompt
    .replace(REFERENCE_TOKEN_PATTERN, (token, rawNumber: string) => {
      const number = Number(rawNumber);
      if (number === removedNumber) return removedMarker;
      if (number > removedNumber) return `@图片${number - 1}`;
      return token.startsWith('@图') && !token.startsWith('@图片')
        ? `@图片${number}`
        : token;
    })
    .replace(new RegExp(`\\s*${removedMarker}\\s*[，、,]?\\s*`, 'g'), ' ')
    .replace(/[ \t]{2,}/g, ' ')
    .replace(/，\s*([，。！？、：；,.!?;:]|$)/g, '$1')
    .trim();
}