"""Page rasters for the heat map, the weld map and the VOID stamp, drawn with Pillow.

Ports of the old renderMapCanvas, drawArrow, htDrawCanvas and stampVoidCanvas, and of the new frontend's
lib/heatTable.ts drawHeatTable, lib/voidStamp.ts and lib/geometry.ts bpos. Pillow's built-in sized font stands in
for Arial and Courier New (no font files are shipped); bold is drawn with a thin stroke of the same colour.
"""

from __future__ import annotations

import io
import math
from collections.abc import Iterable, Mapping, Sequence
from functools import lru_cache
from pathlib import Path
from typing import Any

from PIL import Image, ImageDraw, ImageFont

HEAT = (0xB0, 0x78, 0x00)
WELD = (0xC2, 0x18, 0x5B)
HEAT_MISSING = (0xB0, 0x00, 0x20)
HEAT_NA = (0x6B, 0x5A, 0x3A)
INK = (0x2B, 0x1D, 0x00)
HEAD_INK = (0x3D, 0x2A, 0x00)
HEAD_BG = (0xF3, 0xE2, 0xB8)
PILL_BG = (0xFF, 0xF8, 0xE6)
NO_HEAT_W = "MM"  # the width of the red dash a missing heat shows (Pillow's font has no em dash glyph)
HT_ROWS = 12
HT_FONT = 0.0075

Font = Any  # ImageFont.FreeTypeFont, or the bitmap default when FreeType is missing


@lru_cache(maxsize=256)
def font(size: float) -> Font:
    px = max(6, round(size))
    try:
        return ImageFont.load_default(size=px)
    except (TypeError, OSError, ImportError):  # Pillow without FreeType: the fixed bitmap font
        return ImageFont.load_default()


def text_w(f: Font, s: str) -> float:
    try:
        return float(f.getlength(s))
    except AttributeError:
        box = f.getbbox(s)
        return float(box[2] - box[0])


def _text(
    d: ImageDraw.ImageDraw,
    xy: tuple[float, float],
    s: str,
    f: Font,
    fill: Sequence[int],
    anchor: str = "mm",
    bold: bool = True,
) -> None:
    size = getattr(f, "size", 12)
    sw = max(1, round(size / 22)) if bold else 0
    d.text(xy, s, font=f, fill=tuple(fill), anchor=anchor, stroke_width=sw, stroke_fill=tuple(fill))


# ---------------------------------------------------------------- the page


def page_points(page: Any) -> tuple[float, float]:
    """A pdfium page's MediaBox size in points, after /Rotate."""
    try:
        llx, lly, urx, ury = page.get_mediabox()
        w, h = abs(urx - llx), abs(ury - lly)
    except Exception:  # an odd box: fall back to pdfium's own size
        w, h = page.get_size()
    if int(page.get_rotation() or 0) % 180 == 90:
        w, h = h, w
    return float(w), float(h)


def render_page(
    path: Path | str, pn: int, max_px: int = 3400, max_scale: float = 200 / 72
) -> tuple[Image.Image, float, float]:
    """Page `pn` (1-based) of a PDF as an RGB image on white, rotation applied, at most `max_px` on the long side and
    at most `max_scale` pixels per point. Returns the image and the page size in points."""
    import pypdfium2 as pdfium

    doc = pdfium.PdfDocument(str(path))
    try:
        if pn < 1 or pn > len(doc):
            raise ValueError(f"The PDF has no page {pn}.")
        page = doc[pn - 1]
        try:
            w_pt, h_pt = page_points(page)
            scale = min(max_scale, max_px / max(w_pt, h_pt, 1.0))
            bitmap = page.render(scale=scale, fill_color=(255, 255, 255, 255))
            img = bitmap.to_pil().convert("RGB")
        finally:
            page.close()
    finally:
        doc.close()
    return img, w_pt, h_pt


def to_jpeg(img: Image.Image, quality: int = 90) -> bytes:
    buf = io.BytesIO()
    img.convert("RGB").save(buf, "JPEG", quality=quality, optimize=True)
    return buf.getvalue()


# ---------------------------------------------------------------- marks


def clamp01(v: float) -> float:
    return max(0.01, min(0.99, v))


def bpos(m: Mapping[str, Any]) -> tuple[float, float, float, float]:
    """(fx, fy, bfx, bfy): the arrow point and the bubble. A stored bubble wins; else the old default spot."""
    fx = float(m["fx"]) if m.get("fx") is not None else 0.5
    fy = float(m["fy"]) if m.get("fy") is not None else 0.5
    bfx = float(m["bfx"]) if m.get("bfx") is not None else clamp01(fx + (-0.05 if fx > 0.85 else 0.05))
    bfy = float(m["bfy"]) if m.get("bfy") is not None else clamp01(fy + (0.07 if fy < 0.12 else -0.07))
    return fx, fy, bfx, bfy


def mark_radius(width: float) -> float:
    return max(18.0, width / 150)


