"""The Watermark of Engine > Post: a text or a picture in one corner of the finished video.

`build(values, width, height)` draws it once (a small RGBA layer and where it goes); `apply(frames, layer)` lays it over a chunk of
frames. It runs last, after the picture's own look, so contrast, colour and smoothing never touch it, and the carried last frame
(which never sees Post) never has it. Everything the layer meets is moved to that tensor's device and dtype first.
"""

from collections import namedtuple

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image, ImageDraw, ImageFont

from ..._core import log, media

MODES = ("off", "text", "picture")
CORNERS = ("top-left", "top-right", "bottom-left", "bottom-right")
MAX_HEIGHT = 0.5                      # the mark never takes more than half the picture's height, whatever its shape
TITLE = "FunUI Save Video"

Layer = namedtuple("Layer", "rgb alpha y x")        # rgb [h, w, 3] and alpha [h, w, 1] floats 0..1 (straight, not premultiplied); where its top-left lands


def picture_options():
    """The pictures in the media bin, for the Picture choice."""
    items = [{"value": "", "label": "None picked"}]
    try:
        items += [{"value": m["id"], "label": m.get("name") or m["id"]} for m in media.listing() if m.get("kind") == "image"]
    except Exception as exc:                                  # noqa: BLE001 -- an unreadable bin is an empty list, said in the log
        log.once("watermark:bin", log.WARNING, TITLE, f"the media bin could not be listed for the Watermark picture ({exc})")
    return items


def _font(px):
    try:
        return ImageFont.load_default(size=px)                # a scalable font (Pillow 10.1+)
    except TypeError:
        return ImageFont.load_default()                       # an older Pillow: a small fixed one, scaled up below


def _text_alpha(text, target_w):
    """-> alpha [h, w] floats of `text` drawn at the size that makes it `target_w` pixels wide."""
    draw = ImageDraw.Draw(Image.new("L", (8, 8)))
    probe = _font(64)
    l, t, r, b = draw.textbbox((0, 0), text, font=probe)
    px = max(6, round(64 * target_w / max(1, r - l)))
    font = _font(px)
    l, t, r, b = draw.textbbox((0, 0), text, font=font)
    img = Image.new("L", (max(1, r - l), max(1, b - t)), 0)
    ImageDraw.Draw(img).text((-l, -t), text, font=font, fill=255)
    img = img.crop(img.getbbox() or (0, 0, img.width, img.height))       # the ink itself: the font's own side room is not part of the width asked for
    if abs(img.width - target_w) > 2:                         # the smallest size, or a fixed font from an older Pillow: scaled to the width asked for
        img = img.resize((target_w, max(1, round(img.height * target_w / img.width))), Image.LANCZOS)
    return np.asarray(img, dtype=np.float32) / 255.0


def _grow(alpha, steps):
    """alpha [1, 1, h, w] -> the shape grown by `steps` pixels (a near-round edge: a square step, then a cross step, and so on). Memory: one map, whatever `steps`."""
    a = alpha
    for i in range(steps):
        if i % 2 == 0:
            a = F.max_pool2d(a, 3, 1, 1)
        else:
            a = torch.maximum(F.max_pool2d(a, (1, 3), 1, (0, 1)), F.max_pool2d(a, (3, 1), 1, (1, 0)))
    return a


def _hex(colour):
    c = str(colour or "").lstrip("#")
    if len(c) == 3:
        c = "".join(ch * 2 for ch in c)
    try:
        return tuple(int(c[i:i + 2], 16) / 255.0 for i in (0, 2, 4))
    except ValueError:
        return (0.0, 0.0, 0.0)


