import { mapOutsideCodeSpans, markdownLink } from "./code";
import type { Source, ToolCallRecord } from "./types";

/**
 * Numbered sources for one turn, and the citation markers that point at them.
 *
 * Vicinae has no footnote support to lean on: it parses markdown with
 * `CMARK_OPT_DEFAULT` and attaches only the table and strikethrough extensions
 * (`src/server/src/ui/quick/markdown-model.cpp` upstream), so a `[^1]` written by
 * a model arrives in the pane as literal text. Everything a footnote would do
 * therefore has to happen here, on the way to the renderer.
 *
 * Imports no Vicinae module, so it unit tests with no runtime.
 */

/** What a turn uses to hand out numbers, one per page it consults. */
export interface SourceList {
  /**
   * The number for a URL, registering it the first time it is seen.
   *
   * `index` is 0 for a result with no usable address, which is the caller's
   * signal to print no marker at all. `source` comes back only when this call
   * registered something new, so a caller records each page exactly once.
   */
  add(title: string, url: string): { index: number; source?: Source };
}

/**
 * A fresh list, numbering from one.
 *
 * Scoped to a turn rather than a conversation: a new question starts again at
 * one, which is what lets a marker mean one thing for the whole answer beside
 * it, and lets the transcript list sources per answer rather than per session.
 */
export function createSourceList(): SourceList {
  const byUrl = new Map<string, Source>();
  let next = 1;

  return {
    add(title, url) {
      const address = url.trim();
      if (address.length === 0) return { index: 0 };

      const known = byUrl.get(address);
      if (known) return { index: known.index };

      const source: Source = { index: next, title: title.trim() || address, url: address };
      next += 1;
      byUrl.set(address, source);

      return { index: source.index, source };
    },
  };
}

/**
 * Every source a turn consulted, in number order.
 *
 * A record holds only what its own call added, so the turn's list is the
 * flattening of its records. That keeps the running total from being stored once
 * per later call, and it means the list can be rebuilt from a stored turn without
 * re-reading a tool result string and hoping its format held.
 */
export function collectSources(calls: ToolCallRecord[]): Source[] {
  const byIndex = new Map<number, Source>();

  for (const call of calls) {
    for (const source of call.sources ?? []) {
      // First wins. A number is assigned once, so a second record carrying it is
      // describing the same page, and the first one is what named it.
      if (!byIndex.has(source.index)) byIndex.set(source.index, source);
    }
  }

  return [...byIndex.values()].sort((left, right) => left.index - right.index);
}

/**
 * The numbers an answer cites that are actually sources.
 *
 * A plain scan, without the fence awareness the linking pass has, because its one
 * caller is the clipboard export: what a copied document needs is the pages the
 * answer leaned on, and a `[1]` inside a snippet there costs a reader nothing.
 */
export function citedIndices(answer: string, sources: Source[]): Set<number> {
  const known = new Set(sources.map((source) => source.index));
  const cited = new Set<number>();

  for (const match of answer.matchAll(/\[\^?(\d+(?:\s*,\s*\d+)*)\]/g)) {
    for (const part of match[1].split(",")) {
      const index = Number(part.trim());
      if (known.has(index)) cited.add(index);
    }
  }

  return cited;
}

/** What linking did to one run of prose. */
export interface CitationResult {
  text: string;
  /** The numbers that became links, which is what the list filters on. */
  cited: Set<number>;
}

/**
 * One citation marker, in the two shapes a model reaches for, plus any markers
 * written straight after it.
 *
 * Adjacent markers are matched as a single run, because a raised digit carries no
 * separator of its own: `[1][2]` drawn as `¹²` reads as twelve, or as twenty one, and
 * either way as one number rather than as two sources.
 *
 * The trailing guards are what keep this off markdown the model wrote on purpose:
 * `](` starts a link, and `]:` is a reference definition. The digit list allows
 * `[1, 3]`, which is how a sentence resting on two sources is often written.
 */
