freezmo/lib/documents.ts

"use client";

/**
 * Client-side extraction for the note formats people actually use:
 *
 *   .pdf                      text layer, or rendered pages when scanned
 *   .docx                     text + embedded images
 *   .pptx                     slide text, speaker notes + embedded images
 *   .odt / .odp / .ods        LibreOffice documents (content.xml)
 *   .txt .md .csv .rtf-lite   plain text
 *
 * Everything runs in the browser: the original file never leaves the device,
 * only the extracted text (and any page images for the vision model) is sent
 * to the API routes. Photos uploaded as images keep working exactly as before.
 */

import { t } from "./i18n";
import { prepareImageFromBlob, type PreparedImage } from "./image";

export const MAX_DOCUMENT_BYTES = 40 * 1024 * 1024; // 40 MB per file
/** How many rendered page images one document may contribute. */
export const MAX_DOCUMENT_IMAGES = 10;
/** Characters kept per document (long books get truncated). */
export const MAX_DOCUMENT_CHARS = 60_000;
/** Pages below this much text are treated as scanned and rendered as images. */
const SCANNED_PAGE_CHAR_THRESHOLD = 60;
const PDF_RENDER_SCALE = 2;

export const ACCEPTED_DOCUMENT_EXTENSIONS = [
  "pdf",
  "docx",
  "pptx",
  "odt",
  "odp",
  "ods",
  "txt",
  "md",
  "markdown",
  "csv",
  "tsv",
  "log",
] as const;

export const DOCUMENT_ACCEPT_ATTR = [
  ...ACCEPTED_DOCUMENT_EXTENSIONS.map((extension) => `.${extension}`),
  "application/pdf",
  "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
  "application/vnd.openxmlformats-officedocument.presentationml.presentation",
  "application/vnd.oasis.opendocument.text",
  "application/vnd.oasis.opendocument.presentation",
  "application/vnd.oasis.opendocument.spreadsheet",
  "text/plain",
  "text/markdown",
  "text/csv",
].join(",");

export type DocumentKind = "pdf" | "docx" | "pptx" | "opendocument" | "text";

export type ExtractedDocument = {
  id: string;
  name: string;
  kind: DocumentKind;
  /** Extracted text, already truncated to MAX_DOCUMENT_CHARS. */
  text: string;
  /** Page-rendered or embedded images (JPEG data URLs, ready to send). */
  images: PreparedImage[];
  words: number;
  characters: number;
  pages?: number;
  truncated: boolean;
  imagesTruncated: boolean;
  warnings: string[];
};

export type DocumentProgress = {
  /** Human readable stage, e.g. "Reading page 4 of 12". */
  message: string;
  /** 0..1 where known. */
  progress?: number;
};

function extensionOf(name: string): string {
  const index = name.lastIndexOf(".");
  return index === -1 ? "" : name.slice(index + 1).toLowerCase();
}

export function documentKindOf(file: File): DocumentKind | null {
  const extension = extensionOf(file.name);
  switch (extension) {
    case "pdf":
      return "pdf";
    case "docx":
      return "docx";
    case "pptx":
      return "pptx";
    case "odt":
    case "odp":
    case "ods":
      return "opendocument";
    case "txt":
    case "md":
    case "markdown":
    case "csv":
    case "tsv":
    case "log":
      return "text";
    default:
      return null;
  }
}

export function isDocumentFile(file: File): boolean {
  return documentKindOf(file) !== null;
}

export function describeDocumentError(file: File): string | null {
  if (!isDocumentFile(file)) {
    return t("file.notDocument", { name: file.name });
  }
  if (file.size > MAX_DOCUMENT_BYTES) {
    return t("file.documentTooLarge", { name: file.name });
  }
  return null;
}

export const KIND_LABELS: Record<DocumentKind, string> = {
  pdf: "PDF",
  docx: "Word",
  pptx: "PowerPoint",
  opendocument: "OpenDocument",
  text: "Text",
};

/* ------------------------------------------------------------------ */
/* Shared helpers                                                      */
/* ------------------------------------------------------------------ */