def build(values, width, height):
    """-> a Layer for these (checked) Post values and this frame size, or None when there is nothing to draw (said when it should have been)."""
    mode = values.get("watermark", "off")
    if mode not in ("text", "picture"):
        return None
    target_w = max(8, round(width * float(values.get("watermark_size", 18.0)) / 100.0))
    if mode == "text":
        text = str(values.get("watermark_text") or "").strip()
        if not text:
            log.once("watermark:empty", log.ALERT, TITLE, "Watermark is set to Text but the text is empty, so the video has none.")
            return None
        alpha = _text_alpha(text, target_w)
        rgb = np.broadcast_to(np.array(_hex(values.get("watermark_colour", "#ffffff")), dtype=np.float32), alpha.shape + (3,))
    else:
        path = media.path_for(str(values.get("watermark_picture") or "")) if values.get("watermark_picture") else None
        if path is None:
            log.once("watermark:nopic", log.ALERT, TITLE, "Watermark is set to Picture but no picture from the media bin is picked (or it was removed), so the video has none.")
            return None
        try:
            img = Image.open(path).convert("RGBA")
        except Exception as exc:                              # noqa: BLE001 -- say which file and why, never a silent skip
            log.once("watermark:badpic", log.ALERT, TITLE, f"the Watermark picture could not be read ({exc}), so the video has none.")
            return None
        img = img.resize((target_w, max(1, round(img.height * target_w / img.width))), Image.LANCZOS)
        arr = np.asarray(img, dtype=np.float32) / 255.0
        rgb, alpha = arr[..., :3], arr[..., 3]
    if alpha.shape[0] > height * MAX_HEIGHT:                  # a tall mark is brought down to fit, keeping its shape
        k = height * MAX_HEIGHT / alpha.shape[0]
        w2, h2 = max(1, round(alpha.shape[1] * k)), max(1, round(alpha.shape[0] * k))
        rgb = np.asarray(Image.fromarray((rgb * 255).astype(np.uint8)).resize((w2, h2), Image.LANCZOS), dtype=np.float32) / 255.0
        alpha = np.asarray(Image.fromarray((alpha * 255).astype(np.uint8)).resize((w2, h2), Image.LANCZOS), dtype=np.float32) / 255.0
    rgb, alpha = torch.from_numpy(np.ascontiguousarray(rgb)).float(), torch.from_numpy(np.ascontiguousarray(alpha)).float()
    mark_a = alpha * float(values.get("watermark_opacity", 0.6))
    ring = int(round(float(values.get("watermark_outline", 0.0)) / 100.0 * alpha.shape[0]))
    out_rgb, out_a = rgb, mark_a
    if ring > 0:
        pad = ring + 1
        a = F.pad(alpha[None, None], (pad, pad, pad, pad))
        rgb = F.pad(rgb.permute(2, 0, 1)[None], (pad, pad, pad, pad))[0].permute(1, 2, 0)
        grown = (_grow(a, ring) > 1e-3).float()
        grown = F.avg_pool2d(F.pad(grown, (1, 1, 1, 1), mode="replicate"), 3, 1)      # a soft edge instead of a staircase
        outline = (grown - a).clamp(0, 1)[0, 0] * float(values.get("watermark_outline_opacity", 1.0))     # only outside the mark: a see-through mark shows no outline behind it
        mark_a = F.pad(mark_a[None, None], (pad, pad, pad, pad))[0, 0]
        total = mark_a + outline * (1 - mark_a)
        col = torch.tensor(_hex(values.get("watermark_outline_colour", "#000000")), dtype=torch.float32)
        out_rgb = (rgb * mark_a[..., None] + col * (outline * (1 - mark_a))[..., None]) / total.clamp(min=1e-6)[..., None]
        out_a = total
    h, w = out_a.shape
    m = round(width * float(values.get("watermark_margin", 2.0)) / 100.0)
    corner = values.get("watermark_corner", "bottom-right")
    x = m if corner.endswith("left") else width - w - m
    y = m if corner.startswith("top") else height - h - m
    x, y = max(0, min(x, width - w)), max(0, min(y, height - h))
    return Layer(out_rgb.clamp(0, 1), out_a.clamp(0, 1)[..., None], y, x)


def blend(region, layer):
    """region [T, h, w, 3] floats (the corner of a chunk, exactly the layer's size) -> the same, with the layer laid over it, in the region's dtype and device."""
    a, c = layer.alpha.to(device=region.device, dtype=torch.float32), layer.rgb.to(device=region.device, dtype=torch.float32)
    return (region.float() * (1 - a) + c * a).to(region.dtype)


def corner(x, layer):
    """The part of x [T, H, W, 3] the layer covers (a view: nothing is copied)."""
    h, w = layer.alpha.shape[0], layer.alpha.shape[1]
    if layer.y + h > x.shape[1] or layer.x + w > x.shape[2]:            # the frame changed size after the layer was drawn: no guess about where it goes
        raise ValueError("the Watermark was drawn for another frame size")
    return x[:, layer.y:layer.y + h, layer.x:layer.x + w, :]


def apply(x, layer):
    """x [T, H, W, 3] floats -> a copy with the layer laid over its corner (same shape, dtype and device). The save path uses corner() and blend() instead: it never copies a whole chunk."""
    if layer is None:
        return x
    out = x.clone()
    view = corner(out, layer)
    view.copy_(blend(view, layer))
    return out