"""Top-down map drawing from sensor data.

Shared by the demo's frame recorder and the live viewer. Only uses what an
agent can observe: the world_query layout (room, obstacles, beacon) and
proprioception (position, heading).
"""

from __future__ import annotations

from typing import Any

from PIL import Image, ImageDraw

ROOM_SIZE = 16.0
MAP_SIZE = 320
SCALE = MAP_SIZE / ROOM_SIZE


def draw_sensor_map(
    layout: dict[str, Any] | None,
    pos: tuple[float, float],
    yaw: float,
    path: list[tuple[float, float]] | None = None,
    *,
    step: int | None = None,
    beacon_active: bool = False,
) -> Image.Image:
    """Draw the room from the agent's sensor view: floor grid, crates, the
    beacon, the capsule's path and current position/heading."""
    path = path or []

    def xy(x: float, z: float) -> tuple[float, float]:
        return (x * SCALE, MAP_SIZE - z * SCALE)

    img = Image.new("RGB", (MAP_SIZE, MAP_SIZE), (52, 52, 58))
    d = ImageDraw.Draw(img)
    for i in range(int(ROOM_SIZE) + 1):
        c = 60 if i % 2 == 0 else 54
        d.line([xy(i, 0), xy(i, ROOM_SIZE)], fill=(c, c, c + 6), width=1)
        d.line([xy(0, i), xy(ROOM_SIZE, i)], fill=(c, c, c + 6), width=1)

    if layout:
        for ob in layout.get("obstacles", []):
            x0, z0 = xy(ob["x"] - ob["width"] / 2, ob["z"] - ob["depth"] / 2)
            x1, z1 = xy(ob["x"] + ob["width"] / 2, ob["z"] + ob["depth"] / 2)
            d.rectangle(
                [min(x0, x1), min(z0, z1), max(x0, x1), max(z0, z1)],
                fill=(150, 90, 60),
                outline=(20, 20, 26),
                width=2,
            )
        b = layout.get("beacon", {})
        bx, bz = xy(b.get("x", 12.5), b.get("z", 12.5))
        core = (255, 200, 90) if beacon_active else (120, 210, 235)
        glow = (255, 240, 180) if beacon_active else (160, 235, 250)
        d.ellipse([bx - 14, bz - 14, bx + 14, bz + 14], fill=glow)
        d.ellipse(
            [bx - 7, bz - 7, bx + 7, bz + 7], fill=core, outline=(20, 20, 26), width=2
        )

    if len(path) > 1:
        pts = [xy(x, z) for x, z in path]
        d.line(pts, fill=(255, 170, 60), width=3)

    cx, cz = xy(*pos)
    d.ellipse(
        [cx - 7, cz - 7, cx + 7, cz + 7],
        fill=(240, 240, 250),
        outline=(20, 20, 26),
        width=2,
    )
    fx, fz = _forward(yaw)
    tip = xy(pos[0] + fx * 0.7, pos[1] + fz * 0.7)
    d.line([(cx, cz), tip], fill=(30, 30, 40), width=3)
    if step is not None:
        d.text(
            (8, 8),
            f"step {step}  pos ({pos[0]:.1f}, {pos[1]:.1f}) yaw {yaw:.0f}",
            fill=(230, 230, 230),
        )
    return img


def _forward(yaw_deg: float) -> tuple[float, float]:
    import math

    rad = math.radians(yaw_deg)
    return math.sin(rad), math.cos(rad)