const MERMAID_STARTERS = [
  "flowchart ",
  "graph ",
  "sequenceDiagram",
  "classDiagram",
  "erDiagram",
  "stateDiagram",
  "journey",
  "gantt",
  "pie ",
  "mindmap",
  "timeline",
  "gitGraph",
];

let mermaidInitialized = false;

function looksLikeMermaidParagraph(paragraph: string): boolean {
  const trimmed = paragraph.trim();
  return MERMAID_STARTERS.some((starter) => trimmed.startsWith(starter));
}

function normalizeCandidateMermaid(paragraph: string): string {
  return paragraph
    .trim()
    .replace(/\\n/g, "\n")
    .replace(/\s+→\s+/g, " --> ")
    .replace(/\r\n?/g, "\n")
    .trim();
}

async function getMermaid() {
  const { default: mermaid } = await import("mermaid");
  if (!mermaidInitialized) {
    mermaid.initialize({
      startOnLoad: false,
      suppressErrorRendering: true,
      securityLevel: "loose",
    });
    mermaidInitialized = true;
  }
  return mermaid;
}

async function validateMermaid(chart: string): Promise<boolean> {
  try {
    const mermaid = await getMermaid();
    await mermaid.parse(chart);
    return true;
  } catch {
    return false;
  }
}

function diagramFallback(chart: string): string {
  return [
    "> Diagram omitted because the generated Mermaid was invalid.",
    "",
    "```text",
    chart,
    "```",
  ].join("\n");
}

export async function normalizeAssistantMarkdownForRender(markdown: string): Promise<string> {
  const paragraphs = markdown.split(/\n\s*\n/);
  let changed = false;

  const normalized = await Promise.all(
    paragraphs.map(async (paragraph) => {
      if (!looksLikeMermaidParagraph(paragraph)) {
        return paragraph;
      }

      const candidate = normalizeCandidateMermaid(paragraph);
      const valid = await validateMermaid(candidate);
      changed = true;
      if (valid) {
        return `\`\`\`mermaid\n${candidate}\n\`\`\``;
      }
      return diagramFallback(candidate);
    }),
  );

  return changed ? normalized.join("\n\n") : markdown;
}