# SPDX-License-Identifier: Apache-2.0
"""Fast served postprocess + response encoding (default; ``MOGE_SERVE_FAST=0`` -> the torch/numpy path).

Same results as ``postprocess.canvas_to_camera`` + ``postprocess.to_original`` + ``np.savez_compressed``
+ ``base64.b64encode`` (``tests/test_serve_fastpath.py`` checks parity for every request-flag
combination), with three changes in HOW the work is done:

* ``camera``: the focal / z-shift solve runs on the 64x64 nearest-neighbour samples only. Upstream
  ``recover_focal_shift`` builds the same samples with ``F.interpolate(mode='nearest')`` over the full
  1080x1920 maps; here the identical rows / columns are gathered directly (indices from torch's own
  'nearest' kernel), and the same scipy solver runs on the same arrays -> bit-identical focal / shift.
* ``finish``: shift, metric scale, letterbox crop, bilinear resize to the original size, normal
  re-normalize, nearest mask resize, re-projection (force_projection) and masking in ONE multithreaded
  C++ pass (``host_ext.moge_finish``) that writes the response arrays directly (normal already float16
  for npz). Arithmetic follows torch's CPU kernels operation by operation (see the C++ comment).
* ``npz_b64``: a valid ``.npz`` (zip of ``.npy`` members, ``np.load``-compatible, same keys / dtypes /
  shapes / bytes) whose deflate streams are produced in parallel: each member is cut into 1 MiB pieces,
  each piece is raw-deflated independently (``Z_SYNC_FLUSH`` between pieces, ``Z_FINISH`` on the last:
  the pigz construction, one valid deflate stream per member) by a thread pool. Compressor (round 4):
  ISA-L's ``isal_deflate`` called through ctypes from python-isal's own shared object (``_IsalCt``,
  level 2): ctypes releases the GIL, python-isal's ``compressobj`` does not (its pool saturated at
  ~1.2 GB/s); output pieces go to one persistent scratch buffer and base64 runs over the segment list
  (``moge_b64_segs``) straight into the JSON body buffer. Fallbacks: python-isal ``isal_zlib`` (level 2) -- ~10x faster than
  zlib on these float maps, payload +4% (35.8 vs 34.5 MB for a 1080p response) -- else stdlib zlib at
  level 1 (as small as ``np.savez_compressed``'s zlib level 6, 2.5x faster). ``MOGE_NPZ_DEFLATE`` /
  ``MOGE_NPZ_LEVEL`` override. Any valid deflate stream inflates to the same bytes, so ``np.load``
  returns identical arrays. base64 is a multithreaded C++ loop.
"""
from __future__ import annotations

import ctypes
import io
import math
import os
import struct
import threading
import zlib
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, Optional, Tuple

import numpy as np
import torch
import torch.nn.functional as F

from tt_moge.tt import host_ext

Region = Tuple[int, int, int, int]
_CHUNK = 1 << 20
_pool: Optional[ThreadPoolExecutor] = None
_pool_lock = threading.Lock()
_cache: Dict[tuple, object] = {}


def enabled() -> bool:
    if os.environ.get("MOGE_SERVE_FAST", "1").strip().lower() in ("0", "false", "no", "off"):
        return False
    return host_ext.enabled()


def nthreads() -> int:
    return int(os.environ.get("MOGE_SERVE_THREADS", "0")) or host_ext.nthreads()


def _get_pool() -> ThreadPoolExecutor:
    global _pool
    with _pool_lock:
        if _pool is None:
            _pool = ThreadPoolExecutor(max_workers=max(1, nthreads()), thread_name_prefix="moge-npz")
        return _pool


def _nearest_idx(n_in: int, n_out: int) -> torch.Tensor:
    """Source indices of torch's 'nearest' resize n_in -> n_out (from the kernel itself)."""
    key = ("nearest", n_in, n_out)
    if key not in _cache:
        src = torch.arange(n_in, dtype=torch.float32)[None, None]
        idx = F.interpolate(src, size=n_out, mode="nearest")[0, 0]
        _cache[key] = idx.to(torch.int64)
    return _cache[key]


def _view_plane_uv_axes(width: int, height: int) -> Tuple[torch.Tensor, torch.Tensor]:
    """The u (width) / v (height) linspaces of moge's normalized_view_plane_uv (same expressions)."""
    key = ("uv", width, height)
    if key not in _cache:
        aspect_ratio = width / height
        span_x = aspect_ratio / (1 + aspect_ratio ** 2) ** 0.5
        span_y = 1 / (1 + aspect_ratio ** 2) ** 0.5
        u = torch.linspace(-span_x * (width - 1) / width, span_x * (width - 1) / width, width, dtype=torch.float32)
        v = torch.linspace(-span_y * (height - 1) / height, span_y * (height - 1) / height, height, dtype=torch.float32)
        _cache[key] = (u, v)
    return _cache[key]


@torch.inference_mode()
def camera(out: Dict[str, torch.Tensor], region: Optional[Region], fov_x: Optional[float] = None,
           downsample: Tuple[int, int] = (64, 64)) -> Dict[str, torch.Tensor]:
    """``postprocess.canvas_to_camera`` without the full-resolution work: returns ``focal``, ``shift``,
    ``metric_scale`` (0-d float32 tensors) and the canvas ``intrinsics`` [3,3]. ``region`` is the
    valid (non-letterbox) rectangle, None = the whole canvas (fit='stretch')."""
    from moge.utils.geometry_numpy import solve_optimal_focal_shift, solve_optimal_shift
    import utils3d

    points = out["points"][0]
    height, width = points.shape[:2]
    aspect_ratio = width / height
    iy, ix = _nearest_idx(height, downsample[0]), _nearest_idx(width, downsample[1])
    p_lr = points.index_select(0, iy).index_select(1, ix).float()                    # (h, w, 3)
    m_lr = out["mask"][0].index_select(0, iy).index_select(1, ix).float() > 0.5      # (h, w)
    if region is not None:
        x0, y0, x1, y1 = region
        m_lr &= ((iy >= y0) & (iy < y1))[:, None] & ((ix >= x0) & (ix < x1))[None, :]
    u, v = _view_plane_uv_axes(width, height)
    uv_lr = torch.stack(torch.meshgrid(u[ix], v[iy], indexing="xy"), dim=-1)          # (h, w, 2)

    m_np = m_lr.numpy()
    p_np = p_lr.numpy()[m_np]
    uv_np = uv_lr.numpy()[m_np]
    if fov_x is not None:
        focal_t = torch.tensor(
            [aspect_ratio / (1 + aspect_ratio ** 2) ** 0.5 / math.tan(math.radians(fov_x) / 2)], dtype=torch.float32)
    if uv_np.shape[0] < 2:
        f_val, s_val = 1, 0
    elif fov_x is None:
        s, f = solve_optimal_focal_shift(uv_np, p_np)
        f_val, s_val = float(f), float(s)
    else:
        s_val = float(solve_optimal_shift(uv_np, p_np, focal_t.numpy()[0]))
    shift = torch.tensor([s_val], dtype=torch.float32).reshape(())
    focal = (torch.tensor([f_val], dtype=torch.float32) if fov_x is None else focal_t).reshape(-1)[0]

    fx = focal / 2 * (1 + aspect_ratio ** 2) ** 0.5 / aspect_ratio
    fy = focal / 2 * (1 + aspect_ratio ** 2) ** 0.5
    half = torch.tensor(0.5, dtype=torch.float32)
    intrinsics = utils3d.pt.intrinsics_from_focal_center(fx[None], fy[None], half, half)[0]
    metric_scale = out["metric_scale"].float().reshape(-1)[0]
    return {"focal": focal, "shift": shift, "metric_scale": metric_scale, "intrinsics": intrinsics}


