#!/usr/bin/env python3
"""Host smoke test for the Tactility MCP Screen Phase 2 API."""

import argparse
import base64
import json
import socket
import struct
import urllib.request


def discover(timeout: float = 2.0) -> str | None:
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    sock.settimeout(timeout)
    try:
        sock.sendto(b"DISCOVER_SCREEN", ("255.255.255.255", 5000))
        data, address = sock.recvfrom(128)
        if data == b"SCREEN_IP_80":
            return address[0]
    except TimeoutError:
        return None
    finally:
        sock.close()
    return None


def rpc(ip: str, request_id: int, method: str, params: dict | None = None) -> dict:
    payload = {
        "jsonrpc": "2.0",
        "id": request_id,
        "method": method,
    }
    if params is not None:
        payload["params"] = params

    request = urllib.request.Request(
        f"http://{ip}/api/mcp",
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=5) as response:
        return json.loads(response.read())


def require_result(response: dict, label: str) -> None:
    if "error" in response:
        raise RuntimeError(f"{label} failed: {response['error']}")
    if "result" not in response:
        raise RuntimeError(f"{label} returned no result: {response}")


def make_bmp() -> bytes:
    width = 4
    height = 4
    row_size = width * 3
    pixels = bytearray()
    colors = [
        (0, 0, 255),
        (0, 255, 0),
        (255, 0, 0),
        (255, 255, 255),
    ]
    for y in range(height):
        for x in range(width):
            blue, green, red = colors[(x + y) % len(colors)]
            pixels.extend((blue, green, red))
        pixels.extend(b"\0" * ((4 - row_size % 4) % 4))

    offset = 54
    file_size = offset + len(pixels)
    header = (
        b"BM"
        + struct.pack("<IHHI", file_size, 0, 0, offset)
        + struct.pack("<IIIHHIIIIII", 40, width, height, 1, 24, 0, len(pixels), 0, 0, 0, 0)
    )
    return header + pixels


def raw_endpoint(ip: str) -> None:
    pixels = bytes(
        [
            0xF8, 0x00,
            0x07, 0xE0,
            0x00, 0x1F,
            0xFF, 0xFF,
        ]
    )
    request = urllib.request.Request(
        f"http://{ip}/api/screen/raw?x=2&y=2&w=2&h=2",
        data=pixels,
        headers={"Content-Type": "application/octet-stream"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=5) as response:
        if response.read() != b"OK":
            raise RuntimeError("Raw endpoint returned an unexpected response")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--ip", help="Device IP; omit to use UDP discovery")
    parser.add_argument(
        "--draw",
        action="store_true",
        help="Also clear the app canvas and draw visible test text",
    )
    args = parser.parse_args()

    ip = args.ip or discover()
    if not ip:
        raise SystemExit("MCP Screen was not discovered. Pass --ip DEVICE_IP to test directly.")

    print(f"Testing MCP Screen at {ip}")

    tools = rpc(ip, 1, "tools/list")
    require_result(tools, "tools/list")
    names = [tool["name"] for tool in tools["result"]["tools"]]
    expected = {"get_capabilities", "clear_screen", "draw_text"}
    if not expected.issubset(names):
        raise RuntimeError(f"Missing tools: {sorted(expected - set(names))}")
    print(f"tools/list: {', '.join(names)}")

    capabilities = rpc(
        ip,
        2,
        "tools/call",
        {"name": "get_capabilities", "arguments": {}},
    )
    require_result(capabilities, "get_capabilities")
    print("get_capabilities:", capabilities["result"]["content"][0]["text"])

    if args.draw:
        cleared = rpc(
            ip,
            3,
            "tools/call",
            {"name": "clear_screen", "arguments": {"color": 1}},
        )
        require_result(cleared, "clear_screen")

        drawn = rpc(
            ip,
            4,
            "tools/call",
            {
                "name": "draw_text",
                "arguments": {
                    "text": "Tactility MCP is alive",
                    "x": 20,
                    "y": 30,
                    "size": 2,
                },
            },
        )
        require_result(drawn, "draw_text")

        raw_rpc = rpc(
            ip,
            5,
            "tools/call",
            {
                "name": "draw_raw_rgb565",
                "arguments": {
                    "rgb565_base64": base64.b64encode(
                        bytes([0xFF, 0xE0, 0xF8, 0x1F, 0x07, 0xFF, 0x00, 0x00])
                    ).decode(),
                    "x": 20,
                    "y": 70,
                    "w": 4,
                    "h": 1,
                },
            },
        )
        require_result(raw_rpc, "draw_raw_rgb565")

        bmp = rpc(
            ip,
            6,
            "tools/call",
            {
                "name": "draw_color_bmp",
                "arguments": {
                    "bmp_base64": base64.b64encode(make_bmp()).decode(),
                    "x": 30,
                    "y": 90,
                },
            },
        )
        require_result(bmp, "draw_color_bmp")

        pbm_bytes = b"P4\n8 2\n" + bytes([0b10101010, 0b01010101])
        pbm = rpc(
            ip,
            7,
            "tools/call",
            {
                "name": "draw_image",
                "arguments": {
                    "pbm_base64": base64.b64encode(pbm_bytes).decode(),
                    "x": 50,
                    "y": 110,
                },
            },
        )
        require_result(pbm, "draw_image")

        raw_endpoint(ip)

        screenshot = rpc(
            ip,
            8,
            "tools/call",
            {"name": "get_screenshot", "arguments": {}},
        )
        require_result(screenshot, "get_screenshot")
        screenshot_text = screenshot["result"]["content"][0]["text"]
        if not screenshot_text.startswith("__PBM_BASE64__:"):
            raise RuntimeError("Screenshot did not return PBM data")
        base64.b64decode(screenshot_text.split(":", 1)[1], validate=True)
        print("Visible Phase 2 draw and screenshot tests sent.")

    print("Phase 2 smoke test passed.")


if __name__ == "__main__":
    main()