# SPDX-License-Identifier: Apache-2.0
"""Image decode and preprocessing shared by the Python API (``tt_sam3.api``) and the HTTP server
(``tt_sam3.server.app``).  Moved here unchanged from ``server/app.py`` so that the Python API does not need
fastapi / pydantic; ``server/app.py`` imports these exact functions under their old names.

Only numpy / torch / PIL are needed (``deflate`` is optional: PNG fast path).
"""
from __future__ import annotations

import struct
import warnings
from dataclasses import dataclass
from typing import Optional, Tuple

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

INPUT_SIZE = 1008  # fixed: pos_embed / SDPA chunking are precomputed for the 72x72 patch grid


@dataclass
class DecodedImage:
    rgb: np.ndarray  # (H, W, 3) uint8, possibly a view into the decoded RGBA buffer

    @property
    def size(self) -> Tuple[int, int]:  # PIL order: (width, height)
        return int(self.rgb.shape[1]), int(self.rgb.shape[0])


try:  # libdeflate inflate for the PNG fast path (demo 3.6 MB PNG: 25 ms zlib-ng in PIL -> 10 ms)
    import deflate as _libdeflate
except ImportError:  # optional dependency
    _libdeflate = None

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


def _png_fast_decode(raw: bytes) -> Optional[np.ndarray]:
    """8-bit non-interlaced RGB / RGBA PNGs whose scanlines use filter types 0 (None), 1 (Sub) or 2 (Up): chunk parse +
    one libdeflate inflate + numpy unfilter -> (H, W, 3|4) uint8, the same pixels as PIL.  Returns None for anything else
    (palette / gray / 16-bit / interlaced / tRNS / APNG / Avg or Paeth rows / malformed data), which then takes the PIL
    path with PIL's semantics.  Like PIL, chunk CRCs are not verified."""
    if _libdeflate is None or not raw.startswith(_PNG_SIG):
        return None
    pos, n_raw, ihdr, idat = 8, len(raw), None, []
    while pos + 12 <= n_raw:
        length = struct.unpack_from(">I", raw, pos)[0]
        tag = raw[pos + 4 : pos + 8]
        if pos + 12 + length > n_raw:
            return None
        if tag == b"IHDR":
            if length != 13 or ihdr is not None:
                return None
            ihdr = struct.unpack_from(">IIBBBBB", raw, pos + 8)
        elif tag == b"IDAT":
            idat.append(raw[pos + 8 : pos + 8 + length])
        elif tag == b"IEND":
            break
        elif tag in (b"tRNS", b"PLTE", b"acTL", b"fcTL", b"fdAT"):
            return None
        elif ihdr is None:
            return None
        pos += 12 + length
    else:
        return None  # no IEND
    if ihdr is None or not idat:
        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 == 0 or h == 0 or w * h > Image.MAX_IMAGE_PIXELS:
        return None
    bpp = 3 if ctype == 2 else 4
    stride = w * bpp
    try:
        data = _libdeflate.zlib_decompress(b"".join(idat), h * (stride + 1))
    except Exception:  # noqa: BLE001 - corrupt / short / long stream: PIL decides
        return None
    if len(data) != h * (stride + 1):
        return None
    rows = np.frombuffer(data, dtype=np.uint8).reshape(h, stride + 1)
    ftype = rows[:, 0]
    px = rows[:, 1:]
    if not ftype.any():
        return px.reshape(h, w, bpp)  # read-only view, no copy
    if (ftype > 2).any():
        return None  # Avg / Paeth (or invalid) rows: sequential per pixel, PIL's C unfilter
    out = np.empty((h, stride), dtype=np.uint8)
    prev = np.zeros(stride, dtype=np.uint8)
    for r in range(h):
        f = ftype[r]
        if f == 0:
            out[r] = px[r]
        elif f == 1:  # Sub: running sum per channel (uint8 wraps mod 256)
            np.cumsum(px[r].reshape(w, bpp), axis=0, dtype=np.uint8, out=out[r].reshape(w, bpp))
        else:  # Up
            np.add(px[r], prev, out=out[r])
        prev = out[r]
    return out.reshape(h, w, bpp)


def to_pixel_values(im) -> torch.Tensor:
    """RGB uint8 (DecodedImage or PIL) -> (1, 3, 1008, 1008) float32 in [-1, 1] (bilinear squash, mean=std=0.5),
    as the port validated.  Same arithmetic as ``preprocess_image(arr / 255)`` (fp32 x / 255, the same
    ``F.interpolate`` call on the same contiguous NCHW tensor, (x - 0.5) / 0.5) with fewer passes / allocations:
    the values are <= 1 by construction (no max() scan) and the normalisation runs in place."""
    rgb = im.rgb if isinstance(im, DecodedImage) else np.asarray(im.convert("RGB"))
    h, w = rgb.shape[:2]
    with warnings.catch_warnings():  # PIL's buffer is read-only; it is only read here
        warnings.simplefilter("ignore", UserWarning)
        u8 = torch.from_numpy(rgb)
    # uint8 HWC (strided RGBA view) -> contiguous uint8 CHW -> fp32 / 255: equal to numpy's float32 x / 255.0 for
    # all 256 byte values (checked exhaustively in perf/dev/host_path_check.py)
    chw = u8.permute(2, 0, 1).contiguous().to(torch.float32).div_(255.0).unsqueeze(0)
    if (h, w) != (INPUT_SIZE, INPUT_SIZE):
        chw = F.interpolate(chw, size=(INPUT_SIZE, INPUT_SIZE), mode="bilinear", align_corners=False)
    return chw.sub_(0.5).div_(0.5)