"""Write click-to-edit changes back into a PDF.

Each edit is {page, engine, box: [x1, y1, x2, y2] in page-image pixels, text,
orig, grp, ts}. For every edit we:

  1. map the box from image pixels to PDF points;
  2. redact that rectangle -- on a scanned page this blanks the pixels of the
     page image itself, so the old word is really gone, not just covered --
     filling it with the background colour sampled just outside the box;
  3. write the new text into the same rectangle, shrinking the font until it
     fits, in the ink colour sampled from inside the box.

Image uploads are first wrapped into a one-page PDF so both inputs share this
path.
"""

from __future__ import annotations

import io
import json
import os
import tempfile

import fitz  # PyMuPDF
from PIL import Image


def parse_edits(raw) -> list[dict]:
    if not raw:
        return []
    try:
        data = json.loads(raw) if isinstance(raw, str) else raw
    except (TypeError, ValueError):
        return []
    out = []
    for e in data if isinstance(data, list) else []:
        try:
            box = [int(round(float(v))) for v in e["box"]][:4]
            if len(box) != 4 or box[2] <= box[0] or box[3] <= box[1]:
                continue
            out.append({
                "page": int(e.get("page", 0)),
                "engine": str(e.get("engine", "")),
                "box": box,
                "text": str(e.get("text", "")),
                "grp": str(e.get("grp") or ""),
                "ts": float(e.get("ts") or 0),
            })
        except (KeyError, TypeError, ValueError):
            continue
    return out


def _iou(a, b) -> float:
    ix = max(0, min(a[2], b[2]) - max(a[0], b[0]))
    iy = max(0, min(a[3], b[3]) - max(a[1], b[1]))
    inter = ix * iy
    if not inter:
        return 0.0
    area = lambda r: (r[2] - r[0]) * (r[3] - r[1])  # noqa: E731
    return inter / float(area(a) + area(b) - inter)


def dedupe(edits: list[dict]) -> list[dict]:
    """The same spot may be edited in both panels. Keep the most recent edit
    per matched group / overlapping area so text is never drawn twice."""
    kept: list[dict] = []
    for e in sorted(edits, key=lambda e: e["ts"], reverse=True):
        clash = any(
            k["page"] == e["page"]
            and ((e["grp"] and k["grp"] == e["grp"]) or _iou(k["box"], e["box"]) > 0.3)
            for k in kept
        )
        if not clash:
            kept.append(e)
    return kept


def sample_colors(img: Image.Image, box) -> tuple[tuple, tuple]:
    """(background, ink) as 0..1 RGB. Background = median of a ring just
    outside the box; ink = median of the darkest tenth of pixels inside it."""
    rgb = img.convert("RGB")
    W, H = rgb.size
    x1, y1, x2, y2 = box
    ring = []
    for d in (2, 4, 6):
        for x in range(max(0, x1), min(W, x2), max(1, (x2 - x1) // 16)):
            for y in (y1 - d, y2 + d):
                if 0 <= y < H:
                    ring.append(rgb.getpixel((x, y)))
        for y in range(max(0, y1), min(H, y2), max(1, (y2 - y1) // 6)):
            for x in (x1 - d, x2 + d):
                if 0 <= x < W:
                    ring.append(rgb.getpixel((x, y)))
    if not ring:
        ring = [(255, 255, 255)]

    def median(px):
        n = len(px) // 2
        return tuple(sorted(c[i] for c in px)[n] / 255.0 for i in range(3))

    bg = median(ring)
    crop = rgb.crop((max(0, x1), max(0, y1), min(W, x2), min(H, y2)))
    px = list(crop.getdata()) or [(0, 0, 0)]
    px.sort(key=lambda c: 0.299 * c[0] + 0.587 * c[1] + 0.114 * c[2])
    ink = median(px[: max(1, len(px) // 10)])
    # Low contrast (e.g. an empty box) -> plain black/white ink for legibility.
    if abs(sum(ink) - sum(bg)) < 0.6:
        ink = (0.0, 0.0, 0.0) if sum(bg) > 1.5 else (1.0, 1.0, 1.0)
    return bg, ink


def _font_for(text: str) -> str:
    """Base-14 Helvetica for Latin text; PyMuPDF's built-in CJK font otherwise."""
    return "helv" if all(ord(ch) < 0x250 for ch in text) else "china-s"


def _write_fitted(page, rect, text, color, rotate) -> bool:
    fontname = _font_for(text)
    fs = max(4.0, rect.height * 0.8)
    while fs >= 3.0:
        # insert_textbox writes nothing and returns < 0 when the text overflows.
        rc = page.insert_textbox(rect, text, fontsize=fs, fontname=fontname,
                                 color=color, align=0, rotate=rotate)
        if rc >= 0:
            return True
        fs *= 0.88
    return False


def _open_source(src_path: str, pages: list[Image.Image]) -> fitz.Document:
    if os.path.splitext(src_path)[1].lower() == ".pdf":
        return fitz.open(src_path)
    # Image upload: wrap the (already loaded) page image in a one-page PDF.
    img = pages[0]
    doc = fitz.open()
    page = doc.new_page(width=img.width, height=img.height)
    buf = io.BytesIO()
    img.convert("RGB").save(buf, format="PNG")
    page.insert_image(page.rect, stream=buf.getvalue())
    return doc


def build_edited_pdf(src_path: str, pages: list[Image.Image], edits: list[dict]):
    """-> (output_path, applied, skipped)."""
    edits = dedupe(edits)
    doc = _open_source(src_path, pages)
    applied = skipped = 0
    try:
        by_page: dict[int, list[dict]] = {}
        for e in edits:
            if 0 <= e["page"] < min(len(pages), doc.page_count):
                by_page.setdefault(e["page"], []).append(e)
            else:
                skipped += 1

        for pno, page_edits in by_page.items():
            page = doc[pno]
            img = pages[pno]
            # Page image pixels -> displayed page points, then into the
            # unrotated space that annotations and text insertion use.
            sx = page.rect.width / img.width
            sy = page.rect.height / img.height
            prepared = []
            for e in page_edits:
                x1, y1, x2, y2 = e["box"]
                bg, ink = sample_colors(img, e["box"])
                rect = fitz.Rect(x1 * sx, y1 * sy, x2 * sx, y2 * sy) * page.derotation_matrix
                rect.normalize()
                page.add_redact_annot(rect, fill=bg)
                prepared.append((rect, e, ink))
            # Blank the covered pixels of any page image (the scan itself) and
            # remove any real text underneath.
            page.apply_redactions(images=fitz.PDF_REDACT_IMAGE_PIXELS)
            for rect, e, ink in prepared:
                text = e["text"].strip()
                if not text:  # an edit to empty = "erase this"
                    applied += 1
                    continue
                if _write_fitted(page, rect, text, ink, page.rotation):
                    applied += 1
                else:
                    skipped += 1

        out = os.path.join(tempfile.mkdtemp(), "edited.pdf")
        doc.save(out, garbage=3, deflate=True)
    finally:
        doc.close()
    return out, applied, skipped