learning-platform/app/services/ingestion/src/pipeline/chunk.ts

import { v4 as uuid } from 'uuid';
import type { Chunk, DocumentFormat } from '../types.js';

const MD_MIN = 100;
const MD_MAX = 1500;
const TXT_WINDOW = 800;
const TXT_OVERLAP = 150;
const PDF_MIN = 100;
const PDF_MAX = 1200;

// ---------------------------------------------------------------------------
// MD chunking — heading-based
// ---------------------------------------------------------------------------

interface HeadingSection {
  level: number;
  heading: string;
  parent: string | null;
  content: string;
}

function parseMdSections(text: string): HeadingSection[] {
  const lines = text.split('\n');
  const sections: HeadingSection[] = [];
  let current: HeadingSection | null = null;
  const parentStack: { level: number; heading: string }[] = [];

  for (const line of lines) {
    const h3 = line.match(/^### (.+)/);
    const h2 = line.match(/^## (.+)/);
    const h1 = line.match(/^# (.+)/);
    const match = h3 ?? h2 ?? h1;

    if (match) {
      if (current) sections.push(current);
      const level = h1 ? 1 : h2 ? 2 : 3;
      const heading = match[1] ?? line;

      // Maintain parent stack
      while (parentStack.length > 0 && (parentStack[parentStack.length - 1]?.level ?? 0) >= level) {
        parentStack.pop();
      }
      const parent = parentStack[parentStack.length - 1]?.heading ?? null;
      parentStack.push({ level, heading });

      current = { level, heading, parent, content: '' };
    } else if (current) {
      current.content += line + '\n';
    }
  }
  if (current) sections.push(current);
  return sections;
}

function splitOnParagraphs(text: string, maxSize: number): string[] {
  const paragraphs = text.split(/\n\n+/);
  const parts: string[] = [];
  let current = '';

  for (const para of paragraphs) {
    if ((current + para).length > maxSize && current.length > 0) {
      parts.push(current.trim());
      current = para;
    } else {
      current = current ? current + '\n\n' + para : para;
    }
  }
  if (current.trim()) parts.push(current.trim());
  return parts;
}

function chunkMd(text: string, documentId: string): Chunk[] {
  const sections = parseMdSections(text);
  const merged: HeadingSection[] = [];

  for (let i = 0; i < sections.length; i++) {
    const sec = sections[i];
    if (sec === undefined) continue;
    const fullText = `${'#'.repeat(sec.level)} ${sec.heading}\n\n${sec.content}`.trim();

    if (fullText.length < MD_MIN && i + 1 < sections.length) {
      // Merge with next sibling by appending to next's content prefix
      const next = sections[i + 1];
      if (next !== undefined) {
        next.content = fullText + '\n\n' + next.content;
        continue;
      }
    }
    merged.push(sec);
  }

  const chunks: Chunk[] = [];
  let globalIndex = 0;

  for (const sec of merged) {
    const fullText = `${'#'.repeat(sec.level)} ${sec.heading}\n\n${sec.content}`.trim();
    const parts = fullText.length > MD_MAX ? splitOnParagraphs(fullText, MD_MAX) : [fullText];

    for (const part of parts) {
      if (part.length < 1) continue;
      chunks.push({
        id: uuid(),
        documentId,
        text: part,
        format: 'md',
        index: globalIndex++,
        metadata: {
          headingLevel: sec.level,
          headingText: sec.heading,
          parentHeading: sec.parent ?? undefined,
        },
      });
    }
  }

  return chunks;
}

// ---------------------------------------------------------------------------
// TXT chunking — sliding window
// ---------------------------------------------------------------------------

function splitSentences(text: string): string[] {
  return text.match(/[^.!?]+[.!?]+\s*/g) ?? [text];
}

function chunkTxt(text: string, documentId: string): Chunk[] {
  const paragraphs = text.split(/\n\n+/).filter(p => p.trim().length > 0);
  const chunks: Chunk[] = [];
  let current = '';
  let index = 0;
  const total = text.length;

  const pushChunk = (t: string) => {
    const pos = index * TXT_WINDOW;
    chunks.push({
      id: uuid(),
      documentId,
      text: t.trim(),
      format: 'txt',
      index: index++,
      metadata: {
        approximatePosition:
          pos < total * 0.25 ? 'start' : pos > total * 0.75 ? 'end' : 'middle',
      },
    });
  };

  for (const para of paragraphs) {
    if ((current + ' ' + para).trim().length <= TXT_WINDOW) {
      current = (current + ' ' + para).trim();
    } else {
      if (current.length >= MD_MIN) pushChunk(current);
      // Para may itself exceed window — split on sentences
      if (para.length > TXT_WINDOW) {
        const sentences = splitSentences(para);
        let buf = '';
        for (const sent of sentences) {
          if ((buf + sent).length > TXT_WINDOW && buf.length >= MD_MIN) {
            pushChunk(buf);
            // Keep overlap
            const words = buf.split(' ');
            buf = words.slice(-Math.floor(TXT_OVERLAP / 5)).join(' ') + ' ' + sent;
          } else {
            buf += sent;
          }
        }
        current = buf;
      } else {
        current = para;
      }
    }
  }
  if (current.trim().length >= MD_MIN) pushChunk(current);

  return chunks;
}

// ---------------------------------------------------------------------------
// PDF chunking — page + paragraph
// ---------------------------------------------------------------------------

function chunkPdf(text: string, documentId: string): Chunk[] {
  const pages = text.split('---PAGE---');
  const chunks: Chunk[] = [];
  let globalIndex = 0;

  for (let pageIdx = 0; pageIdx < pages.length; pageIdx++) {
    const page = pages[pageIdx];
    if (page === undefined || !page.trim()) continue;

    const paragraphs = page.split(/\n\n+/).filter(p => p.trim().length > 0);
    let chunkOnPage = 0;
    let accumulator = '';

    const flushAccumulator = () => {
      const t = accumulator.trim();
      if (t.length >= PDF_MIN) {
        chunks.push({
          id: uuid(),
          documentId,
          text: t,
          format: 'pdf',
          index: globalIndex++,
          metadata: { pageNumber: pageIdx + 1, chunkIndexOnPage: chunkOnPage++ },
        });
      }
      accumulator = '';
    };

    for (const para of paragraphs) {
      if ((accumulator + '\n\n' + para).trim().length > PDF_MAX) {
        flushAccumulator();
        // Para may exceed max on its own — hard split at sentence boundary
        if (para.length > PDF_MAX) {
          const sentences = splitSentences(para);
          for (const sent of sentences) {
            if ((accumulator + sent).length > PDF_MAX && accumulator.length >= PDF_MIN) {
              flushAccumulator();
            }
            accumulator += sent;
          }
        } else {
          accumulator = para;
        }
      } else {
        accumulator = accumulator ? accumulator + '\n\n' + para : para;
      }
    }
    flushAccumulator();
  }

  return chunks;
}

// ---------------------------------------------------------------------------
// Public entry
// ---------------------------------------------------------------------------

export function chunk(text: string, format: DocumentFormat, documentId: string): Chunk[] {
  switch (format) {
    case 'md':  return chunkMd(text, documentId);
    case 'txt': return chunkTxt(text, documentId);
    case 'pdf': return chunkPdf(text, documentId);
  }
}