import {
  PDFDocument,
  StandardFonts,
  rgb,
  degrees,
  PDFString,
  PDFName,
  PDFArray,
  pushGraphicsState,
  popGraphicsState,
  translate,
  rotateDegrees,
  scale,
  type PDFFont,
  type PDFPage,
} from '@cantoo/pdf-lib';
import type {
  YeetDoc,
  Anno,
  TextAnno,
  DrawAnno,
  ShapeAnno,
  MarkAnno,
  NoteAnno,
  ImageAnno,
  LinkAnno,
  FieldAnno,
  PageState,
} from './model';
import { effectiveRotation, markThickness } from './model';
import { bufferOf } from './util';
import { compactPdf } from './mupdf';
import { familyById, fontFileUrl, variantFor, type FontVariant } from './fonts';

function hexToRgb(hex: string): { r: number; g: number; b: number } {
  const m = hex.replace('#', '');
  const v = m.length === 3 ? m.split('').map((c) => c + c).join('') : m;
  const n = parseInt(v, 16);
  return { r: ((n >> 16) & 255) / 255, g: ((n >> 8) & 255) / 255, b: (n & 255) / 255 };
}
function col(hex: string) {
  const c = hexToRgb(hex);
  return rgb(c.r, c.g, c.b);
}

interface FontSet {
  regular: PDFFont;
  bold: PDFFont;
  italic: PDFFont;
  boldItalic: PDFFont;
}

async function loadFonts(doc: PDFDocument, annos: Anno[]): Promise<Record<string, FontSet>> {
  const e = (f: StandardFonts) => doc.embedStandardFont(f);
  const fonts: Record<string, FontSet> = {
    Helvetica: {
      regular: await e(StandardFonts.Helvetica),
      bold: await e(StandardFonts.HelveticaBold),
      italic: await e(StandardFonts.HelveticaOblique),
      boldItalic: await e(StandardFonts.HelveticaBoldOblique),
    },
    'Times New Roman': {
      regular: await e(StandardFonts.TimesRoman),
      bold: await e(StandardFonts.TimesRomanBold),
      italic: await e(StandardFonts.TimesRomanItalic),
      boldItalic: await e(StandardFonts.TimesRomanBoldItalic),
    },
    Courier: {
      regular: await e(StandardFonts.Courier),
      bold: await e(StandardFonts.CourierBold),
      italic: await e(StandardFonts.CourierOblique),
      boldItalic: await e(StandardFonts.CourierBoldOblique),
    },
  };
  await loadCustomFonts(doc, annos, fonts);
  return fonts;
}

/**
 * Embed the bundled (non-core) webfonts that text annotations actually use.
 * Only the variants in use are fetched from /fonts/ and embedded (subsetted, so
 * they stay small). fontkit is imported lazily so it's out of the initial bundle
 * and only loads when a document really contains custom-font text.
 */
async function loadCustomFonts(doc: PDFDocument, annos: Anno[], fonts: Record<string, FontSet>) {
  const need = new Map<string, Set<FontVariant>>();
  for (const a of annos) {
    if (a.type !== 'text') continue;
    const fam = familyById(a.fontFamily);
    if (!fam || fam.core) continue;
    if (!need.has(a.fontFamily)) need.set(a.fontFamily, new Set());
    need.get(a.fontFamily)!.add(variantFor(a.bold, a.italic));
  }
  if (!need.size) return;
  let registered = false;
  for (const [family, variants] of need) {
    const got: Partial<Record<FontVariant, PDFFont>> = {};
    for (const v of variants) {
      const url = fontFileUrl(family, v);
      if (!url) continue;
      try {
        const bytes = new Uint8Array(await (await fetch(url)).arrayBuffer());
        if (!registered) { doc.registerFontkit((await import('@pdf-lib/fontkit')).default); registered = true; }
        got[v] = await doc.embedFont(bytes, { subset: true });
      } catch (err) {
        console.warn('[export] custom font embed failed', family, v, err);
      }
    }
    const fb = got.regular ?? got.bold ?? got.italic ?? got.boldItalic ?? fonts.Helvetica.regular;
    fonts[family] = { regular: got.regular ?? fb, bold: got.bold ?? fb, italic: got.italic ?? fb, boldItalic: got.boldItalic ?? fb };
  }
}

function pickFont(fonts: Record<string, FontSet>, a: TextAnno): PDFFont {
  const set = fonts[a.fontFamily] ?? fonts.Helvetica;
  if (a.bold && a.italic) return set.boldItalic;
  if (a.bold) return set.bold;
  if (a.italic) return set.italic;
  return set.regular;
}

