display.ts

import { markdownTable } from "markdown-table"

import { Chat, type ChatMessage } from "./types.ts"
import { models, systemBase } from "./models.ts"
import { renderMarkdown } from "./md-render.ts"

export async function renderMd(md: string, raw = false) {
  if (Deno.stdout.isTerminal() && !raw) {
    console.log(await renderMarkdown(md))
  } else {
    console.log(md)
  }
}

export const codeBlock = (contents: string, lang = "") =>
  `\`\`\`${lang}\n${contents}\n\`\`\`\n`
export const jsonBlock = (obj: unknown) => codeBlock(JSON.stringify(obj, null, 2), "json")

const codeMd = (s: string) => `\`${s}\``
export const codeListMd = (strs: string[]) => strs.map(codeMd).join(", ")

const moneyFmt = Intl.NumberFormat("en-US", {
  style: "currency",
  currency: "USD",
  minimumFractionDigits: 0,
  maximumFractionDigits: 5,
})

const modelsTable = (verbose: boolean) =>
  markdownTable([
    ["Provider", "ID", ...(verbose ? ["Model key"] : []), "Input", "Cached", "Output"],
    ...models.map((m) => [
      m.provider,
      m.id + (m.default ? " ⭐" : ""),
      ...(verbose ? [m.key.replace(/^meta-llama\//, "")] : []),
      moneyFmt.format(m.input),
      m.input_cached ? moneyFmt.format(m.input_cached) : "",
      moneyFmt.format(m.output),
    ]),
  ])

/** Split, add `"> "` to the beginning of each line, and rejoin */
const quote = (s: string) => s.split("\n").map((line) => "> " + line).join("\n")

export const modelsMd = (verbose = false) =>
  `Models are matched on ID or key. Prices are per million tokens.\n\n${
    modelsTable(verbose)
  }`

// Collapse long user inputs in gists, unless the input already contains a
// <details> block (e.g., from `cb`), in which case the user has already
// decided what to hide and we leave it alone rather than nesting.
const COLLAPSE_LINES = 40
// only match beginning of line to avoid mentions in prose
const hasDetails = (s: string) => /^<details[\s>]/m.test(s)

// split from message content because we only want this in show or gist mode
function messageHeaderMd(msg: ChatMessage, msgNum: number, msgCount: number) {
  return `# ${msg.role} (${msgNum}/${msgCount})\n\n`
}

const timeFmt = new Intl.NumberFormat(undefined, { maximumFractionDigits: 2 })

// turn time into `1m20s` string
export function formatElapsed(ms: number) {
  const totalSeconds = ms / 1000
  const minutes = Math.floor(totalSeconds / 60)
  const seconds = totalSeconds % 60
  if (minutes === 0) return `${timeFmt.format(seconds)}s`
  return `${minutes}m${Math.floor(seconds)}s`
}

const escapeThinkTags = (content: string) =>
  content
    .replace("<think>", "\\<think>")
    .replace("</think>", "\\</think>")

// TODO: add `verbose` mode and hide reasoning from `nice` mode

/**
 * - `cli` is the default show output, includes meta but not reasoning
 * - `verbose` is like `cli` + reasoning
 * - `raw` is for insertion in, e.g., a text editor
 * - `gist` includes meta but collapses reasoning under `<details>`
 */
export type DisplayMode = "cli" | "verbose" | "raw" | "gist"

export function messageContentMd(msg: ChatMessage, mode: DisplayMode) {
  let output = ""

  if (msg.role === "assistant") {
    // show metadata line in all modes except raw
    if (mode !== "raw") {
      // only show stop reason if it's not a natural stop
      const showStopReason = !["stop", "end_turn", "completed"].includes(
        msg.stop_reason.toLowerCase(),
      )

      output += codeMd(msg.model)
      output += ` | ${formatElapsed(msg.timeMs)}`
      output += ` | ${moneyFmt.format(msg.cost)}`

      // show cached tokens in parens if there are any

      const input = msg.tokens.input +
        (msg.tokens.input_cache_hit ? ` (${msg.tokens.input_cache_hit})` : "")
      output += ` | ${input} -> ${msg.tokens.output}`

      if (msg.searches) {
        output += ` | 🌐`
        if (msg.searches > 1) output += `×${msg.searches}`
      }

      if (showStopReason) output += ` | **Stop reason:** ${msg.stop_reason}`

      output += "\n\n"
    }

    // only show reasoning in gist or verbose mode
    if (msg.reasoning) {
      if (mode === "gist") {
        output += tag("details", tag("summary", "Reasoning"), quote(msg.reasoning))
      } else if (mode === "verbose") {
        output += quote(msg.reasoning) + "\n\n"
      }
    }
  }

  if (msg.role === "user" && msg.outputSchema && mode !== "raw") {
    output += `**Schema:** \`${msg.outputSchema}\`\n\n`
  }
  // For long user inputs in gist mode, collapse in a details block
  if (
    msg.role === "user" && mode === "gist" &&
    !hasDetails(msg.content) &&
    msg.content.split("\n").length > COLLAPSE_LINES
  ) {
    output += tag(
      "details",
      tag("summary", "Long input collapsed"),
      escapeThinkTags(msg.content),
    )
  } else {
    output += mode === "raw" ? msg.content : escapeThinkTags(msg.content)
  }
  if (msg.role === "user" && msg.image_url) {
    output += msg.image_url.startsWith("data:")
      ? `\n\n[Image (inline)]`
      : `\n\n[Image](${msg.image_url})`
  }
  if (mode !== "raw") output += "\n\n"
  return output
}

type ChatToMd = { chat: Chat; lastN?: number; indices?: number[]; mode?: DisplayMode }

const tag = (t: string, ...children: string[]) =>
  `<${t}>\n${children.join("\n")}\n</${t}>\n\n`

export function chatToMd({ chat, lastN = 0, indices, mode = "cli" }: ChatToMd): string {
  const msgCount = chat.messages.length

  // Determine which messages to include
  let selectedIndices: number[]
  if (indices && indices.length > 0) {
    selectedIndices = indices
  } else if (lastN > 0) {
    const cappedLastN = Math.min(lastN, msgCount)
    selectedIndices = Array.from(
      { length: cappedLastN },
      (_, i) => msgCount - cappedLastN + i,
    )
  } else {
    selectedIndices = Array.from({ length: msgCount }, (_, i) => i)
  }

  const messages = selectedIndices.map((i) => ({ msg: chat.messages[i], idx: i }))

  if (mode === "raw") {
    return messages.map(({ msg }) => messageContentMd(msg, "raw")).join("\n\n")
  }

  let output = `**Chat started:** ${longDateFmt.format(chat.createdAt)}\n\n`

  if (mode === "gist") {
    // always print system prompt in gist mode, but collapse it
    output += tag("details", tag("summary", "System prompt"), chat.systemPrompt)
  } else if (chat.systemPrompt !== systemBase) {
    // otherwise only print system prompt if it's non-default
    output += `**System prompt:** ${chat.systemPrompt}\n\n`
  }

  messages.forEach(({ msg, idx }) => {
    output += messageHeaderMd(msg, idx + 1, msgCount)
    output += messageContentMd(msg, mode)
  })
  return output
}

export const longDateFmt = new Intl.DateTimeFormat("en-US", {
  dateStyle: "medium",
  timeStyle: "short",
})

export const shortDateFmt = new Intl.DateTimeFormat("en-US", {
  month: "numeric",
  day: "numeric",
  hour: "2-digit",
  minute: "2-digit",
  hour12: false,
})