function cleanText(raw: string): string {
  return raw
    .replace(/\r\n?/g, "\n")
    .replace(/\u00a0/g, " ")
    .replace(/[ \t]+\n/g, "\n")
    .replace(/\n{3,}/g, "\n\n")
    .replace(/[ \t]{2,}/g, " ")
    .trim();
}

function countWords(text: string): number {
  return text.split(/\s+/).filter(Boolean).length;
}

function withLimits(
  name: string,
  kind: DocumentKind,
  text: string,
  images: PreparedImage[],
  extra: { pages?: number; warnings?: string[]; imagesTruncated?: boolean },
): ExtractedDocument {
  const cleaned = cleanText(text);
  const truncated = cleaned.length > MAX_DOCUMENT_CHARS;
  const finalText = truncated ? cleaned.slice(0, MAX_DOCUMENT_CHARS) : cleaned;
  return {
    id: `${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 8)}`,
    name,
    kind,
    text: finalText,
    images,
    words: countWords(finalText),
    characters: finalText.length,
    pages: extra.pages,
    truncated,
    imagesTruncated: extra.imagesTruncated ?? false,
    warnings: extra.warnings ?? [],
  };
}

/* ------------------------------------------------------------------ */
/* PDF                                                                 */
/* ------------------------------------------------------------------ */

type PdfPage = {
  getTextContent: () => Promise<{ items: Array<{ str?: string; transform?: number[] }> }>;
  getViewport: (options: { scale: number }) => { width: number; height: number };
  render: (options: { canvasContext: CanvasRenderingContext2D; viewport: unknown }) => { promise: Promise<void> };
  cleanup?: () => void;
};

/** Join pdf.js text items, inserting line breaks when the baseline jumps. */
function textFromPdfItems(items: Array<{ str?: string; transform?: number[] }>): string {
  let output = "";
  let lastY: number | null = null;
  for (const item of items) {
    const value = item.str ?? "";
    const y = item.transform?.[5];
    if (lastY !== null && typeof y === "number" && Math.abs(y - lastY) > 2) {
      output += "\n";
    }
    output += value;
    if (value && !value.endsWith(" ")) output += " ";
    if (typeof y === "number") lastY = y;
  }
  return output;
}

async function extractPdf(
  file: File,
  onProgress?: (update: DocumentProgress) => void,
): Promise<ExtractedDocument> {
  // Loaded lazily so pdf.js (and its worker) never hits the initial bundle.
  const pdfjs = await import("pdfjs-dist");
  pdfjs.GlobalWorkerOptions.workerSrc = "/pdf.worker.min.mjs";

  const data = new Uint8Array(await file.arrayBuffer());
  const document = await pdfjs.getDocument({ data }).promise;

  const pages: string[] = [];
  const images: PreparedImage[] = [];
  const warnings: string[] = [];
  let imagesTruncated = false;

  const totalPages = Math.min(document.numPages, 80);
  if (document.numPages > totalPages) {
    warnings.push(`Only the first ${totalPages} of ${document.numPages} pages were read.`);
  }

  for (let pageNumber = 1; pageNumber <= totalPages; pageNumber++) {
    onProgress?.({
      message: t("doc.readingPage", { n: pageNumber, total: totalPages }),
      progress: pageNumber / totalPages,
    });

    const page = (await document.getPage(pageNumber)) as unknown as PdfPage;
    let pageText = "";
    try {
      const content = await page.getTextContent();
      pageText = textFromPdfItems(content.items);
    } catch {
      warnings.push(`Page ${pageNumber} text could not be extracted.`);
    }

    if (pageText.trim().length >= SCANNED_PAGE_CHAR_THRESHOLD) {
      pages.push(`--- page ${pageNumber} ---\n${pageText.trim()}`);
    } else if (images.length < MAX_DOCUMENT_IMAGES) {
      // Scanned or image-only page: render it so the vision model can read it.
      try {
        const viewport = page.getViewport({ scale: PDF_RENDER_SCALE });
        const canvas = window.document.createElement("canvas");
        canvas.width = Math.floor(viewport.width);
        canvas.height = Math.floor(viewport.height);
        const context = canvas.getContext("2d");
        if (!context) throw new Error("no 2d context");
        context.fillStyle = "#ffffff";
        context.fillRect(0, 0, canvas.width, canvas.height);
        await page.render({ canvasContext: context, viewport }).promise;
        const blob = await new Promise<Blob | null>((resolve) => canvas.toBlob(resolve, "image/jpeg", 0.9));
        if (blob) {
          images.push(await prepareImageFromBlob(blob, `${file.name} p${pageNumber}`));
        }
      } catch {
        warnings.push(`Page ${pageNumber} looks scanned but could not be rendered.`);
      }
    } else {
      imagesTruncated = true;
    }

    if (pageText.trim().length < SCANNED_PAGE_CHAR_THRESHOLD) {
      // Nothing to add to the text; the image carries the content.
      imagesTruncated = images.length >= MAX_DOCUMENT_IMAGES;
    }

    page.cleanup?.();
  }

  if (images.length > 0) {
    warnings.push(
      images.length === 1
        ? "1 scanned page was turned into an image for the vision model."
        : `${images.length} scanned pages were turned into images for the vision model.`,
    );
  }

  return withLimits(file.name, "pdf", pages.join("\n\n"), images, {
    pages: document.numPages,
    warnings,
    imagesTruncated,
  });
}