def _arrow(
    d: ImageDraw.ImageDraw, ax: float, ay: float, bx: float, by: float, size: float, color: Sequence[int]
) -> None:
    """The old drawArrow: a filled head at (ax, ay) pointing away from (bx, by)."""
    ang = math.atan2(ay - by, ax - bx)
    ca, sa = math.cos(ang), math.sin(ang)
    pts = [(0.0, 0.0), (-size * 1.6, -size * 0.6), (-size * 1.6, size * 0.6)]
    d.polygon([(ax + x * ca - y * sa, ay + x * sa + y * ca) for x, y in pts], fill=tuple(color))


def _leaders(d: ImageDraw.ImageDraw, marks: Iterable[Mapping[str, Any]], w: int, h: int, color: Sequence[int]) -> None:
    r = mark_radius(w)
    for m in marks:
        fx, fy, bfx, bfy = bpos(m)
        ax, ay, bx, by = fx * w, fy * h, bfx * w, bfy * h
        d.line([(bx, by), (ax, ay)], fill=tuple(color), width=max(2, round(r / 7)))
        _arrow(d, ax, ay, bx, by, r / 2.4, color)


def draw_weld_marks(img: Image.Image, welds: Iterable[Mapping[str, Any]], page_n: int) -> Image.Image:
    """The weld bubbles of one page: the leader to the arrow point, the arrowhead, a white disc with a pink ring of
    r/6 and the id W{n} bold in the middle."""
    on = [wd for wd in welds if wd.get("page") == page_n]
    if not on:
        return img
    w, h = img.size
    d = ImageDraw.Draw(img)
    r = mark_radius(w)
    _leaders(d, on, w, h, WELD)
    f = font(round(r * 0.85))
    for wd in on:
        _fx, _fy, bfx, bfy = bpos(wd)
        x, y = bfx * w, bfy * h
        d.ellipse([x - r, y - r, x + r, y + r], fill=(255, 255, 255), outline=WELD, width=max(1, round(r / 6)))
        label = f"W{wd.get('weld_no')}"
        fl = f if text_w(f, label) <= r * 1.6 else font(r * 0.85 * r * 1.6 / max(text_w(f, label), 1.0))
        _text(d, (x, y), label, fl, WELD)
    return img


def heat_of(bom: Iterable[Mapping[str, Any]], item: Any) -> str:
    for b in bom:
        if str(b.get("item")) == str(item):
            heat = str(b.get("heat") or "").strip()
            return heat or ("N/A" if b.get("na") else "")
    return ""


def draw_heat_tags(
    img: Image.Image, tags: Iterable[Mapping[str, Any]], bom: Sequence[Mapping[str, Any]], page_n: int
) -> Image.Image:
    """The older heat tags of one page: an amber pill with the item in a filled circle and the heat (or NO HEAT)."""
    on = [t for t in tags if t.get("page") == page_n]
    if not on:
        return img
    w, h = img.size
    d = ImageDraw.Draw(img)
    r = mark_radius(w)
    _leaders(d, on, w, h, HEAT)
    f_txt = font(round(r * 0.8))
    f_item = font(round(r * 0.75))
    for t in on:
        _fx, _fy, bfx, bfy = bpos(t)
        x, y = bfx * w, bfy * h
        txt = heat_of(bom, t.get("item")) or "NO HEAT"
        tw = text_w(f_txt, txt)
        ph = r * 1.7
        cr = ph * 0.42
        pw = tw + ph + r * 0.9
        x0, y0 = x - pw / 2, y - ph / 2
        d.rounded_rectangle(
            [x0, y0, x0 + pw, y0 + ph], radius=ph / 2, fill=PILL_BG, outline=HEAT, width=max(1, round(r / 6))
        )
        cx = x0 + ph / 2
        d.ellipse([cx - cr, y - cr, cx + cr, y + cr], fill=HEAT)
        _text(d, (cx, y), str(t.get("item")), f_item, (255, 255, 255))
        _text(d, (x0 + ph + r * 0.2, y), txt, f_txt, HEAD_INK, anchor="lm")
    return img


# ---------------------------------------------------------------- the heat table


def heat_rows(bom: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]:
    """ITEM and HEAT NO. per BOM line: the heat, N/A on an N/A line, or empty when missing."""
    out = []
    for b in bom:
        heat = str(b.get("heat") or "").strip()
        out.append(
            {"item": str(b.get("item")), "heat": heat or ("N/A" if b.get("na") else ""), "na": bool(b.get("na"))}
        )
    return out


def heat_groups(rows: Sequence[Any]) -> list[list[Any]]:
    """At most 12 rows per group, the groups as even as they can be. Always at least one group."""
    n = len(rows)
    per = math.ceil(n / math.ceil(n / HT_ROWS)) if n > HT_ROWS else HT_ROWS
    out = [list(rows[i : i + per]) for i in range(0, n, per)]
    return out or [[]]