/** Wrap text into lines honoring explicit newlines + a max width. */
function wrapLines(text: string, font: PDFFont, size: number, maxWidth: number): string[] {
  const out: string[] = [];
  for (const para of text.split('\n')) {
    if (para === '') { out.push(''); continue; }
    const words = para.split(/(\s+)/);
    let line = '';
    for (const w of words) {
      const test = line + w;
      if (font.widthOfTextAtSize(test.trimEnd(), size) > maxWidth && line.trim()) {
        out.push(line.trimEnd());
        line = w.trimStart();
      } else {
        line = test;
      }
    }
    out.push(line.trimEnd());
  }
  return out;
}

function drawText(page: PDFPage, a: TextAnno, fonts: Record<string, FontSet>) {
  const ph = page.getHeight();
  const font = pickFont(fonts, a);
  const size = a.fontSize;
  const lineHeight = size * 1.18;
  const pad = 2;
  const maxWidth = Math.max(4, a.w - pad * 2);

  if (a.bg) {
    page.drawRectangle({
      x: a.x, y: ph - a.y - a.h, width: a.w, height: a.h,
      color: col(a.bg), opacity: a.opacity ?? 1,
    });
  }

  const lines = wrapLines(a.text, font, size, maxWidth);
  const ascent = font.heightAtSize(size, { descender: false });
  let baseline = ph - a.y - ascent - pad;
  for (const line of lines) {
    const lineW = font.widthOfTextAtSize(line, size);
    let x = a.x + pad;
    if (a.align === 'center') x = a.x + (a.w - lineW) / 2;
    else if (a.align === 'right') x = a.x + a.w - lineW - pad;
    page.drawText(line, { x, y: baseline, size, font, color: col(a.color), opacity: a.opacity ?? 1 });
    if (a.underline && line.trim()) {
      const uy = baseline - size * 0.12;
      page.drawLine({ start: { x, y: uy }, end: { x: x + lineW, y: uy }, thickness: Math.max(0.5, size * 0.06), color: col(a.color), opacity: a.opacity ?? 1 });
    }
    baseline -= lineHeight;
  }
}

function drawInk(page: PDFPage, a: DrawAnno) {
  const ph = page.getHeight();
  const c = col(a.color);
  for (let i = 1; i < a.points.length; i++) {
    const p0 = a.points[i - 1];
    const p1 = a.points[i];
    page.drawLine({
      start: { x: p0.x, y: ph - p0.y },
      end: { x: p1.x, y: ph - p1.y },
      thickness: a.strokeWidth,
      color: c,
      opacity: a.opacity ?? 1,
      lineCap: 1 as any,
    });
  }
}

function drawShape(page: PDFPage, a: ShapeAnno) {
  const ph = page.getHeight();
  const stroke = col(a.stroke);
  const fill = a.fill ? col(a.fill) : undefined;
  if (a.shape === 'rect') {
    page.drawRectangle({
      x: a.x, y: ph - a.y - a.h, width: a.w, height: a.h,
      borderColor: stroke, borderWidth: a.strokeWidth,
      color: fill, opacity: fill ? a.opacity ?? 1 : undefined, borderOpacity: a.opacity ?? 1,
    });
  } else if (a.shape === 'ellipse') {
    page.drawEllipse({
      x: a.x + a.w / 2, y: ph - a.y - a.h / 2, xScale: Math.abs(a.w / 2), yScale: Math.abs(a.h / 2),
      borderColor: stroke, borderWidth: a.strokeWidth, color: fill,
      opacity: fill ? a.opacity ?? 1 : undefined, borderOpacity: a.opacity ?? 1,
    });
  } else {
    // line / arrow — x,y = start; w,h are signed deltas to the end
    const start = { x: a.x, y: ph - a.y };
    const end = { x: a.x + a.w, y: ph - (a.y + a.h) };
    page.drawLine({ start, end, thickness: a.strokeWidth, color: stroke, opacity: a.opacity ?? 1, lineCap: 1 as any });
    if (a.shape === 'arrow') {
      const ang = Math.atan2(end.y - start.y, end.x - start.x);
      const len = Math.max(8, a.strokeWidth * 3.2);
      for (const off of [Math.PI - 0.5, Math.PI + 0.5]) {
        page.drawLine({
          start: end,
          end: { x: end.x + len * Math.cos(ang + off), y: end.y + len * Math.sin(ang + off) },
          thickness: a.strokeWidth, color: stroke, opacity: a.opacity ?? 1, lineCap: 1 as any,
        });
      }
    }
  }
}