/* ------------------------------------------------------------------ */
/* DOCX                                                                */
/* ------------------------------------------------------------------ */

async function extractDocx(file: File): Promise<ExtractedDocument> {
  const mammoth = await import("mammoth");
  const arrayBuffer = await file.arrayBuffer();
  const warnings: string[] = [];
  const images: PreparedImage[] = [];

  const textResult = await mammoth.extractRawText({ arrayBuffer });
  for (const message of textResult.messages.slice(0, 3)) {
    if (message.type === "warning") warnings.push(message.message);
  }

  // Pull out embedded images (diagrams, pasted screenshots) for the vision step.
  try {
    const htmlResult = await mammoth.convertToHtml(
      { arrayBuffer },
      {
        convertImage: mammoth.images.imgElement(async (image: { read: (encoding: string) => Promise<string>; contentType: string }) => {
          if (images.length >= MAX_DOCUMENT_IMAGES) return { src: "" };
          try {
            const base64 = await image.read("base64");
            const binary = atob(base64);
            const bytes = new Uint8Array(binary.length);
            for (let i = 0; i < binary.length; i++) bytes[i] = binary.charCodeAt(i);
            const blob = new Blob([bytes], { type: image.contentType || "image/png" });
            images.push(await prepareImageFromBlob(blob, `${file.name} image ${images.length + 1}`));
          } catch {
            /* skip unreadable images */
          }
          return { src: "" };
        }),
      },
    );
    void htmlResult;
  } catch {
    warnings.push("Embedded images could not be extracted (text was still read).");
  }

  const imagesTruncated = images.length >= MAX_DOCUMENT_IMAGES;
  if (images.length > 0) {
    warnings.push(`${images.length} embedded image${images.length === 1 ? "" : "s"} sent to the vision model.`);
  }

  return withLimits(file.name, "docx", textResult.value, images, { warnings, imagesTruncated });
}

/* ------------------------------------------------------------------ */
/* PPTX                                                                */
/* ------------------------------------------------------------------ */

const XML_TEXT_RUN = /<a:t[^>]*>([\s\S]*?)<\/a:t>/g;
const XML_PARAGRAPH = /<\/a:p>/g;