const CITATION_RUN = /\[\^?\d+(?:\s*,\s*\d+)*\](?![(:])(?:[ \t]*\[\^?\d+(?:\s*,\s*\d+)*\](?![(:]))*/g;

/** The numbers inside one run, in the order they were written. */
const RUN_NUMBERS = /\d+/g;

/** The digits a superscript marker is built from, in order. */
const DIGITS = "0123456789";
const SUPERSCRIPT = "⁰¹²³⁴⁵⁶⁷⁸⁹";

/**
 * A number as a superscript, which is what a citation mark is.
 *
 * Unicode rather than a `<sup>` tag on purpose. This is the label of a markdown
 * link, and a plain label cannot be broken by any part of the pipeline, while
 * inline HTML in a label asks two renderers to agree about something neither this
 * extension nor its tests can see. The glyphs are raised by design, so the mark
 * still reads as a superscript even when a font lacks the rarer ones and Qt falls
 * back for them.
 */
function superscript(index: number): string {
  let out = "";

  for (const digit of String(index)) {
    const at = DIGITS.indexOf(digit);
    out += at >= 0 ? SUPERSCRIPT[at] : digit;
  }

  return out;
}

/**
 * A footnote definition line, which is `[^1]: title, url` on a line of its own.
 *
 * The caret is required rather than optional, and the two shapes mean different
 * things. `[1]: something` is markdown's own reference definition syntax, which
 * backs a link the model wrote, so dropping one because it happens to carry a
 * source number would break that link. `[^1]:` is a footnote definition, which is
 * a duplicate of an entry we render ourselves.
 */
const FOOTNOTE_DEFINITION = /^\s*\[\^(\d+)\]:\s*\S/;

/**
 * A destination this window can safely be asked to open, or nothing.
 *
 * A search result is untrusted input, so the scheme is checked here for the same
 * reason `SAFE_URL` in `html.ts` checks it: a clickable `javascript:` link would
 * be a real vulnerability rather than a rendering quirk. The rest of the test is
 * characters that would break the link or smuggle a second one out of it.
 */
function linkableUrl(url: string): string | undefined {
  const address = url.trim();
  if (!/^https?:\/\//i.test(address)) return undefined;
  if (/[\s<>\\]/.test(address)) return undefined;

  // Control characters by code point. A character class would say the same thing
  // and trip `no-control-regex`, and a NUL or a DEL in a destination is never
  // anything but damage, so it is refused either way.
  for (const character of address) {
    const code = character.codePointAt(0) ?? 0;
    if (code < 0x20 || code === 0x7f) return undefined;
  }

  return address;
}

/**
 * Drops a definition line for a source the turn already lists.
 *
 * The entry is drawn from what the tool returned, so the model's copy of the same
 * thing is a duplicate line rather than information. One for a number we know
 * nothing about is left alone, because then it is the only place that text exists.
 */
function stripDefinitions(text: string, known: Set<number>): string {
  return text
    .split("\n")
    .filter((line) => {
      const definition = FOOTNOTE_DEFINITION.exec(line);
      return !definition || !known.has(Number(definition[1]));
    })
    .join("\n");
}

/**
 * `[1]`, `[^1]`, `[1, 3]` and `[1][2]` as superscript links, plus the numbers that
 * became links.
 *
 * One run of prose only: fences are the caller's business, because the caller
 * already walks the document in segments for the maths pass and there is no
 * reason for two ideas of where code starts.
 */
export function linkCitations(text: string, sources: Source[]): CitationResult {
  const cited = new Set<number>();
  if (sources.length === 0) return { text, cited };

  const byIndex = new Map(sources.map((source) => [source.index, source]));
  const prose = stripDefinitions(text, new Set(byIndex.keys()));

  const linked = mapOutsideCodeSpans(prose, (chunk) =>
    chunk.replace(CITATION_RUN, (whole: string) => {
      const wanted = [...whole.matchAll(RUN_NUMBERS)].map((match) => Number(match[0]));
      const links: string[] = [];

      for (const index of wanted) {
        const source = byIndex.get(index);
        const href = source ? linkableUrl(source.url) : undefined;
        // All or nothing, and nothing at all if a destination is unusable: a
        // marker that is partly a link reads as a citation that works.
        if (!href) return whole;

        links.push(markdownLink(superscript(index), href));
      }

      for (const index of wanted) cited.add(index);
      // Separated even when the model wrote them touching: `¹ ²` is two citations,
      // where `¹²` is a number, and the reader has no way to tell it from a real one.
      // A space rather than a comma, which sits at the baseline and so drops away from
      // digits that are raised above it.
      return links.join(" ");
    }),
  );

  return { text: linked, cited };
}

/**
 * The sources under an answer, or nothing when there is nothing to say.
 *
 * The marker is the link, matching the inline marker character for character, so
 * one token means one thing in both places. Numbers are the model's rather than
 * the reader's: `[3]` is the third source the turn found whether or not the answer
 * cited the two before it, which keeps the transcript agreeing with the tool
 * output and with anything the model said about its own numbering.
 */
export function renderSources(sources: Source[], cited: Set<number>, streaming: boolean): string {
  // While an answer is arriving the cited set is still filling up, so the list
  // waits rather than appearing in full and then dropping most of itself. An
  // answer that cites nothing at all is the one case left to fall back on
  // everything consulted, because there is then nothing to filter by.
  const shown =
    cited.size > 0 ? sources.filter((source) => cited.has(source.index)) : streaming ? [] : sources;
  if (shown.length === 0) return "";

  const entries = shown.map((source) => {
    const href = linkableUrl(source.url);
    const marker = href ? markdownLink(`[${source.index}]`, href) : `[${source.index}]`;
    return `- ${marker} ${source.title}`;
  });

  return ["_Sources_", "", ...entries].join("\n");
}