function drawMark(page: PDFPage, a: MarkAnno) {
  const ph = page.getHeight();
  const c = col(a.color);
  if (a.type === 'highlight') {
    page.drawRectangle({ x: a.x, y: ph - a.y - a.h, width: a.w, height: a.h, color: c, opacity: 0.4 });
  } else if (a.type === 'underline') {
    const y = ph - a.y - a.h + 1.5;
    page.drawLine({ start: { x: a.x, y }, end: { x: a.x + a.w, y }, thickness: markThickness(a), color: c });
  } else {
    const y = ph - a.y - a.h / 2;
    page.drawLine({ start: { x: a.x, y }, end: { x: a.x + a.w, y }, thickness: markThickness(a), color: c });
  }
}

function addAnnotDict(page: PDFPage, dict: any) {
  const ctx = page.doc.context;
  const ref = ctx.register(ctx.obj(dict));
  let annots = page.node.get(PDFName.of('Annots')) as PDFArray | undefined;
  if (!annots) {
    annots = ctx.obj([]) as PDFArray;
    page.node.set(PDFName.of('Annots'), annots);
  }
  annots.push(ref);
}

function drawNote(page: PDFPage, a: NoteAnno) {
  const ph = page.getHeight();
  const c = hexToRgb(a.color);
  // real, clickable /Text annotation carrying the note contents
  addAnnotDict(page, {
    Type: 'Annot', Subtype: 'Text', Name: 'Comment', Open: false,
    Rect: [a.x, ph - a.y - 20, a.x + 20, ph - a.y],
    Contents: PDFString.of(a.text || ''),
    C: [c.r, c.g, c.b],
  });
  // plus a visible marker so it shows in flat viewers
  page.drawRectangle({
    x: a.x, y: ph - a.y - 18, width: 22, height: 18,
    color: col(a.color), borderColor: rgb(1, 1, 1), borderWidth: 1, opacity: 0.95,
  });
}

function drawLink(page: PDFPage, a: LinkAnno) {
  const ph = page.getHeight();
  addAnnotDict(page, {
    Type: 'Annot', Subtype: 'Link',
    Rect: [a.x, ph - a.y - a.h, a.x + a.w, ph - a.y],
    Border: [0, 0, 0],
    A: { Type: 'Action', S: 'URI', URI: PDFString.of(a.url) },
  });
}

async function drawImageAnno(doc: PDFDocument, page: PDFPage, a: ImageAnno) {
  const ph = page.getHeight();
  const comma = a.dataUrl.indexOf(',');
  const meta = a.dataUrl.slice(5, comma);
  const bytes = Uint8Array.from(atob(a.dataUrl.slice(comma + 1)), (ch) => ch.charCodeAt(0));
  const img = meta.includes('png') ? await doc.embedPng(bytes) : await doc.embedJpg(bytes);
  page.drawImage(img, { x: a.x, y: ph - a.y - a.h, width: a.w, height: a.h, opacity: a.opacity ?? 1 });
}

/** Translate a field's style into pdf-lib widget appearance options. */
function fieldAppearance(a: FieldAnno): Record<string, any> {
  const opt: Record<string, any> = {};
  if (a.textColor) opt.textColor = col(a.textColor);
  const style = a.fieldStyle || 'box';
  if (style === 'filled') {
    if (a.bgColor) opt.backgroundColor = col(a.bgColor);
    opt.borderWidth = 0;
  } else if (style === 'underline' || style === 'minimal') {
    opt.borderWidth = 0;
  } else {
    if (a.borderColor) opt.borderColor = col(a.borderColor);
    opt.borderWidth = a.borderWidth ?? 1;
    if (a.bgColor) opt.backgroundColor = col(a.bgColor);
  }
  return opt;
}

