import { PDFDocument, StandardFonts, TextRenderingMode } from '@cantoo/pdf-lib';
import type { Word } from 'tesseract.js';
import { registerFeature } from '../registry';
import { featureModal, pdfSource, footerButtons, deliver, type LoadedFile } from '../common';
import { bufferOf, baseName } from '../../core/util';
import { loadPdf, renderPage } from '../../core/pdf';
import { field, select } from '../../ui/controls';
import { toast } from '../../ui/toast';
import { h } from '../../core/dom';

const LANGS: { value: string; label: string }[] = [
  { value: 'eng', label: 'English' },
  { value: 'spa', label: 'Spanish' },
  { value: 'fra', label: 'French' },
  { value: 'deu', label: 'German' },
  { value: 'ita', label: 'Italian' },
  { value: 'por', label: 'Portuguese' },
];

/** Flatten the tesseract.js v7 block → paragraph → line → word tree into a flat word list. */
function collectWords(blocks: unknown): Word[] {
  const out: Word[] = [];
  if (!Array.isArray(blocks)) return out;
  for (const block of blocks as any[]) {
    for (const para of block?.paragraphs ?? []) {
      for (const line of para?.lines ?? []) {
        for (const word of line?.words ?? []) {
          if (word && typeof word.text === 'string' && word.bbox) out.push(word as Word);
        }
      }
    }
  }
  return out;
}

registerFeature({
  id: 'ocr',
  title: 'OCR — make searchable',
  desc: 'Recognize text in a scanned PDF and add a searchable layer.',
  icon: 'ocr',
  category: 'ocr',
  run(ctx) {
    let files: LoadedFile[] = [];
    let lang = 'eng';

    const m = featureModal(
      'OCR — make searchable',
      'Add an invisible text layer so a scanned PDF becomes searchable — all in your browser.',
      'ocr',
      560,
    );

    m.body.append(
      pdfSource(ctx, {
        multiple: false,
        onChange: (f) => { files = f; primary.disabled = files.length === 0; },
      }),
      field(
        'Document language',
        select({ options: LANGS, value: lang, onchange: (v) => { lang = v; } }),
        'Pick the language of the text in the scan for the best accuracy.',
      ),
      h('p.hint', 'OCR quality depends on how clear the scan is — faint, skewed, or low-resolution pages recognize less reliably.'),
    );

    const { primary } = footerButtons(m, 'Run OCR & download', async () => {
      const src = files[0];
      if (!src) return;

      const t = toast('Preparing OCR…', { type: 'progress' });
      // Tesseract is large; load it on demand so opening the editor stays light.
      const { createWorker } = await import('tesseract.js');
      let worker: Awaited<ReturnType<typeof createWorker>> | null = null;
      try {
        // Load the original with pdf-lib so we keep the existing page images intact,
        // and add the invisible text layer on top.
        const doc = await PDFDocument.load(bufferOf(src.bytes), { ignoreEncryption: true });
        const font = await doc.embedFont(StandardFonts.Helvetica);
        const pdfPages = doc.getPages();

        // Load with pdf.js for rendering each page to a canvas.
        const pjs = await loadPdf(src.bytes);
        const total = Math.min(pjs.numPages, pdfPages.length);

        t.update(`Loading the ${lang} recognizer…`);
        worker = await createWorker(lang);

        const canvas = document.createElement('canvas');
        const s = 2; // render scale

        for (let i = 0; i < total; i++) {
          t.update(`Recognizing page ${i + 1} of ${total}…`);
          t.progress(i / total);

          const pjsPage = await pjs.getPage(i + 1);
          await renderPage(pjsPage, s, canvas, { dpr: 1 });

          // Tesseract.js v7 omits word/block geometry unless explicitly requested.
          const { data } = await worker.recognize(canvas, {}, { blocks: true, text: true });
          const words = collectWords((data as { blocks?: unknown }).blocks);

          const page = pdfPages[i];
          const pageH = page.getHeight();

          for (const word of words) {
            if (word.confidence <= 30) continue;
            const text = word.text?.trim();
            if (!text) continue;
            try {
              const { x0, x1, y0, y1 } = word.bbox;
              const x = x0 / s;
              const wpx = (x1 - x0) / s;
              const hpx = (y1 - y0) / s;
              const y = pageH - y1 / s;

              let size = Math.max(4, hpx);
              // Fit the invisible text width to the recognized box width so
              // selection/highlight rectangles line up with the underlying image.
              const drawn = font.widthOfTextAtSize(text, size);
              if (drawn > 0 && wpx > 0) size = Math.max(1, (size * wpx) / drawn);

              page.drawText(text, {
                x,
                y,
                size,
                font,
                renderMode: TextRenderingMode.Invisible,
              });
            } catch {
              // Skip words with glyphs the standard font can't encode.
            }
          }
        }

        t.progress(1);
        t.update('Saving your searchable PDF…');

        await worker.terminate();
        worker = null;

        const bytes = await doc.save({ useObjectStreams: true });
        t.close();
        m.close();
        deliver(bytes, baseName(src.name) + '-ocr.pdf');
      } catch (e) {
        t.close();
        if (worker) { try { await worker.terminate(); } catch { /* ignore */ } }
        m.setBusy(false);
        toast('OCR failed — the file may be damaged or the language data could not load.', { type: 'error' });
        console.error(e);
      }
    }, 'ocr');

    primary.disabled = true;
  },
});
