import { RichText } from "@skyware/bot";

const MARKER = /\[(\d+)\]/g;

/**
 * Everything the model has been shown by tools in one request. Each source gets
 * a number the model can cite as `[n]`.
 */
export class SourceBook {
  readonly urls: string[] = [];

  /** Returns the number for a URL, registering it the first time it's seen. */
  add(url: string) {
    const index = this.urls.indexOf(url);
    if (index >= 0) return index + 1;
    this.urls.push(url);
    return this.urls.length;
  }

  /**
   * Adds a `source` number to everything citable inside a tool result: web
   * pages (`url`) and Bluesky posts (`uri`), including nested replies.
   */
  annotate<T>(result: T): T {
    if (Array.isArray(result)) {
      result.forEach((item) => this.annotate(item));
    } else if (result && typeof result === "object") {
      const item = result as Record<string, unknown>;
      const url = typeof item.url === "string" && /^https?:\/\//.test(item.url)
        ? item.url
        : typeof item.uri === "string"
        ? postUrl(item.uri)
        : undefined;
      if (url) item.source = this.add(url);

      if (Array.isArray(item.replies)) this.annotate(item.replies);
    }
    return result;
  }
}

/** Turns an at:// post URI into a bsky.app link. */
export function postUrl(uri: string) {
  const match = uri.match(/^at:\/\/([^/]+)\/app\.bsky\.feed\.post\/([^/]+)$/);
  return match ? `https://bsky.app/profile/${match[1]}/post/${match[2]}` : undefined;
}

/**
 * Cleans up the `[n]` markers in a model reply. Markers are renumbered in order
 * of appearance, and ones that don't match a real source are dropped. If the
 * model cited nothing, the pages it read are listed at the end instead.
 *
 * @param fallback Source numbers to list when the reply has no markers.
 * @returns The text and the URL behind each marker, where `links[0]` is `[1]`.
 */
export function finalizeCitations(
  text: string,
  book: SourceBook,
  fallback: number[],
) {
  const order: number[] = [];

  let result = text.replace(MARKER, (_, digits: string) => {
    const id = Number(digits);
    if (id < 1 || id > book.urls.length) return "";

    let position = order.indexOf(id);
    if (position < 0) position = order.push(id) - 1;
    return `[${position + 1}]`;
  });
  // Dropping a bogus marker can leave stray spaces behind
  result = result.replace(/[ \t]{2,}/g, " ").replace(/ +([.,;:!?])/g, "$1");

  if (order.length === 0 && fallback.length > 0) {
    const ids = [...new Set(fallback)];
    result = `${result.trimEnd()} ${ids.map((_, i) => `[${i + 1}]`).join(", ")}`;
    order.push(...ids);
  }

  return { text: result.trim(), links: order.map((id) => book.urls[id - 1]!) };
}

/** Removes citation markers, for text that is stored rather than shown. */
export const stripCitations = (text: string) =>
  text.replace(/ ?\[\d+\]/g, "").trim();

/** Builds a message where each `[n]` marker links to `links[n - 1]`. */
export function renderCitations(text: string, links: string[]) {
  const richText = new RichText();
  let last = 0;

  for (const match of text.matchAll(MARKER)) {
    const url = links[Number(match[1]) - 1];
    if (!url) continue;

    richText.addText(text.slice(last, match.index));
    richText.addLink(match[0], url);
    last = match.index + match[0].length;
  }

  return richText.addText(text.slice(last));
}