let fieldSeq = 0;
function addField(doc: PDFDocument, page: PDFPage, a: FieldAnno, fonts: Record<string, FontSet>) {
  const ph = page.getHeight();
  const form = doc.getForm();
  const rect = { x: a.x, y: ph - a.y - a.h, width: a.w, height: a.h };
  const app = fieldAppearance(a);
  const name = `${a.name || 'field'}_${++fieldSeq}`;
  const helv = fonts.Helvetica.regular;
  const size = a.fontSize ?? 12;

  // an "underline" style needs the line drawn manually (widgets only do full borders)
  if ((a.fieldStyle || 'box') === 'underline') {
    page.drawLine({
      start: { x: a.x, y: ph - a.y - a.h }, end: { x: a.x + a.w, y: ph - a.y - a.h },
      thickness: a.borderWidth ?? 1, color: col(a.borderColor || '#2f6bff'),
    });
  }

  if (a.field === 'text') {
    const f = form.createTextField(name);
    if (typeof a.value === 'string') f.setText(a.value);
    try { f.setFontSize(size); } catch { /* ignore */ }
    f.addToPage(page, { ...rect, ...app });
    try { f.updateAppearances(helv); } catch { /* ignore */ }
  } else if (a.field === 'checkbox') {
    const f = form.createCheckBox(name);
    f.addToPage(page, { ...rect, ...app });
    if (a.value === true) f.check();
  } else if (a.field === 'dropdown') {
    const f = form.createDropdown(name);
    f.addOptions(a.options?.length ? a.options : ['Option 1', 'Option 2']);
    if (typeof a.value === 'string' && a.value) f.select(a.value);
    try { f.setFontSize(size); } catch { /* ignore */ }
    f.addToPage(page, { ...rect, ...app });
    try { f.updateAppearances(helv); } catch { /* ignore */ }
  } else if (a.field === 'radio') {
    const group = form.createRadioGroup(name);
    const opts = a.options?.length ? a.options : ['Yes', 'No'];
    const n = opts.length;
    const rowH = a.h / n;
    opts.forEach((opt, i) => {
      group.addOptionToPage(opt, page, {
        x: a.x, y: ph - a.y - rowH * (i + 1), width: Math.min(a.w, rowH), height: rowH * 0.9, ...app,
      });
    });
    if (typeof a.value === 'string' && a.value) { try { group.select(a.value); } catch { /* ignore */ } }
  }
}

/** True once the page *set* differs from the source (reorder / delete / insert
 *  blank) — that needs a rebuilt document. A pure rotation does NOT count: it's
 *  applied in place via `page.setRotation`. */
function pageSetChanged(doc: YeetDoc, livePages: PageState[]): boolean {
  if (doc.pages.some((p) => p.deleted)) return true; // includes tail deletions
  let expected = 0;
  for (const ps of livePages) {
    if (ps.blank) return true;
    if (ps.srcIndex !== expected) return true; // reorder or mid-document deletion
    expected++;
  }
  return false;
}
function hasRotation(livePages: PageState[]): boolean {
  return livePages.some((ps) => ((ps.rotation % 360) + 360) % 360 !== 0);
}
function hasMetaEdits(m: YeetDoc['meta']): boolean {
  return !!(m.title || m.author || m.subject || m.keywords);
}

/** Draw every overlay annotation onto its page (preserves insertion z-order). */
async function drawAll(
  out: PDFDocument,
  doc: YeetDoc,
  pageMap: Map<string, PDFPage>,
  fonts: Record<string, FontSet>,
) {
  for (const a of doc.annos) {
    const page = pageMap.get(a.page);
    if (!page) continue;
    // draw treating the page as unrotated: pdf-lib draws in unrotated space and
    // the page /Rotate handles display, so no per-anno rotation is needed.
    await drawAnno(out, page, a, fonts);
  }
}

/**
 * Serialize the full editing model into final PDF bytes — page reorder / rotate /
 * delete / insert-blank + all overlay annotations + metadata.
 *
 * The engine is chosen to keep memory bounded on LARGE documents (pdf-lib holds
 * the whole object graph in the JS heap — ~30–50× the file size — so rebuilding a
 * big PDF through it OOM-crashes the tab):
 *  - Tier A — nothing for pdf-lib to do (no overlays, no page-set change, no
 *    rotation, no metadata edits): the native edits already live in `doc.bytes`,
 *    so we just flatten them in the mupdf worker (memory-light) and skip pdf-lib.
 *  - Tier B — overlays / rotation / metadata but the page set is unchanged: draw
 *    IN PLACE on the loaded document (no `copyPages` — roughly halves peak heap).
 *  - Tier C — the page set changed: rebuild via a single batched `copyPages`.
 */