def _kinv(intrinsics: torch.Tensor) -> torch.Tensor:
    """Rows 0..2 of inv(4x4 intrinsics) exactly as utils3d.unproject_cv builds and inverts it."""
    k = intrinsics[None]
    k4 = torch.cat([
        torch.cat([k, torch.zeros((*k.shape[:-2], 3, 1), dtype=k.dtype)], dim=-1),
        torch.tensor([[0, 0, 0, 1]], dtype=k.dtype).expand(*k.shape[:-2], 1, 4)], dim=-2)
    return torch.linalg.inv(k4)[0, :3, :].contiguous()


def _uv_map_axes(height: int, width: int) -> Tuple[torch.Tensor, torch.Tensor]:
    """utils3d uv_map linspaces (u: width, v: height)."""
    key = ("uvmap", height, width)
    if key not in _cache:
        u = torch.linspace(0.5 / width, 1 - 0.5 / width, width, dtype=torch.float32)
        v = torch.linspace(0.5 / height, 1 - 0.5 / height, height, dtype=torch.float32)
        _cache[key] = (u.contiguous(), v.contiguous())
    return _cache[key]


class _Finish:
    """Arguments of one ``host_ext.moge_finish_rows`` call sequence (all rows or row bands)."""

    @torch.inference_mode()
    def __init__(self, out, cam, canvas, region, size, force_projection, apply_mask, normal_half):
        from tt_moge.server import postprocess as pp
        H, W = size
        self.pts = out["points"][0].float().contiguous()
        self.nrm = out["normal"][0].float().contiguous()
        self.msk = out["mask"][0].float().contiguous()
        self.Hc, self.Wc = self.pts.shape[:2]
        self.region, self.size = region, size
        self.intrinsics = pp.crop_intrinsics(cam["intrinsics"], canvas, region)
        self.kinv = _kinv(self.intrinsics)
        self.u, self.v = _uv_map_axes(H, W)
        x0, y0, x1, y1 = region
        self.ny = _nearest_idx(y1 - y0, H).to(torch.int32).contiguous()
        self.nx = _nearest_idx(x1 - x0, W).to(torch.int32).contiguous()
        self.shift, self.ms = float(cam["shift"]), float(cam["metric_scale"])
        self.fp, self.am, self.half = int(force_projection), int(apply_mask), int(normal_half)
        self.depth = np.empty((H, W), np.float32)
        self.points = np.empty((H, W, 3), np.float32)
        self.normal = np.empty((H, W, 3), np.float16 if normal_half else np.float32)
        self.mask = np.empty((H, W), np.uint8)
        self.cnt, self.mn, self.mx = 0, math.inf, -math.inf

    def rows(self, ya: int, yb: int, nthreads_: int):
        H, W = self.size
        x0, y0, x1, y1 = self.region
        stats = np.zeros(3, np.float64)
        rc = host_ext.lib().moge_finish_rows(
            self.pts.data_ptr(), self.nrm.data_ptr(), self.msk.data_ptr(), self.Hc, self.Wc, x0, y0, x1, y1,
            self.shift, self.ms, H, W, self.ny.data_ptr(), self.nx.data_ptr(), self.u.data_ptr(), self.v.data_ptr(),
            self.kinv.data_ptr(), self.fp, self.am, self.depth.ctypes.data, self.points.ctypes.data,
            self.normal.ctypes.data, self.half, self.mask.ctypes.data, stats.ctypes.data, nthreads_, ya, yb)
        if rc != 0:
            raise ValueError(f"moge_finish: bad region {self.region} for canvas {(self.Wc, self.Hc)}")
        self.cnt += int(stats[0])
        self.mn, self.mx = min(self.mn, float(stats[1])), max(self.mx, float(stats[2]))

    def result(self) -> Dict[str, object]:
        return {"depth": self.depth, "points": self.points, "normal": self.normal, "mask": self.mask.view(np.bool_),
                "intrinsics": self.intrinsics, "valid_count": self.cnt, "depth_min": self.mn, "depth_max": self.mx}


@torch.inference_mode()
def finish(out: Dict[str, torch.Tensor], cam: Dict[str, torch.Tensor], canvas: Tuple[int, int], region: Region,
           size: Tuple[int, int], force_projection: bool = True, apply_mask: bool = True,
           normal_half: bool = True) -> Dict[str, object]:
    """``postprocess.to_original(canvas_to_camera(...))`` as numpy arrays: ``depth`` f32 [H,W],
    ``points`` f32 [H,W,3], ``normal`` [H,W,3] (f16 if ``normal_half`` else f32), ``mask`` bool [H,W],
    ``intrinsics`` torch [3,3] (cropped, normalized), plus ``valid_count`` / ``depth_min`` / ``depth_max``."""
    f = _Finish(out, cam, canvas, region, size, force_projection, apply_mask, normal_half)
    f.rows(0, size[0], nthreads())
    return f.result()


def gather_segments(segs) -> np.ndarray:
    """``b"".join(segs)`` into one uint8 numpy buffer (served as a memoryview body, no bytes copy)."""
    views = [np.frombuffer(memoryview(x).cast("B"), np.uint8) for x in segs]
    out = np.empty(sum(v.size for v in views), np.uint8)
    o = 0
    for v in views:
        out[o:o + v.size] = v
        o += v.size
    return out


