import base64
import numpy as np
from fastapi.testclient import TestClient
from vlm.webconsole.qa_console.server import create_app


class FakeSource:
    is_mock = True
    is_connected = False

    def get_latest_frame(self):
        return np.zeros((48, 64, 3), dtype=np.uint8)


class FakeEngine:
    loaded = True

    def stream_infer(self, frame, prompt):
        yield "I see "
        yield "a phone [100, 100, 500, 500]."


def make_client():
    app = create_app(FakeSource(), FakeEngine())
    return TestClient(app)


def test_root_serves_html():
    r = make_client().get("/")
    assert r.status_code == 200
    assert "<html" in r.text.lower()


def test_status_endpoint():
    r = make_client().get("/status")
    assert r.json() == {"connected": False, "mock": True}


def test_ws_streams_tokens_and_image():
    client = make_client()
    with client.websocket_connect("/ws") as ws:
        ws.send_json({"prompt": "find the phone"})
        msgs = []
        while True:
            m = ws.receive_json()
            msgs.append(m)
            if m["type"] in ("done", "error"):
                break
    types = [m["type"] for m in msgs]
    assert "token" in types
    assert types[-1] == "done"
    tokens = "".join(m["text"] for m in msgs if m["type"] == "token")
    assert "phone" in tokens
    images = [m for m in msgs if m["type"] == "image"]
    assert len(images) == 1
    raw = base64.b64decode(images[0]["data"])
    assert raw[:2] == b"\xff\xd8"


def test_ws_no_box_no_image():
    class NoBoxEngine(FakeEngine):
        def stream_infer(self, frame, prompt):
            yield "nothing here"

    app = create_app(FakeSource(), NoBoxEngine())
    client = TestClient(app)
    with client.websocket_connect("/ws") as ws:
        ws.send_json({"prompt": "anything"})
        msgs = []
        while True:
            m = ws.receive_json()
            msgs.append(m)
            if m["type"] in ("done", "error"):
                break
    assert not any(m["type"] == "image" for m in msgs)
    assert msgs[-1]["type"] == "done"