function decodeXmlEntities(value: string): string {
  return value
    .replace(/&lt;/g, "<")
    .replace(/&gt;/g, ">")
    .replace(/&quot;/g, '"')
    .replace(/&apos;/g, "'")
    .replace(/&#(\d+);/g, (_, code: string) => String.fromCodePoint(Number(code)))
    .replace(/&#x([0-9a-f]+);/gi, (_, code: string) => String.fromCodePoint(Number.parseInt(code, 16)))
    .replace(/&amp;/g, "&");
}

function textFromSlideXml(xml: string): string {
  const withBreaks = xml.replace(XML_PARAGRAPH, "\n");
  const parts: string[] = [];
  for (const match of withBreaks.matchAll(XML_TEXT_RUN)) {
    parts.push(decodeXmlEntities(match[1]));
  }
  return cleanText(parts.join(" "));
}

async function extractOoxmlPackage(
  file: File,
  kind: "pptx" | "opendocument",
): Promise<ExtractedDocument> {
  const JSZip = (await import("jszip")).default;
  const zip = await JSZip.loadAsync(await file.arrayBuffer());
  const warnings: string[] = [];
  const images: PreparedImage[] = [];
  const sections: string[] = [];

  if (kind === "pptx") {
    const slideNames = Object.keys(zip.files)
      .filter((name) => /^ppt\/slides\/slide\d+\.xml$/.test(name))
      .sort((a, b) => {
        const na = Number(a.match(/(\d+)/)?.[1] ?? 0);
        const nb = Number(b.match(/(\d+)/)?.[1] ?? 0);
        return na - nb;
      });

    for (const [index, name] of slideNames.entries()) {
      const xml = await zip.file(name)?.async("string");
      if (!xml) continue;
      const body = textFromSlideXml(xml);
      const notesName = `ppt/notesSlides/notesSlide${index + 1}.xml`;
      const notesXml = zip.file(notesName) ? await zip.file(notesName)?.async("string") : null;
      const notes = notesXml ? textFromSlideXml(notesXml) : "";
      const chunk = [`--- slide ${index + 1} ---`, body || "(no text on this slide)"];
      if (notes) chunk.push(`Speaker notes: ${notes}`);
      sections.push(chunk.join("\n"));
    }

    const media = Object.keys(zip.files).filter(
      (name) => /^ppt\/media\/.+\.(png|jpe?g|gif|webp|bmp)$/i.test(name) && !name.includes("/_rels/"),
    );
    for (const name of media) {
      if (images.length >= MAX_DOCUMENT_IMAGES) break;
      try {
        const blob = await zip.file(name)?.async("blob");
        if (blob) images.push(await prepareImageFromBlob(blob, `${file.name} ${name.split("/").pop()}`));
      } catch {
        /* skip */
      }
    }
    if (media.length > images.length) {
      warnings.push(`Only ${images.length} of ${media.length} slide images were kept.`);
    }
  } else {
    const content = await zip.file("content.xml")?.async("string");
    if (!content) throw new Error("This OpenDocument file has no readable content.");
    const text = decodeXmlEntities(
      content
        .replace(/<text:p[^>]*>/g, "\n")
        .replace(/<text:h[^>]*>/g, "\n")
        .replace(/<\/text:p>/g, "\n")
        .replace(/<\/text:h>/g, "\n")
        .replace(/<[^>]+>/g, ""),
    );
    sections.push(text);
  }

  return withLimits(file.name, kind, sections.join("\n\n"), images, {
    pages: kind === "pptx" ? sections.length : undefined,
    warnings,
    imagesTruncated: images.length >= MAX_DOCUMENT_IMAGES,
  });
}

/* ------------------------------------------------------------------ */
/* Plain text                                                          */
/* ------------------------------------------------------------------ */

async function extractPlainText(file: File): Promise<ExtractedDocument> {
  const text = await file.text();
  return withLimits(file.name, "text", text, [], {});
}

/* ------------------------------------------------------------------ */
/* Dispatcher                                                          */
/* ------------------------------------------------------------------ */

export async function extractDocument(
  file: File,
  onProgress?: (update: DocumentProgress) => void,
): Promise<ExtractedDocument> {
  const kind = documentKindOf(file);
  if (!kind) throw new Error(`"${file.name}" is not a supported document.`);

  switch (kind) {
    case "pdf":
      return extractPdf(file, onProgress);
    case "docx":
      onProgress?.({ message: t("doc.readingWord") });
      return extractDocx(file);
    case "pptx":
      onProgress?.({ message: t("doc.readingSlides") });
      return extractOoxmlPackage(file, "pptx");
    case "opendocument":
      onProgress?.({ message: t("doc.readingOpen") });
      return extractOoxmlPackage(file, "opendocument");
    case "text":
      onProgress?.({ message: t("doc.readingText") });
      return extractPlainText(file);
  }
}

/** Merge several documents into one text block with file headers. */
export function documentsToText(documents: ExtractedDocument[]): string {
  return documents
    .filter((document) => document.text.trim().length > 0)
    .map((document) => `### ${document.name}\n${document.text.trim()}`)
    .join("\n\n")
    .slice(0, MAX_DOCUMENT_CHARS * 2);
}