import { Marked } from 'marked';
import markedKatex from 'marked-katex-extension';
import { markedHighlight } from 'marked-highlight';
import { gfmHeadingId } from 'marked-gfm-heading-id';
import hljs from 'highlight.js/lib/core';
import bash from 'highlight.js/lib/languages/bash';
import css from 'highlight.js/lib/languages/css';
import diff from 'highlight.js/lib/languages/diff';
import go from 'highlight.js/lib/languages/go';
import javascript from 'highlight.js/lib/languages/javascript';
import json from 'highlight.js/lib/languages/json';
import markdown from 'highlight.js/lib/languages/markdown';
import python from 'highlight.js/lib/languages/python';
import rust from 'highlight.js/lib/languages/rust';
import shell from 'highlight.js/lib/languages/shell';
import sql from 'highlight.js/lib/languages/sql';
import typescript from 'highlight.js/lib/languages/typescript';
import xml from 'highlight.js/lib/languages/xml';
import yaml from 'highlight.js/lib/languages/yaml';
import DOMPurify from 'isomorphic-dompurify';

hljs.registerLanguage('bash', bash);
hljs.registerLanguage('css', css);
hljs.registerLanguage('diff', diff);
hljs.registerLanguage('go', go);
hljs.registerLanguage('javascript', javascript);
hljs.registerLanguage('json', json);
hljs.registerLanguage('markdown', markdown);
hljs.registerLanguage('python', python);
hljs.registerLanguage('rust', rust);
hljs.registerLanguage('shell', shell);
hljs.registerLanguage('sql', sql);
hljs.registerLanguage('typescript', typescript);
hljs.registerLanguage('xml', xml);
hljs.registerLanguage('yaml', yaml);

function escapeHtml(s: string): string {
  return s
    .replace(/&/g, '&')
    .replace(/</g, '&lt;')
    .replace(/>/g, '&gt;')
    .replace(/"/g, '&quot;')
    .replace(/'/g, '&#39;');
}

const renderer = new Marked()
  .setOptions({ gfm: true, breaks: false })
  .use(gfmHeadingId())
  .use(
    markedHighlight({
      langPrefix: 'hljs language-',
      highlight(code, lang) {
        const language = lang && hljs.getLanguage(lang) ? lang : 'plaintext';
        return hljs.highlight(code, { language }).value;
      },
    }),
  )
  .use(markedKatex({ throwOnError: false, nonStandard: true }))
  .use({
    renderer: {
      image({ href, title, text }: { href: string; title: string | null; text: string }) {
        const safeHref = escapeHtml(href);
        const safeAlt = escapeHtml(text || '');
        const safeTitle = title ? escapeHtml(title) : '';
        const caption = title || text || '';
        const figcaption = caption.trim()
          ? `<figcaption>${escapeHtml(caption)}</figcaption>`
          : '';
        const titleAttr = safeTitle ? ` title="${safeTitle}"` : '';
        return `<figure><img src="${safeHref}" alt="${safeAlt}"${titleAttr} loading="lazy" />${figcaption}</figure>`;
      },
    },
  });

const KATEX_TAGS = [
  'math', 'annotation', 'semantics', 'mrow', 'mi', 'mo', 'mn', 'mtext',
  'msup', 'msub', 'msubsup', 'mfrac', 'msqrt', 'mroot', 'mover', 'munder',
  'munderover', 'mtable', 'mtr', 'mtd', 'mspace', 'mstyle', 'mphantom',
  'mpadded', 'menclose',
];

export interface ImageDim { w: number; h: number }
export type ImageDims = Record<string, ImageDim>;

function injectDimensions(html: string, dims?: ImageDims): string {
  if (!dims) return html;
  return html.replace(/<img\s+src="([^"]+)"([^>]*?)\s*\/?>/g, (match, src, rest) => {
    const d = dims[src];
    if (!d) return match;
    if (/\swidth\s*=/.test(rest) || /\sheight\s*=/.test(rest)) return match;
    return `<img src="${src}" width="${d.w}" height="${d.h}"${rest} />`;
  });
}

// Marked emits `<p><figure>…</figure></p>` because the image renderer returns
// a block element from an inline slot. Strip the `<p>` wrapper when the
// paragraph contains nothing but figures (and whitespace / <br>), so the
// figure-grouping step below can see contiguous runs.
function unwrapFiguresFromParagraphs(html: string): string {
  return html.replace(/<p>([\s\S]*?)<\/p>/g, (match, inner: string) => {
    const stripped = inner.replace(/<br\s*\/?>/g, '').trim();
    if (!stripped) return match;
    if (/^(?:\s*<figure>[\s\S]*?<\/figure>\s*)+$/.test(stripped)) {
      return stripped;
    }
    return match;
  });
}

// Wrap runs of 2+ consecutive <figure> elements in a `.figure-row` flex
// container. Each figure gets `flex: <aspect-ratio>` so widths divide the
// row proportionally and the final heights match.
function groupFigures(html: string): string {
  return html.replace(
    /(?:<figure>[\s\S]*?<\/figure>\s*){2,}/g,
    (run) => {
      const figures = run.match(/<figure>[\s\S]*?<\/figure>/g) ?? [];
      const items = figures.map((fig) => {
        const m = fig.match(/<img[^>]*\swidth="(\d+)"[^>]*\sheight="(\d+)"/);
        const ratio = m ? Number(m[1]) / Number(m[2]) : 1;
        const safe = Number.isFinite(ratio) && ratio > 0 ? ratio : 1;
        const r = safe.toFixed(3);
        // flex-basis tracks target row height (var --row-h) × aspect ratio,
        // so the browser fits as many figures per row as it can while keeping
        // each at roughly the target height; the rest wrap. max-width caps
        // any leftover row so a lone last image doesn't balloon.
        const style = [
          `flex:${r} ${r} calc(${r} * var(--row-h, 16rem))`,
          `max-width:calc(${r} * var(--row-max, 30rem))`,
        ].join(';');
        return fig.replace('<figure>', `<figure style="${style}">`);
      });
      return `<div class="figure-row">${items.join('')}</div>`;
    },
  );
}

export function renderMarkdown(src: string, dims?: ImageDims): string {
  let html = renderer.parse(src, { async: false }) as string;
  html = injectDimensions(html, dims);
  html = unwrapFiguresFromParagraphs(html);
  html = groupFigures(html);
  return DOMPurify.sanitize(html, {
    ADD_TAGS: [...KATEX_TAGS, 'figure', 'figcaption'],
    ADD_ATTR: ['aria-hidden', 'style', 'id', 'class', 'encoding', 'mathvariant', 'displaystyle', 'scriptlevel', 'loading'],
  });
}