import os
import socket
import sys

import numpy as np
import pytest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import treppe_gl as t


def test_encode_frame_known_bytes():
    rgbw = np.zeros((16, 4), dtype=np.float32)
    rgbw[0] = [1.0, 0.0, 0.5, 0.0]
    frame = t.encode_frame(rgbw)
    assert len(frame) == 2 + 16 * 8
    assert frame[:2] == bytes([0x42, 0x0A])
    # r=65535, g=0, b=round(0.5*65535)=32768, w=0 big endian
    assert frame[2:10] == bytes([0xFF, 0xFF, 0, 0, 0x80, 0x00, 0, 0])


def test_white_modes():
    px = np.array([[1.0, 0.5, 0.25, 1.0]] + [[0, 0, 0, 0]] * 15)
    assert t.to_rgbw(px, "extract")[0].tolist() == [0.75, 0.25, 0.0, 0.25]
    assert t.to_rgbw(px, "alpha")[0].tolist() == [1.0, 0.5, 0.25, 1.0]
    assert t.to_rgbw(px, "off")[0].tolist() == [1.0, 0.5, 0.25, 0.0]


@pytest.fixture(scope="module")
def renderer():
    ctx, _ = t.create_context(window=False)
    return t.Renderer(ctx)


def test_strip_pixels(renderer):
    renderer.load("""
    void mainImage(out vec4 c, in vec2 p) {
        c = vec4(p.x / iResolution.x, 0.0, iTime, 1.0);
    }""")
    px = renderer.render_strip(0.25)
    assert px.shape == (16, 4)
    np.testing.assert_allclose(px[:, 0], (np.arange(16) + 0.5) / 16, atol=1e-5)
    np.testing.assert_allclose(px[:, 2], 0.25, atol=1e-5)


def test_bad_shader_keeps_old(renderer):
    renderer.load("void mainImage(out vec4 c, in vec2 p) { c = vec4(1.0); }")
    with pytest.raises(Exception):
        renderer.load("void mainImage(out vec4 c, in vec2 p) { nope }")
    px = renderer.render_strip(0.0)
    np.testing.assert_allclose(px[:, 0], 1.0)


def test_udp_loopback():
    rx = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    rx.bind(("127.0.0.1", 0))
    rx.settimeout(2)
    tx = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    tx.sendto(t.BLACK_FRAME, rx.getsockname())
    assert rx.recv(2048) == t.BLACK_FRAME


def test_brightness_scales_and_clamps():
    px = np.array([[1.0, 0.5, 0.25, 1.0]] + [[0, 0, 0, 0]] * 15)
    half = t.to_rgbw(px, "extract", 50)[0]
    np.testing.assert_allclose(half, [0.375, 0.125, 0.0, 0.125])
    assert t.to_rgbw(px, "extract", 0).max() == 0
    np.testing.assert_allclose(t.to_rgbw(px, "extract", 500),
                               t.to_rgbw(px, "extract", 100))


CANVAS_SHADER = """
void mainImage(out vec4 c, in vec2 p) {
    // red in the two lowest rows, and in the two leftmost columns (step 0)
    c = vec4(p.y < 2.0 ? 1.0 : 0.0, p.x < 2.0 ? 1.0 : 0.0, 0.0, 1.0);
}"""


def test_canvas_compaction_modes():
    ctx, _ = t.create_context(window=False)
    r = t.Renderer(ctx)
    r.load(CANVAS_SHADER)
    r.configure((32, 4), "mean")      # 2 columns per step, 4 rows
    px = r.render_strip(0.0)
    assert px.shape == (16, 4)
    np.testing.assert_allclose(px[:, 0], 0.5)                 # 2 of 4 rows
    np.testing.assert_allclose(px[0, 1], 1.0)                 # step 0 green
    np.testing.assert_allclose(px[1:, 1], 0.0)
    r.configure((32, 4), "max")
    np.testing.assert_allclose(r.render_strip(0.0)[:, 0], 1.0)
    r.configure((32, 4), "row")                               # row 2 is dark
    np.testing.assert_allclose(r.render_strip(0.0)[:, 0], 0.0)
    r.configure(None)                                         # back to 1D
    assert r.render_strip(0.0).shape == (16, 4)


def test_canvas_parsing():
    assert t.parse_canvas("256x32") == (256, 32)
    for bad in ("100x10", "abc", "8x8", "256x0"):
        with pytest.raises(ValueError):
            t.parse_canvas(bad)
    assert t.directive("// x\n// canvas: 128x16 compact: max\nvoid f(){}") == (
        (128, 16), "max")
    assert t.directive("void mainImage() {}") == (None, None)
    with pytest.raises(ValueError):
        t.directive("// canvas: 128x16 compact: nope")