markdown-utils.ts

import TurndownService from 'turndown';

export type MarkdownSection = {
  header: string;
  depth: number;
  content: string;
  children?: MarkdownSection[];
};

export type SectionMatch = {
  path: string[];
  section: MarkdownSection;
  content: string;
};

class NotFoundException extends Error {}

export const getTurndownService = (): TurndownService => {
  const turndown = new TurndownService({
    codeBlockStyle: 'fenced',
    headingStyle: 'atx',
  });

  // Remove non-content elements (defense in depth — Readability strips most of these)
  turndown.remove(['script', 'style', 'noscript', 'iframe']);

  return turndown;
};

/**
 * Parses markdown into a nested tree of sections based on headings.
 * Text before the first heading is assigned the header "Introduction".
 */
export const parseMarkdownSections = (markdown: string): MarkdownSection[] => {
  const lines = markdown.split('\n');
  const headingRegex = /^(#{1,6})\s+(.+)$/;

  // Split into flat segments: { depth, header, contentLines }
  const segments: { depth: number; header: string; contentLines: string[] }[] =
    [];
  let current: (typeof segments)[number] | null = null;

  for (const line of lines) {
    const match = line.match(headingRegex);
    if (match) {
      if (current) segments.push(current);
      current = {
        depth: match[1].length,
        header: match[2].trim(),
        contentLines: [],
      };
    } else {
      if (!current) {
        current = { depth: 2, header: 'Introduction', contentLines: [] };
      }
      current.contentLines.push(line);
    }
  }
  if (current) segments.push(current);

  if (segments.length === 0) return [];

  // Normalize depths so the shallowest heading becomes depth 1
  const minDepth = Math.min(...segments.map((s) => s.depth));
  for (const seg of segments) {
    seg.depth = seg.depth - minDepth + 1;
  }

  // Build nested tree using a stack
  const root: MarkdownSection[] = [];
  const stack: { section: MarkdownSection; depth: number }[] = [];

  for (const seg of segments) {
    const section: MarkdownSection = {
      header: seg.header,
      depth: seg.depth,
      content: seg.contentLines.join('\n').trim(),
    };

    // Pop stack until we find a parent with shallower depth
    while (stack.length > 0 && stack[stack.length - 1].depth >= seg.depth) {
      stack.pop();
    }

    if (stack.length === 0) {
      root.push(section);
    } else {
      const parent = stack[stack.length - 1].section;
      if (!parent.children) parent.children = [];
      parent.children.push(section);
    }

    stack.push({ section, depth: seg.depth });
  }

  return root;
};

/**
 * Finds all sections matching a header and returns each with its ancestor path
 * and reconstructed markdown content (including nested children).
 * Throws NotFoundException if no matches are found.
 */
export const getSection = (
  header: string,
  sections: MarkdownSection[],
): SectionMatch[] => {
  const matches: SectionMatch[] = [];
  findSections(header, sections, [], matches);

  if (matches.length === 0) {
    throw new NotFoundException(`Section "${header}" not found`);
  }

  return matches;
};

const findSections = (
  header: string,
  sections: MarkdownSection[],
  ancestors: string[],
  results: SectionMatch[],
): void => {
  for (const section of sections) {
    if (section.header === header) {
      results.push({
        path: [...ancestors, section.header],
        section,
        content: sectionToMarkdown(section),
      });
    }
    if (section.children) {
      findSections(
        header,
        section.children,
        [...ancestors, section.header],
        results,
      );
    }
  }
};

/**
 * Builds a condensed digest: each section's heading followed by its first paragraph.
 * Produces a much smaller input for LLM summarization while preserving key context.
 */
export const buildSectionDigest = (sections: MarkdownSection[]): string => {
  const parts: string[] = [];

  const walk = (items: MarkdownSection[]) => {
    for (const item of items) {
      const heading = '#'.repeat(item.depth);
      const firstParagraph = extractFirstParagraph(item.content);
      if (firstParagraph) {
        parts.push(`${heading} ${item.header}\n${firstParagraph}`);
      } else {
        parts.push(`${heading} ${item.header}`);
      }
      if (item.children) walk(item.children);
    }
  };

  walk(sections);
  return parts.join('\n\n');
};

const extractFirstParagraph = (content: string): string => {
  if (!content) return '';
  // Split on double newlines (paragraph breaks) and take the first non-empty block
  const paragraphs = content.split(/\n\n+/);
  for (const p of paragraphs) {
    const trimmed = p.trim();
    if (trimmed) return trimmed;
  }
  return '';
};

const sectionToMarkdown = (section: MarkdownSection): string => {
  const heading = '#'.repeat(section.depth);
  const parts: string[] = [`${heading} ${section.header}`];

  if (section.content) parts.push(section.content);

  if (section.children) {
    for (const child of section.children) {
      parts.push(sectionToMarkdown(child));
    }
  }

  return parts.join('\n\n');
};