# SPDX-License-Identifier: Apache-2.0
"""Round 14 (P1): native host preprocess -- ``resize.c`` (crop + grey pad + PIL BILINEAR box resize + mirror in one call,
bit-identical to the PIL path), built lazily with the host C compiler into a per-user cache and loaded with ctypes.

Nothing here is required: if there is no compiler, the build fails, or the load-time self-check (random crops vs the PIL
path, bit for bit) fails, :func:`crop_resize` is ``None`` and preprocess keeps the PIL path (same bytes, slower).
``HAMER_PRE_NATIVE=0`` disables it.
"""
from __future__ import annotations

import ctypes
import hashlib
import os
import subprocess
import tempfile
import threading
from pathlib import Path
from typing import Optional

import numpy as np

_SRCS = [Path(__file__).with_name("resize.c")]
# no -ffast-math, no FP contraction: the filter coefficients must be computed with Pillow's exact double arithmetic
_CFLAGS = ["-O3", "-march=native", "-ffp-contract=off", "-fno-fast-math", "-shared", "-fPIC"]
_LOCK = threading.Lock()
_FN = None
_FN_PS = None  # round 15: hamer_crop_resize_ps (any source pixel stride) when the zero-copy PIL path self-checked
_IM2COL = None
_TRIED = False
STATUS = "not loaded"


def _cache_dir() -> Path:
    d = os.environ.get("HAMER_NATIVE_DIR")
    if d:
        return Path(d)
    base = os.environ.get("XDG_CACHE_HOME") or str(Path.home() / ".cache")
    return Path(base) / "hamer-native"


def _build() -> Optional[Path]:
    src = b"".join(p.read_bytes() for p in _SRCS)
    cc = os.environ.get("CC", "cc")
    tag = hashlib.sha1(src + " ".join(_CFLAGS).encode() + os.uname().machine.encode()).hexdigest()[:16]
    for d in (_cache_dir(), Path(tempfile.gettempdir()) / f"hamer-native-{os.getuid()}"):
        so = d / f"hamer_resize_{tag}.so"
        if so.exists():
            return so
        try:
            d.mkdir(parents=True, exist_ok=True)
            tmp = d / f".{so.name}.{os.getpid()}"
            subprocess.run([cc, *_CFLAGS, "-o", str(tmp), *map(str, _SRCS)], check=True, capture_output=True, timeout=120)
            os.replace(tmp, so)
            return so
        except Exception:  # noqa: BLE001 - no compiler / read-only dir: try the next dir, then give up
            continue
    return None


def _self_check(fn) -> bool:
    """The native path must reproduce the PIL path bit for bit (random frames / boxes / both hands, edge-crossing boxes)."""
    from PIL import Image

    from hamer import preprocess as pp

    rng = np.random.default_rng(1234)
    for i in range(24):
        H, W = (int(v) for v in rng.integers(48, 400, 2))
        a = rng.integers(0, 256, (H, W, 3), dtype=np.uint8)
        img = Image.fromarray(a)
        w, h = rng.uniform(4, W), rng.uniform(4, H)
        x1, y1 = rng.uniform(-w / 2, W - w / 2), rng.uniform(-h / 2, H - h / 2)
        bbox, rf, right = (x1, y1, x1 + w, y1 + h), float(rng.choice([2.0, 1.0, 3.3])), bool(i % 2)
        if pp.hamer_bbox_size(bbox, rf) < 2:
            continue
        crop, meta = pp._square_crop(img, bbox, rf)  # noqa: SLF001
        ref = pp.resize_crop_hw(crop, right, bands=1)
        got = _call(fn, a, bbox, rf, right)
        got_chw = _call(fn, a, bbox, rf, right, chw=True)
        if got is None or not np.array_equal(ref, got[0]) or got[1] != meta:
            return False
        if got_chw is None or not np.array_equal(ref.transpose(2, 0, 1), got_chw[0]) or got_chw[1] != meta:
            return False
    return True