def finish_npz_b64(out, cam, canvas, region, size, force_projection: bool, apply_mask: bool,
                   metric_scale: float, headroom: int = 0, bands: Optional[int] = None,
                   timing: Optional[Dict[str, float]] = None, raw_npz: bool = False):
    """``finish`` + ``npz_b64_bytes`` of the npz response arrays, overlapped: the finish pass runs in row
    bands and each band's rows of points / depth / normal / mask are deflated (thread pool) while the
    next band is computed. Same arrays, same npz members (bit-identical decoded values). Returns
    (finish result dict, base64) -- or with ``raw_npz`` the npz file itself as a uint8 array (npz_binary).
    ``timing`` (optional) gets ``finish`` = seconds until the last band."""
    import time
    t0 = time.perf_counter()
    bands = bands or int(os.environ.get("MOGE_FINISH_BANDS", "4"))
    f = _Finish(out, cam, canvas, region, size, force_projection, apply_mask, True)
    H = size[0]
    intr = f.intrinsics.numpy().astype(np.float32)
    names = ("points", "depth", "normal", "mask", "intrinsics", "metric_scale")
    arrs = {"points": f.points, "depth": f.depth, "normal": f.normal, "mask": f.mask,
            "intrinsics": intr, "metric_scale": np.float32(metric_scale)}
    mod, default_level = _deflate_backend()
    level = int(os.environ.get("MOGE_NPZ_LEVEL", str(default_level)))
    ct = isinstance(mod, _IsalCt)
    edges = [round(H * b / bands) for b in range(bands + 1)]
    edges = sorted(set(edges))
    # piece plan per member: [header] + per band the band's byte range cut into <= 1 MiB pieces
    hdrs, plan = {}, {}
    for nm in names:
        a = np.asanyarray(arrs[nm])
        hb = io.BytesIO()
        np.lib.format.write_array_header_1_0(hb, np.lib.format.header_data_from_array_1_0(a))
        hdrs[nm] = hb.getvalue()
        raw = memoryview(a.reshape(-1).view(np.uint8)) if a.size else memoryview(b"")
        per_band = []
        if nm in ("points", "depth", "normal", "mask"):
            rowb = len(raw) // H
            for b in range(len(edges) - 1):
                lo, hi = edges[b] * rowb, edges[b + 1] * rowb
                per_band.append([raw[i:min(i + _CHUNK, hi)] for i in range(lo, hi, _CHUNK)])
        else:
            per_band.append([raw[i:i + _CHUNK] for i in range(0, len(raw), _CHUNK)])
        plan[nm] = (raw, per_band)
    with _scratch_lock:
        scratch = None
        if ct:
            need = sum(mod.bound(len(hdrs[nm])) + sum(mod.bound(len(p)) for ps in plan[nm][1] for p in ps)
                       for nm in names)
            scratch = _scratch(need)
        soff = 0
        pool = _get_pool()
        futs = {nm: [] for nm in names}

        def submit(nm, piece, last):
            nonlocal soff
            dst = None
            if ct:
                cap = mod.bound(len(piece))
                dst, soff = scratch[soff:soff + cap], soff + cap
            futs[nm].append(pool.submit(_deflate_piece, mod, piece, level, last, dst))

        def n_pieces(nm):
            return sum(len(ps) for ps in plan[nm][1])

        for nm in names:                       # headers (+ the small members) right away
            submit(nm, hdrs[nm], n_pieces(nm) == 0)
            if nm in ("intrinsics", "metric_scale"):
                ps = plan[nm][1][0]
                for j, p in enumerate(ps):
                    submit(nm, p, j == len(ps) - 1)
        for b in range(len(edges) - 1):
            f.rows(edges[b], edges[b + 1], nthreads())
            for nm in ("points", "depth", "normal", "mask"):
                ps = plan[nm][1][b]
                for j, p in enumerate(ps):
                    submit(nm, p, b == len(edges) - 2 and j == len(ps) - 1)
        if timing is not None:
            timing["finish"] = time.perf_counter() - t0
        stats_fut = None
        if _STATS_OVERLAP:   # median of the valid depths while the deflate tail runs (same values)
            stats_fut = _get_aux_pool().submit(depth_stats, f.result())
        crcs = {nm: pool.submit(_crc, mod, [hdrs[nm]] + [p for ps in plan[nm][1] for p in ps]) for nm in names}
        out_segs, central = [], []
        off = 0
        for nm in names:
            comp = [x.result() for x in futs[nm]]
            csize = sum(memoryview(c).nbytes for c in comp)
            usize = len(hdrs[nm]) + len(plan[nm][0])
            crc = crcs[nm].result()
            fb = (nm + ".npy").encode()
            lh = struct.pack("<IHHHHHIIIHH", 0x04034B50, 20, 0, 8, 0, 0x21, crc, csize, usize, len(fb), 0) + fb
            central.append(struct.pack("<IHHHHHHIIIHHHHHII", 0x02014B50, 20, 20, 0, 8, 0, 0x21, crc, csize, usize,
                                       len(fb), 0, 0, 0, 0, 0o600 << 16, off) + fb)
            out_segs.append(lh)
            out_segs.extend(comp)
            off += len(lh) + csize
        if timing is not None:
            timing["deflated"] = time.perf_counter() - t0
        cd = b"".join(central)
        out_segs.append(cd)
        out_segs.append(struct.pack("<IHHHHIIH", 0x06054B50, 0, 0, len(central), len(central), len(cd), off, 0))
        b64v = gather_segments(out_segs) if raw_npz else b64_segments(out_segs, headroom)
        if timing is not None:
            timing["b64"] = time.perf_counter() - t0
    res = f.result()
    if stats_fut is not None:
        res["depth_stats"] = stats_fut.result()
    return res, b64v


_STATS_OVERLAP = os.environ.get("MOGE_STATS_OVERLAP", "1").strip().lower() not in ("0", "false", "no", "off")
_aux_pool: Optional[ThreadPoolExecutor] = None


def _get_aux_pool() -> ThreadPoolExecutor:
    global _aux_pool
    with _pool_lock:
        if _aux_pool is None:
            _aux_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="moge-aux")
        return _aux_pool


