"""Synthetic source node: the test's stand-in for a camera and a detector.

Emits an animated frame, a depth map of the same scene, the intrinsics, the box
around the square moving in it, and a caption -- every input dora-image-view draws --
so dataflow.yml needs no camera, no detector and no network.

This is a test fixture, not an example of how to wire the node; see the README for
that. It is a dora node rather than a pytest case because what it exists to exercise
is the wiring and the event loop, which only a running dataflow covers. pytest does
not collect it: the filename does not start with `test_`.

The depth map describes the same scene as the image: a flat back wall with the square
standing in front of it. `camera_info` is published once, on the first tick, the way a
real camera publishes it once per session -- which is exactly the case the viewer has
to cope with by reading intrinsics from frame metadata as well.
"""

import os

import numpy as np
import pyarrow as pa
from dora import Node

from dora_image_view.decode import IMAGE_TYPE

WIDTH = int(os.getenv("SOURCE_WIDTH", "640"))
HEIGHT = int(os.getenv("SOURCE_HEIGHT", "480"))
# 0 runs until the viewer quits; a positive count is what lets CI run the
# dataflow to completion instead of killing it on a timeout.
FRAMES = int(os.getenv("SOURCE_FRAMES", "0"))

BOX = 96
SPEED = 6
CAPTION = "dora-image-view synthetic source"

WALL_MM = 2000.0
SQUARE_MM = 800.0

# A plausible pinhole for a frame this size.
FX = FY = float(WIDTH)
CX, CY = WIDTH / 2, HEIGHT / 2
DISTORTION = [0.0, 0.0, 0.0, 0.0, 0.0]
# Attached to every frame as well as published once, which is what dora-zed does
# and what lets a viewer that started late still learn the camera model.
INTRINSICS_METADATA = {"fx": FX, "fy": FY, "cx": CX, "cy": CY, "distortion": DISTORTION}


def main() -> None:
    """Send one image, depth map, bbox and text per tick, and camera_info once."""
    node = Node("source")
    sent = 0

    for event in node:
        if event["type"] != "INPUT":
            continue

        if sent == 0:
            node.send_output(
                "camera_info",
                pa.array(np.array([FX, FY, CX, CY], dtype=np.float32)),
                {"shape": [4], "dtype": "float32", "layout": "C", "distortion": DISTORTION},
            )

        frame, corners = render(sent)
        node.send_output("image", image_struct(frame, "bgr8"), INTRINSICS_METADATA)
        node.send_output(
            "bbox",
            pa.array([{
                "bbox": corners,
                "conf": np.array([0.99]),
                "labels": np.array(["square"]),
            }]),
            {"format": "xyxy"},
        )
        depth = render_depth(corners)
        node.send_output(
            "depth",
            image_struct(depth, "32FC1"),
            {"units": "mm", **INTRINSICS_METADATA},
        )
        node.send_output("text", pa.array([CAPTION]))

        sent += 1
        if FRAMES and sent >= FRAMES:
            break


def image_struct(frame: np.ndarray, encoding: str) -> pa.Array:
    """Pack a frame as std/media/v1/Image.

    The struct carries the width, height and encoding a receiver needs, so none
    of it has to travel as message metadata. `frame` may be any dtype: the
    encoding names how to read the bytes back.
    """
    return pa.array(
        [{
            "width": frame.shape[1],
            "height": frame.shape[0],
            "encoding": encoding,
            "data": frame.tobytes(),
        }],
        type=IMAGE_TYPE,
    )


def render(index: int) -> tuple[np.ndarray, np.ndarray]:
    """Draw frame `index` and return it with the xyxy corners of the square in it."""
    frame = np.zeros((HEIGHT, WIDTH, 3), dtype=np.uint8)
    frame[:, :, 0] = np.linspace(40, 140, HEIGHT, dtype=np.uint8)[:, None]

    travel = WIDTH - BOX
    # A triangle wave over `index`: the square starts at the left edge, runs to
    # the right edge, and comes back.
    left = travel - abs((index * SPEED) % (2 * travel) - travel)
    top = (HEIGHT - BOX) // 2
    frame[top : top + BOX, left : left + BOX] = (60, 200, 255)

    corners = np.array([left, top, left + BOX, top + BOX], dtype=np.float64)
    return frame, corners


def render_depth(corners: np.ndarray) -> np.ndarray:
    """Depth of the scene in `render`: a flat wall with the square in front of it."""
    depth = np.full((HEIGHT, WIDTH), WALL_MM, dtype=np.float32)
    left, top, right, bottom = (int(value) for value in corners)
    depth[top:bottom, left:right] = SQUARE_MM
    return depth


if __name__ == "__main__":
    main()