def _call(fn, a: np.ndarray, bbox, rescale_factor: float, is_right: bool, chw: bool = False):
    """(u8 HWC [256, 192, 3] or planar CHW [3, 256, 192] when chw, (cx, cy, bbox_size)) or None when the PIL path must handle it (degenerate boxes)."""
    if a.ndim != 3 or a.shape[2] != 3 or a.dtype != np.uint8 or a.strides[1:] != (3, 1):
        return None
    H, W = a.shape[:2]
    return _call_ptr(fn, a.ctypes.data, H, W, a.strides[0], 3, bbox, rescale_factor, is_right, chw)


def _call_ptr(fn, ptr: int, H: int, W: int, row_stride: int, ps: int, bbox, rescale_factor: float, is_right: bool,
              chw: bool = False):
    """_call on a raw u8 image pointer with ``ps`` bytes per pixel (3: packed RGB; 4: Pillow's RGBX rows)."""
    from hamer.preprocess import hamer_bbox_size

    x1, y1, x2, y2 = bbox
    cx, cy = (x1 + x2) / 2.0, (y1 + y2) / 2.0
    bs = hamer_bbox_size(bbox, rescale_factor)
    half = bs / 2.0
    S = int(round(bs))
    # exactly _square_crop's integer boxes
    l, t = max(0, int(round(cx - half))), max(0, int(round(cy - half)))
    r, b = min(W, int(round(cx + half))), min(H, int(round(cy + half)))
    px, py = max(0, int(round(-(cx - half)))), max(0, int(round(-(cy - half))))
    if S < 2 or r <= l or b <= t:
        return None
    out = np.empty((3, 256, 192) if chw else (256, 192, 3), np.uint8)
    args = (S, l, t, r - l, b - t, px, py, 0 if is_right else 1, 1 if chw else 0, out.ctypes.data)
    rc = fn(ptr, H, W, row_stride, *args) if ps == 3 and fn is not _FN_PS else fn(ptr, H, W, row_stride, ps, *args)
    if rc != 0:
        return None
    return out, (cx, cy, bs)


def load():
    """The ctypes function (or None).  Built / checked once per process."""
    global _FN, _FN_PS, _IM2COL, _TRIED, STATUS
    with _LOCK:
        if _TRIED:
            return _FN
        _TRIED = True
        pre_off = os.environ.get("HAMER_PRE_NATIVE", "1").strip().lower() in ("0", "false", "no", "off")
        try:
            so = _build()
            if so is None:
                STATUS = "no compiler / build failed: PIL path"
                return None
            lib = ctypes.CDLL(str(so))
            fn = lib.hamer_crop_resize
            fn.restype = ctypes.c_int
            fn.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_long] + [ctypes.c_int] * 9 + [ctypes.c_void_p]
            ic = lib.hamer_im2col_tile_bf16
            ic.restype = ctypes.c_int
            ic.argtypes = [ctypes.c_void_p, ctypes.c_long, ctypes.c_long, ctypes.c_long, ctypes.c_int, ctypes.c_int,
                           ctypes.c_int, ctypes.c_void_p]
            _IM2COL = ic  # checked by its caller (TtHamer._setup_fast_io) against the torch path
            if pre_off:
                STATUS = "off (HAMER_PRE_NATIVE=0)"
                return None
            if not _self_check(fn):
                STATUS = "self-check FAILED (not bit-identical to PIL): PIL path"
                return None
            _FN = fn
            STATUS = f"on ({so.name}, self-check bit-identical)"
            if os.environ.get("HAMER_PRE_ZEROCOPY", "1").strip().lower() not in ("0", "false", "no", "off"):
                fps = lib.hamer_crop_resize_ps
                fps.restype = ctypes.c_int
                fps.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_long, ctypes.c_int] \
                    + [ctypes.c_int] * 9 + [ctypes.c_void_p]
                _FN_PS = fps
                if _self_check_zc():
                    STATUS += "; zero-copy PIL buffer on"
                else:
                    _FN_PS = None
                    STATUS += "; zero-copy PIL buffer off (self-check failed)"
        except Exception as e:  # noqa: BLE001
            STATUS = f"load failed ({e!r}): PIL path"
            _FN = None
        return _FN


def im2col_tile_fn():
    """Round 14: ``hamer_im2col_tile_bf16`` (u8 crop -> bf16 TILE im2col of the centred crop) or None."""
    load()
    return _IM2COL