def depth_stats(res: Dict[str, object]) -> Optional[Dict[str, float]]:
    """{"min", "median", "max"} of the valid depths, == the numpy expressions in app.predict (precomputed
    by finish_npz_b64 when overlapped)."""
    if "depth_stats" in res:
        return res["depth_stats"]
    n = res["valid_count"]
    if n == 0:
        return None
    depth, mask = res["depth"], res["mask"]
    L = host_ext.lib()
    kth = lambda k: np.float32(L.moge_masked_kth(depth.ctypes.data, mask.ctypes.data, depth.size, k, nthreads()))  # noqa: E731
    if n % 2:
        med = kth(n // 2)
    else:
        med = np.float32(kth(n // 2 - 1) + kth(n // 2)) / np.float32(2)
    return {"min": res["depth_min"], "median": float(med), "max": res["depth_max"]}


# ------------------------------------------------------------------------------------- encoding

class _IsalCt:
    """ISA-L ``isal_deflate`` called through ctypes from the python-isal extension's own shared object.
    ctypes releases the GIL for the whole call (python-isal's ``compressobj.compress`` holds it: its pool
    throughput saturated at ~1.2 GB/s on this host vs ~5 GB/s here with 8 threads). Same compressor, same
    level semantics; the produced raw-deflate pieces are valid deflate streams (checked by a round trip at
    load). Thread-local stream state + level buffer."""

    STRM = 1 << 19        # >= sizeof(struct isal_zstream) (ISA-L 2.31: ~100 KB incl. the internal state)
    LBUF = 1 << 20        # level buffer (>= ISAL_DEF_LVL3_EXTRA_LARGE)
    Z_FINISH = Z_SYNC_FLUSH = None   # (only for _deflate_piece's signature compatibility)

    def __init__(self):
        import ctypes
        from isal import _isal, isal_zlib
        if int(_isal.ISAL_MAJOR_VERSION) != 2:
            raise ImportError(f"isal_deflate layout unverified for ISA-L {_isal.ISAL_VERSION}")
        self.ct = ctypes
        L = ctypes.CDLL(isal_zlib.__file__)
        L.isal_deflate_init.argtypes = [ctypes.c_void_p]
        L.isal_deflate_init.restype = None
        L.isal_deflate.argtypes = [ctypes.c_void_p]
        L.isal_deflate.restype = ctypes.c_int
        self.L = L
        self.crc32 = isal_zlib.crc32
        self.tl = threading.local()

        class Stream(ctypes.Structure):   # struct isal_zstream (igzip_lib.h), public prefix + opaque state
            _fields_ = [("next_in", ctypes.c_void_p), ("avail_in", ctypes.c_uint32), ("total_in", ctypes.c_uint32),
                        ("next_out", ctypes.c_void_p), ("avail_out", ctypes.c_uint32), ("total_out", ctypes.c_uint32),
                        ("hufftables", ctypes.c_void_p), ("level", ctypes.c_uint32),
                        ("level_buf_size", ctypes.c_uint32), ("level_buf", ctypes.c_void_p),
                        ("end_of_stream", ctypes.c_uint16), ("flush", ctypes.c_uint16),
                        ("gzip_flag", ctypes.c_uint16), ("hist_bits", ctypes.c_uint16),
                        ("internal_state", ctypes.c_uint8 * (self.STRM - 64))]
        assert Stream.end_of_stream.offset == 56 and Stream.internal_state.offset == 64
        self.Stream = Stream
        rng = np.random.default_rng(0)                     # self test: multi-piece stream round trip
        data = (rng.standard_normal(300_000).astype(np.float32) * 10).view(np.uint8)
        for level in (0, 1, 2, 3):
            pieces = [memoryview(data)[i:i + 400_000] for i in range(0, data.size, 400_000)]
            comp = b"".join(self.deflate(p, level, j == len(pieces) - 1) for j, p in enumerate(pieces))
            d = zlib.decompressobj(-15)
            if d.decompress(comp) != data.tobytes() or not d.eof:
                raise ImportError("isal_deflate (ctypes) self test failed")

    @staticmethod
    def bound(n: int) -> int:
        return n + n // 2 + 65536                # level 0 (huffman only) can expand by ~21 %

    def deflate(self, data, level: int, last: bool, out: Optional[np.ndarray] = None) -> np.ndarray:
        ct, tl = self.ct, self.tl
        if not hasattr(tl, "s"):
            tl.s = self.Stream()
            tl.lbuf = (ct.c_uint8 * self.LBUF)()
        s = tl.s
        self.L.isal_deflate_init(ct.byref(s))
        src = np.frombuffer(data, np.uint8) if len(data) else np.zeros(1, np.uint8)[:0]
        n = src.size
        cap = self.bound(n)
        if out is None:
            out = np.empty(cap, np.uint8)
        assert out.size >= cap and out.flags.c_contiguous
        s.next_in = src.ctypes.data if n else None
        s.avail_in = n
        s.next_out = out.ctypes.data
        s.avail_out = cap
        s.level = level
        s.level_buf = ct.addressof(tl.lbuf) if level > 0 else None
        s.level_buf_size = self.LBUF if level > 0 else 0
        s.end_of_stream = 1 if last else 0
        s.flush = 1                              # SYNC_FLUSH (byte-aligned end for the next piece)
        s.gzip_flag = 0                          # raw deflate
        s.hist_bits = 0                          # default window
        rc = self.L.isal_deflate(ct.byref(s))
        if rc != 0 or s.avail_in != 0 or s.avail_out == 0:
            raise RuntimeError(f"isal_deflate rc={rc} avail_in={s.avail_in} avail_out={s.avail_out}")
        return out[:cap - s.avail_out]


_isalct_obj = None


def _isalct():
    global _isalct_obj
    if _isalct_obj is None:
        _isalct_obj = _IsalCt()
    return _isalct_obj


def _deflate_backend():
    """(module, default level): ISA-L through ctypes (GIL released; ``isalct``), python-isal's
    ``isal_zlib`` (``isal``), else stdlib zlib. Override with MOGE_NPZ_DEFLATE=isalct|isal|zlib and
    MOGE_NPZ_LEVEL."""
    want = os.environ.get("MOGE_NPZ_DEFLATE", "auto").strip().lower()
    if want in ("auto", "isalct"):
        try:
            return _isalct(), 2
        except Exception:   # noqa: BLE001 -- not installed / unexpected ABI: next backend
            if want == "isalct":
                raise
    if want in ("auto", "isal", "isalct"):
        try:
            from isal import isal_zlib
            return isal_zlib, 2
        except ImportError:
            if want == "isal":
                raise
    return zlib, 1


def deflate_backend_name() -> str:
    mod, level = _deflate_backend()
    name = "isal (ctypes)" if isinstance(mod, _IsalCt) else mod.__name__
    return f"{name} level {int(os.environ.get('MOGE_NPZ_LEVEL', str(level)))}"


def _inflate_backend():
    """isal_zlib (python-isal) when installed, else zlib: decompress + crc32 for the request decode."""
    try:
        from isal import isal_zlib
        return isal_zlib
    except ImportError:
        return zlib


def _deflate_piece(mod, data, level: int, last: bool, out=None):
    if isinstance(mod, _IsalCt):
        return mod.deflate(data, level, last, out)
    co = mod.compressobj(level, mod.DEFLATED, -15)
    return co.compress(data) + co.flush(mod.Z_FINISH if last else mod.Z_SYNC_FLUSH)


def _crc(mod, pieces) -> int:
    c = 0
    for p in pieces:
        c = mod.crc32(p, c)
    return c


def npz_segments(arrays: Dict[str, np.ndarray], level: Optional[int] = None) -> list:
    """A .npz (zip, deflate) equivalent to ``np.savez_compressed(**arrays)`` for ``np.load``, as the list
    of byte segments whose concatenation is the file (``npz_bytes`` joins them). The segments may view a
    shared scratch buffer: hold ``_scratch_lock`` until they are consumed (``npz_bytes`` / ``npz_b64_bytes``)."""
    mod, default_level = _deflate_backend()
    if level is None:
        level = int(os.environ.get("MOGE_NPZ_LEVEL", str(default_level)))
    pool = _get_pool()
    ct = isinstance(mod, _IsalCt)
    plan = []
    for name, a in arrays.items():
        a = np.asanyarray(a)
        if not a.flags.c_contiguous:
            a = a.copy(order="C")      # (np.ascontiguousarray would turn a 0-d array into 1-d)
        hdr = io.BytesIO()
        np.lib.format.write_array_header_1_0(hdr, np.lib.format.header_data_from_array_1_0(a))
        raw = memoryview(a.reshape(-1).view(np.uint8)) if a.size else memoryview(b"")
        plan.append((name, [hdr.getvalue()] + [raw[i:i + _CHUNK] for i in range(0, len(raw), _CHUNK)]))
    # isalct: the compressed pieces land in one persistent scratch buffer (no per-piece allocation; the
    # caller holds _scratch_lock until the returned segments are consumed)
    scratch = _scratch(sum(mod.bound(len(p)) for _, ps in plan for p in ps)) if ct else None
    soff = 0
    members = []
    for name, pieces in plan:
        futs = []
        for j, p in enumerate(pieces):
            dst = None
            if ct:
                cap = mod.bound(len(p))
                dst, soff = scratch[soff:soff + cap], soff + cap
            futs.append(pool.submit(_deflate_piece, mod, p, level, j == len(pieces) - 1, dst))
        members.append((name + ".npy", sum(len(p) for p in pieces), pool.submit(_crc, mod, pieces), futs))
    out, central = [], []
    off = 0
    for fname, usize, crc_f, futs in members:
        comp = [f.result() for f in futs]
        csize = sum(memoryview(c).nbytes for c in comp)
        crc = crc_f.result()
        nb = fname.encode()
        # local file header: version 2.0, no flags, deflate, DOS time 0 / date 1980-01-01
        lh = struct.pack("<IHHHHHIIIHH", 0x04034B50, 20, 0, 8, 0, 0x21, crc, csize, usize, len(nb), 0) + nb
        central.append(struct.pack("<IHHHHHHIIIHHHHHII", 0x02014B50, 20, 20, 0, 8, 0, 0x21, crc, csize, usize,
                                   len(nb), 0, 0, 0, 0, 0o600 << 16, off) + nb)
        out.append(lh)
        out.extend(comp)
        off += len(lh) + csize
    cd = b"".join(central)
    out.append(cd)
    out.append(struct.pack("<IHHHHIIH", 0x06054B50, 0, 0, len(central), len(central), len(cd), off, 0))
    return out


_scratch_buf = np.empty(0, np.uint8)
_scratch_lock = threading.RLock()


def _scratch(n: int) -> np.ndarray:
    """Persistent output buffer of the parallel deflate (callers hold ``_scratch_lock`` until the
    segments it backs are consumed)."""
    global _scratch_buf
    if _scratch_buf.size < n:
        _scratch_buf = np.empty(n, np.uint8)
    return _scratch_buf


def npz_bytes(arrays: Dict[str, np.ndarray], level: Optional[int] = None) -> bytes:
    with _scratch_lock:
        return b"".join(npz_segments(arrays, level))


class Spliced:
    """A base64 value written at ``buf[start:end]`` of a larger buffer with ``start`` bytes of headroom in
    front and ``buf.size - end`` behind (the JSON body is assembled around it in place)."""

    def __init__(self, buf: np.ndarray, start: int, end: int):
        self.buf, self.start, self.end = buf, start, end

    def tobytes(self) -> bytes:
        return self.buf[self.start:self.end].tobytes()


_body_bufs: list = []
_BODY_POOL = int(os.environ.get("MOGE_BODY_POOL", "3"))


def _body_buffer(n: int) -> np.ndarray:
    """uint8 [n] for a response body: a view of a pooled buffer that no live response references any
    more (a 39 MB fresh allocation is mmap'd and page-faults on first touch: ~1.7 ms of the base64 pass).
    A buffer is reused only when its refcount shows no view of it alive (bodies are memoryviews of it)."""
    import sys
    with _pool_lock:
        for a in _body_bufs:
            if a.size >= n and sys.getrefcount(a) <= 3:     # the list + `a` + getrefcount's argument
                return a[:n]
        a = np.empty(max(n + n // 16, 1 << 20), np.uint8)
        if len(_body_bufs) < _BODY_POOL:
            _body_bufs.append(a)
        else:   # replace a free (too small) pooled buffer
            for i in range(len(_body_bufs)):
                if sys.getrefcount(_body_bufs[i]) <= 2:      # the list + getrefcount's argument
                    _body_bufs[i] = a
                    break
        return a[:n]


def b64_segments(segs, headroom: int = 0):
    """``b64(b"".join(segs))`` without the join (multithreaded C++ over the segment list). With
    ``headroom`` > 0 the result is a ``Spliced`` with that many free bytes before and after the value."""
    import ctypes
    views = [np.frombuffer(memoryview(x).cast("B"), np.uint8) for x in segs]
    views = [v for v in views if v.size]
    n = sum(v.size for v in views)
    m = 4 * ((n + 2) // 3)
    res = _body_buffer(m + 2 * headroom) if headroom else np.empty(m, np.uint8)
    if n:
        ptrs = (ctypes.c_void_p * len(views))(*[v.ctypes.data for v in views])
        lens = (ctypes.c_int64 * len(views))(*[v.size for v in views])
        host_ext.lib().moge_b64_segs(ptrs, lens, len(views), res.ctypes.data + headroom, nthreads())
    return Spliced(res, headroom, headroom + m) if headroom else res


def b64(data) -> np.ndarray:
    """base64 (standard alphabet, padded) as an ASCII uint8 array; multithreaded C++ (the output buffer is
    np.empty: a zero-filled bytearray of ~48 MB alone costs ~20 ms)."""
    mv = memoryview(data).cast("B")
    n = len(mv)
    res = np.empty(4 * ((n + 2) // 3), np.uint8)
    if n:
        src = np.frombuffer(mv, np.uint8)
        host_ext.lib().moge_b64(src.ctypes.data, n, res.ctypes.data, nthreads())
    return res


def npz_b64_bytes(arrays: Dict[str, np.ndarray], level: Optional[int] = None, headroom: int = 0):
    with _scratch_lock:
        return b64_segments(npz_segments(arrays, level), headroom)


def npz_b64(arrays: Dict[str, np.ndarray], level: Optional[int] = None) -> str:
    return npz_b64_bytes(arrays, level).tobytes().decode("ascii")


def tune_malloc() -> str:
    """MOGE_MALLOPT=1 (default): keep freed multi-MB blocks in the glibc heap instead of mmap/munmap per
    allocation (M_MMAP_THRESHOLD 32 MB = glibc's max, M_TRIM_THRESHOLD 256 MB): the served path allocates
    and frees tens of MB per request (decoded image, response arrays, npz / base64 buffers) and every fresh
    mmap'd page is a page fault on first touch (~0.25 us per 4 KB). Only allocation placement changes."""
    if os.environ.get("MOGE_MALLOPT", "1").strip().lower() in ("0", "false", "no", "off"):
        return "off"
    try:
        libc = ctypes.CDLL("libc.so.6")
        ok = libc.mallopt(-3, 32 << 20) == 1 and libc.mallopt(-1, 256 << 20) == 1   # M_MMAP_/M_TRIM_THRESHOLD
        return "mmap 32 MB / trim 256 MB" if ok else "mallopt refused"
    except OSError as e:
        return f"unavailable ({e})"


# ------------------------------------------------------------------------------------- request decode

_PNG_SIG = b"\x89PNG\r\n\x1a\n"


def decode_png_rgb(raw: bytes) -> Optional[np.ndarray]:
    """uint8 [H, W, 3] of an 8-bit, non-interlaced RGB / RGBA PNG (== PIL ``open().convert('RGB')``:
    PNG is lossless and the alpha channel is dropped), or None for anything else (the caller then uses
    PIL). Chunk CRCs are verified; IDAT is inflated with ISA-L (zlib fallback), unfiltered in C++."""
    if bytes(memoryview(raw)[:8]) != _PNG_SIG:
        return None
    try:
        from PIL import Image
        mod = _inflate_backend()
        p, n = 8, len(raw)
        mv = memoryview(raw)
        ihdr, idat, idat_off, end = None, [], [], False
        while p + 12 <= n:
            ln, = struct.unpack_from(">I", raw, p)
            typ = bytes(mv[p + 4:p + 8])
            if p + 12 + ln > n:
                return None
            data = mv[p + 8:p + 8 + ln]
            crc, = struct.unpack_from(">I", raw, p + 8 + ln)
            if mod.crc32(data, mod.crc32(typ)) & 0xFFFFFFFF != crc:
                return None
            if typ == b"IHDR":
                ihdr = struct.unpack(">IIBBBBB", data)
            elif typ == b"IDAT":
                idat.append(data)
                idat_off.append((p + 8, ln))
            elif typ == b"IEND":
                end = True
                break
            elif typ == b"PLTE" or not (typ[0] & 0x20):   # unknown critical chunk -> let PIL decide
                return None
            p += 12 + ln
        if ihdr is None or not idat or not end:
            return None
        W, H, depth, ctype, comp, filt, interlace = ihdr
        if depth != 8 or ctype not in (2, 6) or comp or filt or interlace or W < 1 or H < 1:
            return None
        if W * H > (Image.MAX_IMAGE_PIXELS or 1 << 62):
            return None
        bpp = 3 if ctype == 2 else 4
        rowbytes = W * bpp
        out = np.empty((H, W, 3), np.uint8)
        isal = _isal_inflate_fns()
        if isal is not None:
            rc = _png_isal_ct(isal, raw, idat_off, H, rowbytes, bpp, out)
            return out if rc == 0 else None
        if _png_unfilter_pipelined(mod, idat, H, rowbytes, bpp, out):
            return out
        return None
    except Exception:   # noqa: BLE001 -- anything odd: PIL decides (and reports the error)
        return None


_unf_pool: Optional[ThreadPoolExecutor] = None
_isal_fns = None          # None = not probed yet; False = unavailable
_png_tls = threading.local()


def _png_isal_ct(fns, raw, idat_off, H: int, rowbytes: int, bpp: int, out: np.ndarray) -> int:
    """IDAT inflate (ISA-L, C++, no GIL) pipelined with the C++ unfilter: moge_png_inflate_unfilter."""
    L = host_ext.lib()
    total = H * (rowbytes + 1)
    z = getattr(_png_tls, "z", None)
    if z is None or z.size < total:       # per-thread persistent scratch (no page faults per request)
        z = _png_tls.z = np.empty(total, np.uint8)
    base = np.frombuffer(raw, np.uint8).ctypes.data
    k = len(idat_off)
    ptrs = (ctypes.c_void_p * k)(*[base + o for o, _ in idat_off])
    lens = (ctypes.c_int64 * k)(*[ln for _, ln in idat_off])
    times = _png_times if _PNG_TIMES else None
    return L.moge_png_inflate_unfilter(fns[0], fns[1], ptrs, lens, k, H, rowbytes, bpp, z.ctypes.data,
                                       out.ctypes.data, times.ctypes.data if times is not None else None)


_PNG_TIMES = os.environ.get("MOGE_DECODE_LOG") == "1"      # opt tool: inflate / unfilter end times (us)
_png_times = np.zeros(2, np.float64)


def _isal_inflate_fns():
    """(isal_inflate_init, isal_inflate) addresses from python-isal's shared object for the C++ PNG decode,
    or None (MOGE_PNG_ISALCT=0, no python-isal / host_ext, ISA-L other than 2.31, or the self test fails)."""
    global _isal_fns
    if _isal_fns is None:
        _isal_fns = False
        try:
            if os.environ.get("MOGE_PNG_ISALCT", "1").strip().lower() in ("0", "false", "no", "off"):
                raise ImportError("disabled")
            from isal import _isal, isal_zlib
            if not str(_isal.ISAL_VERSION).startswith("2.31."):
                raise ImportError(f"inflate_state layout unverified for ISA-L {_isal.ISAL_VERSION}")
            if not host_ext.enabled():
                raise ImportError("host_ext unavailable")
            so = ctypes.CDLL(isal_zlib.__file__)
            fns = (ctypes.cast(so.isal_inflate_init, ctypes.c_void_p).value,
                   ctypes.cast(so.isal_inflate, ctypes.c_void_p).value)
            if not _isal_selftest(fns):
                raise ImportError("self test failed")
            _isal_fns = fns
        except Exception:   # noqa: BLE001 -- fall back to the python-isal decompressobj path
            _isal_fns = False
    return _isal_fns or None


def _isal_selftest(fns) -> bool:
    """All 5 filter types, RGB + RGBA, multi-chunk IDAT, vs zlib + the serial C++ unfilter; a corrupted
    adler32 must be rejected."""
    rng = np.random.default_rng(0)
    L = host_ext.lib()
    for bpp in (3, 4):
        H, W = 37, 53
        rowbytes = W * bpp
        rows = rng.integers(0, 256, (H, rowbytes + 1), dtype=np.uint8)
        rows[:, 0] = np.arange(H) % 5
        rows[5:9, 1:] = rows[4, 1:]          # some compressible runs
        comp = zlib.compress(rows.tobytes(), 6)
        ref = np.empty((H, W, 3), np.uint8)
        if L.moge_png_unfilter(rows.ctypes.data, H, rowbytes, bpp, 3, ref.ctypes.data) != 0:
            return False
        cut = [0, 7, 100, len(comp) // 2, len(comp)]
        raw = bytes(comp)
        offs = [(cut[i], cut[i + 1] - cut[i]) for i in range(len(cut) - 1)]
        out = np.empty((H, W, 3), np.uint8)
        if _png_isal_ct(fns, raw, offs, H, rowbytes, bpp, out) != 0 or not np.array_equal(out, ref):
            return False
        bad = bytearray(comp)
        bad[-1] ^= 1
        if _png_isal_ct(fns, bytes(bad), [(0, len(bad))], H, rowbytes, bpp, out) == 0:
            return False
    return True
_PNG_ROWS = 96       # rows per unfilter step


def _png_unfilter_pipelined(mod, idat, H: int, rowbytes: int, bpp: int, out: np.ndarray) -> bool:
    """Inflate the IDAT stream piecewise on this thread while a worker thread unfilters the rows that are
    complete (ctypes releases the GIL); same bytes as inflating everything first and unfiltering after
    (the unfilter is sequential over rows and keeps its previous row in its own state)."""
    global _unf_pool
    if _unf_pool is None:
        _unf_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="moge-png")
    L = host_ext.lib()
    stride = rowbytes + 1
    total = H * stride
    z = np.empty(total, np.uint8)
    st = L.moge_unf_new(rowbytes, bpp, 3)
    futs = []
    try:
        d = mod.decompressobj()
        have, done = 0, 0
        rows_out = out.reshape(H, -1)

        def flush(upto: int):
            nonlocal done
            if upto > done:
                futs.append(_unf_pool.submit(L.moge_unf_rows, st, z[done * stride:].ctypes.data, upto - done,
                                             rows_out[done].ctypes.data))
                done = upto
        for chunk in idat:
            piece = d.decompress(chunk)
            n = len(piece)
            if have + n > total:
                n = total - have        # (trailing bytes beyond the image are ignored, as before)
            z[have:have + n] = np.frombuffer(piece, np.uint8, n)
            have += n
            if have // stride - done >= _PNG_ROWS:
                flush(have // stride)
        piece = d.flush()
        n = min(len(piece), total - have)
        z[have:have + n] = np.frombuffer(piece, np.uint8, n)
        have += n
        if have < total or not d.eof:    # (truncated stream: PIL decides, as with the one-shot inflate)
            return False
        flush(H)
        return all(f.result() == 0 for f in futs)
    finally:
        for f in futs:
            f.result()
        L.moge_unf_free(st)


_PyUnicode_AsUTF8AndSize = None


def _ascii_view(s: str):
    """(address, length) of an ASCII str's own UTF-8 data (== its compact ASCII storage, no copy)."""
    global _PyUnicode_AsUTF8AndSize
    if _PyUnicode_AsUTF8AndSize is None:
        f = ctypes.pythonapi.PyUnicode_AsUTF8AndSize
        f.argtypes = [ctypes.py_object, ctypes.POINTER(ctypes.c_ssize_t)]
        f.restype = ctypes.c_void_p
        _PyUnicode_AsUTF8AndSize = f
    n = ctypes.c_ssize_t()
    addr = _PyUnicode_AsUTF8AndSize(s, ctypes.byref(n))
    return addr, n.value


def b64decode_view(s):
    """``b64decode`` into a per-thread persistent buffer: a memoryview valid until this thread's next call
    (no 2-3 MB fresh allocations per request; an ASCII str is read in place). Same bytes as b64decode."""
    if isinstance(s, str) and s.isascii() and len(s) >= 4096 and host_ext.enabled():
        addr, n = _ascii_view(s)
        if addr and n == len(s):
            buf = getattr(_png_tls, "b64", None)
            need = n // 4 * 3
            if buf is None or buf.size < need:
                buf = _png_tls.b64 = np.empty(max(need, 1 << 20), np.uint8)
            k = host_ext.lib().moge_b64dec(addr, n, buf.ctypes.data, nthreads())
            if k >= 0:
                return memoryview(buf)[:k]
    return memoryview(b64decode(s))


def b64decode(s) -> bytes:
    """``base64.b64decode(s, validate=False)``: multithreaded C++ for a clean standard-alphabet string,
    Python's decoder for anything else (which skips foreign characters / raises on bad padding)."""
    import base64
    data = s.encode("ascii") if isinstance(s, str) and s.isascii() else None
    if data is not None and len(data) >= 4096 and host_ext.enabled():
        buf = np.empty(len(data) // 4 * 3, np.uint8)
        src = np.frombuffer(data, np.uint8)
        n = host_ext.lib().moge_b64dec(src.ctypes.data, src.size, buf.ctypes.data, nthreads())
        if n >= 0:
            return buf[:n].tobytes()
    return base64.b64decode(s, validate=False)


def _deflate_parallel(mod, data: memoryview, level: int):
    """Raw-deflate pieces of one stream (futures, in order)."""
    pool = _get_pool()
    pieces = [data[i:i + _CHUNK] for i in range(0, len(data), _CHUNK)] or [data]
    return [pool.submit(_deflate_piece, mod, p, level, j == len(pieces) - 1) for j, p in enumerate(pieces)]


def png_bytes(arr: np.ndarray, level: Optional[int] = None) -> bytes:
    """Lossless PNG of a uint8 [H,W] / [H,W,3] or uint16 [H,W] array (decodes to exactly ``arr``, like
    PIL's ``Image.fromarray(arr).save(PNG)``): 'Up' filter on every row (vectorized), the zlib stream
    deflated in parallel pieces (pigz construction), adler32 trailer."""
    mod, default_level = _deflate_backend()
    if level is None:
        level = int(os.environ.get("MOGE_NPZ_LEVEL", str(default_level)))
    if arr.dtype == np.uint16 and arr.ndim == 2:
        depth, ctype, rows = 16, 0, np.ascontiguousarray(arr.astype(">u2")).view(np.uint8)
    elif arr.dtype == np.uint8 and arr.ndim == 2:
        depth, ctype, rows = 8, 0, np.ascontiguousarray(arr)
    elif arr.dtype == np.uint8 and arr.ndim == 3 and arr.shape[2] == 3:
        depth, ctype, rows = 8, 2, np.ascontiguousarray(arr).reshape(arr.shape[0], -1)
    else:
        raise ValueError(f"png_bytes: unsupported {arr.dtype} {arr.shape}")
    H, W = arr.shape[:2]
    rows = rows.reshape(H, -1)
    filt = np.empty((H, rows.shape[1] + 1), np.uint8)
    filt[:, 0] = 2                                   # Up (row 0: the previous row is all zero)
    filt[0, 1:] = rows[0]
    np.subtract(rows[1:], rows[:-1], out=filt[1:, 1:])
    raw = memoryview(filt.reshape(-1))
    futs = _deflate_parallel(mod, raw, level)
    adler = zlib.adler32(raw)
    comp = b"".join(memoryview(f.result()).cast("B") for f in futs)

    def chunk(typ: bytes, data: bytes) -> bytes:
        return struct.pack(">I", len(data)) + typ + data + struct.pack(">I", zlib.crc32(data, zlib.crc32(typ)) & 0xFFFFFFFF)

    ihdr = struct.pack(">IIBBBBB", W, H, depth, ctype, 0, 0, 0)
    idat = b"\x78\x01" + comp + struct.pack(">I", adler & 0xFFFFFFFF)
    return _PNG_SIG + chunk(b"IHDR", ihdr) + chunk(b"IDAT", idat) + chunk(b"IEND", b"")


def fit_to_canvas(rgb: np.ndarray, canvas: Tuple[int, int], fit: str) -> Tuple[torch.Tensor, Region]:
    """``postprocess.fit_to_canvas`` on a uint8 [H, W, 3] array in one multithreaded C++ pass (Pillow's
    BILINEAR resize arithmetic, paste on the gray canvas, / 255 as float32, planar NCHW output);
    bit-identical (tests/test_serve_fastpath.py)."""
    from tt_moge.server import postprocess as pp
    cw, ch = canvas
    H, W = rgb.shape[:2]
    if fit == "stretch":
        rw, rh, x0, y0 = cw, ch, 0, 0
    elif fit == "pad":
        s = min(cw / W, ch / H)
        rw, rh = max(1, round(W * s)), max(1, round(H * s))
        x0, y0 = (cw - rw) // 2, (ch - rh) // 2
    else:
        raise ValueError(f"unknown fit mode {fit!r}")
    rgb = np.ascontiguousarray(rgb, dtype=np.uint8)
    out = torch.empty(1, 3, ch, cw, dtype=torch.float32)
    pad = float(np.float32(pp.PAD_GRAY[0]) / np.float32(255.0))
    assert pp.PAD_GRAY[0] == pp.PAD_GRAY[1] == pp.PAD_GRAY[2]
    rc = host_ext.lib().moge_fit_canvas(rgb.ctypes.data, H, W, rgb.strides[0], rw, rh, x0, y0, cw, ch, pad,
                                        out.data_ptr(), nthreads())
    if rc != 0:
        raise ValueError("moge_fit_canvas: bad geometry")
    return out, (x0, y0, x0 + rw, y0 + rh)


def fit_to_canvas_u8(rgb: np.ndarray, canvas: Tuple[int, int], fit: str) -> Tuple[torch.Tensor, Region]:
    """Round 12 (MOGE_SERVE_U8): ``fit_to_canvas`` without the float32 canvas: the uint8 canvas [1, Ch, Cw, 3]
    (NHWC) whose u8 / 255 IS fit_to_canvas's float32 output (TtMoGe's input prep reads it that way: the same
    model input bit for bit). A canvas-sized image is passed through without a copy."""
    from tt_moge.server import postprocess as pp
    cw, ch = canvas
    H, W = rgb.shape[:2]
    if fit == "stretch":
        rw, rh, x0, y0 = cw, ch, 0, 0
    elif fit == "pad":
        s = min(cw / W, ch / H)
        rw, rh = max(1, round(W * s)), max(1, round(H * s))
        x0, y0 = (cw - rw) // 2, (ch - rh) // 2
    else:
        raise ValueError(f"unknown fit mode {fit!r}")
    rgb = np.ascontiguousarray(rgb, dtype=np.uint8)
    region = (x0, y0, x0 + rw, y0 + rh)
    if (rw, rh, x0, y0) == (W, H, 0, 0) and (W, H) == (cw, ch):
        return torch.from_numpy(rgb)[None], region            # identity: the decoded pixels are the canvas
    out = torch.empty(1, ch, cw, 3, dtype=torch.uint8)
    assert pp.PAD_GRAY[0] == pp.PAD_GRAY[1] == pp.PAD_GRAY[2]
    rc = host_ext.lib().moge_fit_canvas_u8(rgb.ctypes.data, H, W, rgb.strides[0], rw, rh, x0, y0, cw, ch,
                                           float(pp.PAD_GRAY[0]), out.data_ptr(), nthreads())
    if rc != 0:
        raise ValueError("moge_fit_canvas_u8: bad geometry")
    return out, region