export async function exportDoc(doc: YeetDoc): Promise<Uint8Array> {
  const livePages = doc.pages.filter((p) => !p.deleted);
  const rebuild = pageSetChanged(doc, livePages);

  // Tier A: pure native-edit (or untouched) export — flatten in the worker.
  // Falls back to the in-place pdf-lib path if mupdf can't handle these bytes.
  if (!rebuild && !doc.annos.length && !hasRotation(livePages) && !hasMetaEdits(doc.meta)) {
    try {
      return await compactPdf(doc.bytes);
    } catch (err) {
      console.warn('[export] mupdf compaction failed, falling back to pdf-lib', err);
    }
  }

  fieldSeq = 0;
  if (!rebuild) return exportInPlace(doc, livePages);
  return exportRebuild(doc, livePages);
}

/** Tier B — page set unchanged: draw overlays + apply rotations directly on the
 *  loaded document, no `copyPages`. */
async function exportInPlace(doc: YeetDoc, livePages: PageState[]): Promise<Uint8Array> {
  const out = await PDFDocument.load(bufferOf(doc.bytes), { ignoreEncryption: true, updateMetadata: false });
  const fonts = await loadFonts(out, doc.annos);
  const pageMap = new Map<string, PDFPage>();
  for (const ps of livePages) {
    const page = out.getPage(ps.srcIndex);
    page.setRotation(degrees(effectiveRotation(ps) % 360));
    pageMap.set(ps.id, page);
  }
  await drawAll(out, doc, pageMap, fonts);
  applyMeta(out, doc);
  return out.save({ useObjectStreams: true });
}

/** Tier C — page set changed (reorder/delete/insert): rebuild into a fresh doc.
 *  One batched `copyPages` shares the copier (far lighter than per-page copies). */
async function exportRebuild(doc: YeetDoc, livePages: PageState[]): Promise<Uint8Array> {
  const src = await PDFDocument.load(bufferOf(doc.bytes), { ignoreEncryption: true, updateMetadata: false });
  const out = await PDFDocument.create();
  const fonts = await loadFonts(out, doc.annos);
  const pageMap = new Map<string, PDFPage>();

  const srcIdx = livePages.filter((p) => !p.blank).map((p) => p.srcIndex);
  const copied = srcIdx.length ? await out.copyPages(src, srcIdx) : [];
  let ci = 0;
  for (const ps of livePages) {
    const page = ps.blank ? out.addPage([ps.width, ps.height]) : out.addPage(copied[ci++]);
    page.setRotation(degrees(effectiveRotation(ps) % 360));
    pageMap.set(ps.id, page);
  }
  await drawAll(out, doc, pageMap, fonts);
  applyMeta(out, doc);
  return out.save({ useObjectStreams: true, addDefaultPage: false });
}

async function drawAnno(doc: PDFDocument, page: PDFPage, a: Anno, fonts: Record<string, FontSet>) {
  // rotation / flip: wrap the drawing in a transform about the element's center
  const needsTransform = !!(a.rotation || a.flipH || a.flipV);
  if (needsTransform) {
    const ph = page.getHeight();
    const cx = a.x + a.w / 2;
    const cy = ph - (a.y + a.h / 2);
    page.pushOperators(
      pushGraphicsState(),
      translate(cx, cy),
      rotateDegrees(-(a.rotation || 0)), // CSS rotates clockwise; PDF is counter-clockwise
      scale(a.flipH ? -1 : 1, a.flipV ? -1 : 1),
      translate(-cx, -cy),
    );
  }
  await drawAnnoInner(doc, page, a, fonts);
  if (needsTransform) page.pushOperators(popGraphicsState());
}

async function drawAnnoInner(doc: PDFDocument, page: PDFPage, a: Anno, fonts: Record<string, FontSet>) {
  switch (a.type) {
    case 'text': drawText(page, a, fonts); break;
    case 'draw': drawInk(page, a); break;
    case 'shape': drawShape(page, a); break;
    case 'highlight':
    case 'underline':
    case 'strike': drawMark(page, a); break;
    case 'note': drawNote(page, a); break;
    case 'image':
    case 'signature': await drawImageAnno(doc, page, a); break;
    case 'link': drawLink(page, a); break;
    case 'field': addField(doc, page, a, fonts); break;
  }
}

function applyMeta(out: PDFDocument, doc: YeetDoc) {
  const m = doc.meta;
  out.setProducer('YeetPDF (yeetpdf.com)');
  out.setCreator('YeetPDF');
  if (m.title != null) out.setTitle(m.title);
  if (m.author != null) out.setAuthor(m.author);
  if (m.subject != null) out.setSubject(m.subject);
  if (m.keywords != null) out.setKeywords(m.keywords.split(/[,;]\s*/).filter(Boolean));
  try { out.setModificationDate(new Date()); } catch { /* ignore */ }
}