def crop_resize(a: np.ndarray, bbox, rescale_factor: float, is_right: bool, chw: bool = False):
    fn = load()
    if fn is None:
        return None
    return _call(fn, a, bbox, rescale_factor, is_right, chw)


# ---------------------------------------------------------------------------------------------------------------------
# round 15: the decoded frame read IN PLACE (no np.asarray copy, ~0.08 ms per 512x334 frame).  Pillow stores an RGB image
# as 4-byte RGBX pixels; its Arrow C export (Pillow >= 11.2, a single-block image) hands out the pixel buffer pointer.
class _ArrowArray(ctypes.Structure):
    pass


_ArrowArray._fields_ = [("length", ctypes.c_int64), ("null_count", ctypes.c_int64), ("offset", ctypes.c_int64),
                        ("n_buffers", ctypes.c_int64), ("n_children", ctypes.c_int64),
                        ("buffers", ctypes.POINTER(ctypes.c_void_p)),
                        ("children", ctypes.POINTER(ctypes.POINTER(_ArrowArray))), ("dictionary", ctypes.c_void_p),
                        ("release", ctypes.c_void_p), ("private_data", ctypes.c_void_p)]
_CAP_GET = ctypes.pythonapi.PyCapsule_GetPointer
_CAP_GET.restype = ctypes.c_void_p
_CAP_GET.argtypes = [ctypes.py_object, ctypes.c_char_p]


def _rgbx_view(img):
    """(capsules to keep alive, pointer, H, W) of an RGB PIL image's RGBX pixel buffer, or None."""
    if img.mode != "RGB" or not hasattr(img, "__arrow_c_array__"):
        return None
    try:
        img.load()
        sch, arr = img.__arrow_c_array__()
        a = ctypes.cast(_CAP_GET(arr, b"arrow_array"), ctypes.POINTER(_ArrowArray)).contents
        W, H = img.size
        if a.n_children != 1 or a.length != H * W or a.offset != 0:
            return None
        c = a.children[0].contents
        if c.n_buffers != 2 or c.length != H * W * 4 or c.offset != 0 or not c.buffers[1]:
            return None
        return (sch, arr), c.buffers[1], H, W
    except Exception:  # noqa: BLE001 - older Pillow / multi-block image: the copying path
        return None


def _self_check_zc() -> bool:
    """The zero-copy path must give the bytes of the np.asarray path (random frames / boxes / both hands)."""
    from PIL import Image

    rng = np.random.default_rng(4321)
    for i in range(16):
        H, W = (int(v) for v in rng.integers(48, 400, 2))
        a = rng.integers(0, 256, (H, W, 3), dtype=np.uint8)
        img = Image.fromarray(a)
        w, h = rng.uniform(4, W), rng.uniform(4, H)
        x1, y1 = rng.uniform(-w / 2, W - w / 2), rng.uniform(-h / 2, H - h / 2)
        bbox, rf, right = (x1, y1, x1 + w, y1 + h), float(rng.choice([2.0, 1.0, 3.3])), bool(i % 2)
        ref = _call(_FN, a, bbox, rf, right, chw=True)
        v = _rgbx_view(img)
        if v is None:
            return False
        got = _call_ptr(_FN_PS, v[1], v[2], v[3], v[3] * 4, 4, bbox, rf, right, chw=True)
        if (ref is None) != (got is None) or (ref is not None and (not np.array_equal(ref[0], got[0]) or ref[1] != got[1])):
            return False
    return True


def crop_resize_image(img, bbox, rescale_factor: float, is_right: bool, chw: bool = False):
    """crop_resize on a decoded PIL image, reading its pixel buffer in place when possible (round 15)."""
    fn = load()
    if fn is None:
        return None
    if _FN_PS is not None:
        v = _rgbx_view(img)
        if v is not None:
            keep, ptr, H, W = v
            res = _call_ptr(_FN_PS, ptr, H, W, W * 4, 4, bbox, rescale_factor, is_right, chw)
            del keep
            return res
    return _call(fn, np.asarray(img), bbox, rescale_factor, is_right, chw)