import remend, { type RemendHandler } from 'remend'

/**
 * Streaming repair for partial Markdown, applied only while a message is
 * still arriving. `remend` closes unterminated emphasis, code, links and
 * fences; the table handler below makes a GFM table render as a table as soon
 * as its header row arrives instead of flashing raw pipes until the
 * `|---|` delimiter row is complete.
 */

const TABLE_ROW = /^\s*\|.*\|\s*$/
const PARTIAL_DELIMITER = /^\s*\|?[\s:|-]*$/

function countColumns(row: string): number {
  return row.trim().replace(/^\||\|$/g, '').split('|').length
}

function delimiterRow(columns: number): string {
  return `|${' --- |'.repeat(columns)}`
}

function isInsideFence(lines: string[]): boolean {
  return lines.filter(line => /^\s*(```|~~~)/.test(line)).length % 2 === 1
}

export function completeStreamingTable(text: string): string {
  const lines = text.split('\n')
  const last = lines[lines.length - 1]
  const previous = lines[lines.length - 2]
  const beforeHeader = lines[lines.length - 3]
  if (isInsideFence(lines)) return text

  // Header row just finished: `| A | B |` → add the delimiter row.
  if (TABLE_ROW.test(last) && (previous === undefined || !TABLE_ROW.test(previous))) {
    return `${text}\n${delimiterRow(countColumns(last))}`
  }

  // Delimiter row still streaming: `| A | B |\n|--` → finish it.
  if (
    previous !== undefined
    && TABLE_ROW.test(previous)
    && (beforeHeader === undefined || !TABLE_ROW.test(beforeHeader))
    && last.trim().length > 0
    && PARTIAL_DELIMITER.test(last)
  ) {
    return [...lines.slice(0, -1), delimiterRow(countColumns(previous))].join('\n')
  }

  return text
}

const tableHandler: RemendHandler = {
  name: 'gfm-table',
  handle: completeStreamingTable,
}

export function repairStreamingMarkdown(text: string): string {
  return remend(text, { linkMode: 'text-only', handlers: [tableHandler] })
}