def draw_heat_table(img: Image.Image, placement: Mapping[str, Any], rows: Sequence[Mapping[str, Any]]) -> Image.Image:
    """The heat table on a page raster: HEAT NUMBERS on an amber band, then ITEM | HEAT NO. groups side by side.
    Font size W x 0.0075 x s; a missing heat is a red dash, N/A is brown."""
    W, H = img.size
    s = float(placement.get("s") or 1)
    fs = W * HT_FONT * s
    pad_x = fs * 0.5
    pad_y = fs * 0.2
    row_h = fs * 1.15 + pad_y * 2
    bw = max(1, round(max(1.5, fs * 0.08)))
    groups = heat_groups(rows)
    th_f, td_f, cap_f = font(round(fs * 0.8)), font(round(fs)), font(round(fs * 0.85))
    widths = []
    for g in groups:
        wi = max([text_w(th_f, "ITEM")] + [text_w(td_f, r["item"]) for r in g])
        wh = max(
            [text_w(th_f, "HEAT NO.")] + [text_w(td_f, r["heat"] or NO_HEAT_W) * (1 if r["heat"] else 0.5) for r in g]
        )
        widths.append((wi + pad_x * 2, wh + pad_x * 2))
    total_w = sum(a + b for a, b in widths)
    cap_h = row_h
    max_rows = max(len(g) for g in groups)
    total_h = cap_h + row_h * (max_rows + 1)
    x0 = float(placement.get("x") or 0) * W
    y0 = float(placement.get("y") or 0) * H
    d = ImageDraw.Draw(img)
    d.rectangle([x0, y0, x0 + total_w, y0 + total_h], fill=(255, 255, 255))
    d.rectangle([x0, y0, x0 + total_w, y0 + cap_h], fill=HEAT)
    _text(d, (x0 + total_w / 2, y0 + cap_h / 2), "HEAT NUMBERS", cap_f, (255, 255, 255))
    gx = x0
    for gi, g in enumerate(groups):
        wi, wh = widths[gi]
        y = y0 + cap_h
        d.rectangle([gx, y, gx + wi + wh, y + row_h], fill=HEAD_BG)
        _text(d, (gx + wi / 2, y + row_h / 2), "ITEM", th_f, HEAD_INK)
        _text(d, (gx + wi + wh / 2, y + row_h / 2), "HEAT NO.", th_f, HEAD_INK)
        d.rectangle([gx, y, gx + wi, y + row_h], outline=HEAT, width=bw)
        d.rectangle([gx + wi, y, gx + wi + wh, y + row_h], outline=HEAT, width=bw)
        y += row_h
        for r in g:
            _text(d, (gx + wi / 2, y + row_h / 2), r["item"], td_f, INK)
            color = (HEAT_NA if r["na"] else INK) if r["heat"] else HEAT_MISSING
            if r["heat"]:
                _text(d, (gx + wi + pad_x, y + row_h / 2), r["heat"], td_f, color, anchor="lm")
            else:
                dash = text_w(td_f, NO_HEAT_W) * 0.5
                ym = y + row_h / 2
                d.line([(gx + wi + pad_x, ym), (gx + wi + pad_x + dash, ym)], fill=color, width=max(1, round(fs / 10)))
            d.rectangle([gx, y, gx + wi, y + row_h], outline=HEAT, width=bw)
            d.rectangle([gx + wi, y, gx + wi + wh, y + row_h], outline=HEAT, width=bw)
            y += row_h
        gx += wi + wh
    d.rectangle([x0, y0, x0 + total_w, y0 + total_h], outline=HEAT, width=bw * 2)
    return img


# ---------------------------------------------------------------- VOID


def stamp_void(img: Image.Image, text: str) -> Image.Image:
    """A big diagonal VOID at low opacity, the void text under it, and a red border (voidStamp.ts)."""
    base = img.convert("RGBA")
    w, h = base.size
    layer = Image.new("RGBA", (w, h), (0, 0, 0, 0))
    d = ImageDraw.Draw(layer)
    big = font(round(w / 5.5))
    d.text(
        (w / 2, h / 2),
        "VOID",
        font=big,
        anchor="mm",
        fill=(220, 0, 0, round(255 * 0.3)),
        stroke_width=max(4, round(w / 260)),
        stroke_fill=(200, 0, 0, round(255 * 0.55)),
    )
    if text:
        small = font(round(w / 55))
        d.text((w / 2, h / 2 + w / 9), text, font=small, anchor="mm", fill=(200, 0, 0, round(255 * 0.8)))
    layer = layer.rotate(
        math.degrees(math.atan2(h, w) * 0.85), center=(w / 2, h / 2), resample=Image.Resampling.BICUBIC
    )
    out = Image.alpha_composite(base, layer)
    lw = max(6, round(w / 200))
    border = Image.new("RGBA", (w, h), (0, 0, 0, 0))
    ImageDraw.Draw(border).rectangle(
        [lw / 2, lw / 2, w - lw / 2, h - lw / 2], outline=(200, 0, 0, round(255 * 0.7)), width=lw
    )
    out = Image.alpha_composite(out, border)
    return out.convert("RGB")