import type { ToolInput } from "../../core/io/input-types";
import type { ToolResult } from "../../core/io/output-types";
import type { ToolContext } from "../../core/plugins/plugin-types";
import type { PdfToDocxOcrOptions } from "./index";

function abortIfNeeded(signal: AbortSignal) {
  if (signal.aborted) throw new DOMException("Operation cancelled", "AbortError");
}

export async function runPdfToDocxOcr(
  input: ToolInput,
  options: PdfToDocxOcrOptions,
  context: ToolContext
): Promise<ToolResult> {
  const file = input.files?.[0];
  if (!file) throw new Error("Add a PDF document.");
  if (file.type && file.type !== "application/pdf") throw new Error("The selected file is not a PDF.");

  abortIfNeeded(context.signal);
  context.reportProgress({ percentage: 1, message: "Opening PDF locally..." });

  const [pdfjs, { createWorker }] = await Promise.all([
    import("pdfjs-dist"),
    import("tesseract.js"),
  ]);
  pdfjs.GlobalWorkerOptions.workerSrc = new URL(
    "pdfjs-dist/build/pdf.worker.min.mjs",
    import.meta.url
  ).toString();

  const loadingTask = pdfjs.getDocument({ data: new Uint8Array(await file.arrayBuffer()) });
  const pdf = await loadingTask.promise;
  const pageTexts: string[] = [];
  const scale = options.quality === "high" ? 2.5 : 1.75;

  const worker = await createWorker(options.language, undefined, {
    logger: (message) => {
      if (message.status === "recognizing text" && typeof message.progress === "number") {
        const completedPages = pageTexts.length;
        context.reportProgress({
          percentage: 8 + ((completedPages + message.progress) / pdf.numPages) * 88,
          message: `Recognizing page ${completedPages + 1} of ${pdf.numPages}...`,
        });
      }
    },
  });

  const cancelWorker = () => void worker.terminate();
  context.signal.addEventListener("abort", cancelWorker, { once: true });

  try {
    for (let pageNumber = 1; pageNumber <= pdf.numPages; pageNumber += 1) {
      abortIfNeeded(context.signal);
      context.reportProgress({
        percentage: 5 + ((pageNumber - 1) / pdf.numPages) * 88,
        message: `Rendering page ${pageNumber} of ${pdf.numPages} locally...`,
      });

      const page = await pdf.getPage(pageNumber);
      const viewport = page.getViewport({ scale });
      const canvas = document.createElement("canvas");
      canvas.width = Math.ceil(viewport.width);
      canvas.height = Math.ceil(viewport.height);
      const canvasContext = canvas.getContext("2d", { alpha: false });
      if (!canvasContext) throw new Error("Your browser could not render this PDF page.");
      await page.render({ canvas, canvasContext, viewport }).promise;
      abortIfNeeded(context.signal);

      const result = await worker.recognize(canvas);
      pageTexts.push(result.data.text.trim());
      page.cleanup();
      canvas.width = 1;
      canvas.height = 1;
    }
  } finally {
    context.signal.removeEventListener("abort", cancelWorker);
    await worker.terminate().catch(() => undefined);
    await loadingTask.destroy();
  }

  abortIfNeeded(context.signal);
  context.reportProgress({ percentage: 97, message: "Building editable DOCX..." });
  const { Document, Packer, Paragraph } = await import("docx");
  const children: Array<InstanceType<typeof Paragraph>> = [];
  pageTexts.forEach((text, pageIndex) => {
    const lines = text.split(/\r?\n/);
    lines.forEach((line, lineIndex) => {
      children.push(new Paragraph({
        text: line,
        pageBreakBefore: options.pageBreaks && pageIndex > 0 && lineIndex === 0,
        spacing: { after: line.trim() ? 120 : 0 },
      }));
    });
  });

  const docxDocument = new Document({ sections: [{ children }] });
  const blob = await Packer.toBlob(docxDocument);
  context.reportProgress({ percentage: 100, message: "DOCX ready." });

  return {
    type: "files",
    files: [{
      name: `${file.name.replace(/\.pdf$/i, "")}_ocr.docx`,
      mimeType: "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
      blob,
      sizeBefore: file.size,
      sizeAfter: blob.size,
    }],
  };
}