# SPDX-License-Identifier: Apache-2.0
"""Host tests for server/hostinput.py: the fused resize + normalise + window-order im2col band kernels give exactly the
bits of the reference path (app._to_pixel_values + fused_tables.im2col_window_order), for any image size and for the
buffer layouts the decoders return (contiguous RGB, RGB view of RGBA, PNG rows with the filter byte)."""
import base64
import os

import numpy as np
import pytest

from tt_sam3.server import hostinput as H

pytestmark = pytest.mark.skipif(not H.available(), reason="numba not installed")
MODEL_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))


def _same(rgb):
    return np.array_equal(H.fused_bands(rgb).view(np.uint32), H.reference_bands(rgb).view(np.uint32))


def test_self_check():
    assert H.self_check(0) and H.self_check(5)


@pytest.mark.parametrize("h,w", [(1, 1), (2, 3), (14, 14), (1007, 1009), (1009, 1008), (4096, 31), (77, 3000), (1080, 1920)])
def test_sizes(h, w):
    rng = np.random.default_rng(h * 7919 + w)
    assert _same(rng.integers(0, 256, size=(h, w, 3), dtype=np.uint8))


def test_png_row_layout():
    """The PNG fast path returns an (H, W, 4) view whose rows are W * 4 + 1 bytes apart (filter byte), RGB = [..., :3]."""
    rng = np.random.default_rng(11)
    h, w = 601, 803
    rows = rng.integers(0, 256, size=(h, w * 4 + 1), dtype=np.uint8)
    rgba = rows[:, 1:].reshape(h, w, 4)
    rgb = rgba[..., :3]
    assert rgb.strides == (w * 4 + 1, 4, 1)
    assert _same(rgb)
    assert _same(np.ascontiguousarray(rgb))
    # non-unit pixel stride other than 3 / 4 -> contiguous copy path
    assert _same(rng.integers(0, 256, size=(300, 400, 5), dtype=np.uint8)[..., 1:4])


def test_extremes():
    for v in (0, 255):
        assert _same(np.full((1080, 1920, 3), v, dtype=np.uint8))
    a = np.zeros((500, 700, 3), dtype=np.uint8)
    a[::2, ::3] = 255
    assert _same(a)


def test_demo_image():
    from tt_sam3.server import app

    b64 = base64.b64encode(open(os.path.join(MODEL_DIR, "media", "source_1.png"), "rb").read()).decode("ascii")
    assert _same(app._decode_image(b64).rgb)


def test_pad_columns_untouched():
    out = np.full((9, 576, 608), np.float32(7.0))
    bw = H.BandWriter(np.zeros((600, 800, 3), dtype=np.uint8))
    for b in range(3):
        bw.write(b, out[3 * b : 3 * b + 3])
    assert (out[..., 588:] == 7.0).all() and (out[..., :588] == -1.0).all()