"""Vision tower + TinyVLM tests (CPU-only, no downloads)."""

import sys
from pathlib import Path

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

torch = pytest.importorskip("torch")

from qwenjev.vision import (  # noqa: E402
    BOS_ID,
    EOT_ID,
    PAD_ID,
    TinyVLM,
    TinyVLMConfig,
    grid_to_tensor,
    render_grid_pil,
    text_to_byte_ids,
)


GRID = [
    [1, 1, 1, 1, 1],
    [1, 2, 0, 3, 1],
    [1, 0, 4, 0, 1],
    [1, 1, 1, 1, 1],
]


def test_text_to_byte_ids_roundtrip():
    text = "STATE:\nqwenjev-state-v1"
    ids = text_to_byte_ids(text, max_len=512)
    assert ids[0] == BOS_ID and ids[-1] == EOT_ID
    body = bytes(ids[1:-1]).decode("utf-8")
    assert body == text


def test_text_truncation():
    ids = text_to_byte_ids("x" * 999, max_len=64)
    assert len(ids) == 64 and ids[-1] == EOT_ID


def test_grid_to_tensor_values():
    t = grid_to_tensor(GRID, image_size=64)
    assert t.shape == (3, 64, 64)
    assert 0.0 <= t.min().item() and t.max().item() <= 1.0
    # a coloured cell maps to a non-black pixel somewhere in the image
    assert t.sum().item() > 0


def test_grid_to_tensor_distinguishes_values():
    g1 = [[0, 0], [0, 0]]
    g2 = [[0, 1], [0, 0]]
    t1, t2 = grid_to_tensor(g1, 32), grid_to_tensor(g2, 32)
    assert not torch.allclose(t1, t2)


def test_render_grid_pil():
    pil = pytest.importorskip("PIL")
    img = render_grid_pil(GRID, cell_px=4)
    assert isinstance(img, pil.Image.Image)
    assert img.size == (5 * 4, 4 * 4)


def test_tinyvlm_forward_shapes():
    m = TinyVLM(TinyVLMConfig(image_size=64, patch_size=8, lm_dim=64))
    px = torch.stack([grid_to_tensor(GRID, 64), grid_to_tensor(GRID, 64)])
    ids = torch.tensor([
        text_to_byte_ids("abcd", 32),
        text_to_byte_ids("defg", 32),
    ])
    out, hs = m(px, ids, output_hidden_states=True)
    n_vis = m.n_vision_tokens
    assert out.shape == (2, n_vis + ids.shape[1], 64)
    # hidden states: [embed, after_layer1, after_layer2]
    assert len(hs) == 3
    assert all(h.shape == out.shape for h in hs)


def test_tinyvlm_all_finite_with_padding():
    """Padding + causal masks must never produce NaN/Inf (regression)."""
    torch.manual_seed(0)
    m = TinyVLM(TinyVLMConfig(image_size=64, patch_size=8, lm_dim=64))
    px = torch.stack([grid_to_tensor(GRID, 64)])
    long_ids = text_to_byte_ids("a long state text " * 3, 128)
    short_ids = text_to_byte_ids("tiny", 128)
    ids = torch.tensor([long_ids, short_ids + [PAD_ID] * (len(long_ids) - len(short_ids))])
    attn = torch.tensor([
        [1] * len(long_ids),
        [1] * len(short_ids) + [0] * (len(long_ids) - len(short_ids)),
    ])
    out, hs = m(px.expand(2, -1, -1, -1), ids, attention_mask=attn,
                output_hidden_states=True)
    for h in hs:
        assert torch.isfinite(h).all(), "non-finite hidden states (mask bug)"
    assert torch.isfinite(out).all()


def test_tinyvlm_pooled_uses_last_text_position():
    m = TinyVLM(TinyVLMConfig(image_size=64, patch_size=8, lm_dim=64))
    px = torch.stack([grid_to_tensor(GRID, 64)])
    ids = torch.tensor([text_to_byte_ids("hello", 32)])
    out, _ = m(px, ids)
    n_vis = m.n_vision_tokens
    pooled = m.pooled(out, torch.tensor([ids.shape[1]]) + n_vis)
    assert torch.allclose(pooled[0], out[0, n_vis + ids.shape[1] - 1])


def test_tinyvlm_vision_changes_representation():
    """The vision tower MUST matter: different frames -> different pooled."""
    torch.manual_seed(3)
    m = TinyVLM(TinyVLMConfig(image_size=64, patch_size=8, lm_dim=64)).eval()
    g1 = [[0] * 8 for _ in range(8)]
    g2 = [[0] * 8 for _ in range(8)]
    g2[3][3] = 7
    ids = torch.tensor([text_to_byte_ids("state", 32)])
    with torch.no_grad():
        o1, _ = m(torch.stack([grid_to_tensor(g1, 64)]), ids)
        o2, _ = m(torch.stack([grid_to_tensor(g2, 64)]), ids)
    assert not torch.allclose(o1[0, 0], o2[0, 0])