"""Headless first-person renderer for the testbed room.

A small raycaster (Pillow-only, no GPU, no display) that produces an honest
frame: it derives every pixel from the actual world state in ``Room``. Used by
the ``vision`` and ``depth`` sensor tools.

Camera model: yaw = heading around Y, pitch = camera tilt. 90° horizontal FOV.
"""

from __future__ import annotations

import math

from PIL import Image, ImageDraw

from testbed.room.world import ROOM_HEIGHT, ROOM_SIZE, Room

# Palette
FOG = (14, 14, 20)
FLOOR_A = (56, 56, 62)
FLOOR_B = (64, 64, 70)
FLOOR_LINE = (42, 42, 48)
CEIL = (40, 40, 46)
SKY = (12, 16, 30)
WALL = (96, 96, 106)
WALL_TOP_BAND = (70, 70, 78)

BEACON_OFF = ((0.85, 1.0, 1.0), (0.42, 0.82, 0.92), (0.10, 0.45, 0.62))
BEACON_ON = ((1.0, 0.98, 0.78), (1.0, 0.75, 0.35), (0.78, 0.43, 0.0))

MAX_VIEW = 20.0
_EPS = 1e-9


class Raycaster:
    """Renders the room from the capsule's camera. Stateless per frame."""

    def __init__(self, width: int = 160, height: int = 120):
        self.width = width
        self.height = height
        self.tan_fx = 1.0  # 90° horizontal FOV
        self.tan_fy = self.tan_fx * height / width

    def render(
        self, room: Room, max_depth: float = 10.0
    ) -> tuple[Image.Image, list[list[float]]]:
        """Render the current view. Returns (RGB image, aligned depth grid)."""
        w, h = self.width, self.height
        cap = room.capsule
        px = cap.x
        pz = cap.z
        eye = cap.eye_height
        yaw = math.radians(cap.yaw_deg)
        pitch = math.radians(cap.pitch_deg)
        tan_pitch = math.tan(pitch)

        fx = math.sin(yaw)
        fz = math.cos(yaw)
        rx = math.cos(yaw)
        rz = -math.sin(yaw)

        pixels = bytearray(w * h * 3)
        depth = [[MAX_VIEW] * w for _ in range(h)]

        slabs = self._slabs(room)

        for c in range(w):
            u = (2.0 * (c + 0.5) / w - 1.0) * self.tan_fx
            dx = fx + rx * u
            dz = fz + rz * u
            inv = 1.0 / math.hypot(dx, dz)
            dx *= inv
            dz *= inv

            dist, surf = self._cast(px, pz, dx, dz, slabs)

            if surf is None:
                t_wall = math.inf
                s_coord = 0.0
            else:
                t_wall = dist
                s_coord = surf["s"]

            row_top = self._row_of(ROOM_HEIGHT, t_wall, eye, tan_pitch)
            row_bot = self._row_of(0.0, t_wall, eye, tan_pitch)
            r_top = max(0, math.ceil(min(row_top, row_bot)))
            r_bot = min(h, math.floor(max(row_top, row_bot)))

            c * 3
            for r in range(h):
                v = self.tan_fy * (1.0 - 2.0 * (r + 0.5) / h)
                down = v + tan_pitch
                o = (r * w + c) * 3
                if r_top <= r < r_bot and t_wall < MAX_VIEW:
                    col = self._shade_wall(surf, dist, s_coord)
                    if r == r_top:
                        col = _lerp(col, WALL_TOP_BAND, 0.55)
                    t = dist
                elif down < -_EPS:
                    t = eye / -down
                    col = self._shade_floor(px + dx * t, pz + dz * t, t)
                elif down > _EPS:
                    t = (ROOM_HEIGHT - eye) / down
                    col = _fog(CEIL, t)
                else:
                    col = SKY
                    t = math.inf
                if t < MAX_VIEW:
                    depth[r][c] = t
                pixels[o] = col[0]
                pixels[o + 1] = col[1]
                pixels[o + 2] = col[2]

        self._draw_beacon(room, pixels, depth, fx, fz, rx, rz, eye, tan_pitch)

        img = Image.frombuffer("RGB", (w, h), bytes(pixels), "raw", "RGB", 0, 1)
        grid = self._depth_grid(depth, max_depth)
        return img, grid

    # ---------- helpers ----------

    def _slabs(self, room: Room) -> list[dict]:
        """Room walls as thin slabs + the obstacle boxes, all axis-aligned."""
        e = 0.01
        slabs: list[dict] = []
        slabs.append(
            {
                "x0": -e,
                "x1": ROOM_SIZE + e,
                "z0": -e,
                "z1": e,
                "kind": "wall",
                "id": "wall_north",
                "top": ROOM_HEIGHT,
                "color": WALL,
                "face": "z0",
            }
        )
        slabs.append(
            {
                "x0": -e,
                "x1": ROOM_SIZE + e,
                "z0": ROOM_SIZE - e,
                "z1": ROOM_SIZE + e,
                "kind": "wall",
                "id": "wall_south",
                "top": ROOM_HEIGHT,
                "color": WALL,
                "face": "z1",
            }
        )
        slabs.append(
            {
                "x0": -e,
                "x1": e,
                "z0": -e,
                "z1": ROOM_SIZE + e,
                "kind": "wall",
                "id": "wall_west",
                "top": ROOM_HEIGHT,
                "color": WALL,
                "face": "x0",
            }
        )
        slabs.append(
            {
                "x0": ROOM_SIZE - e,
                "x1": ROOM_SIZE + e,
                "z0": -e,
                "z1": ROOM_SIZE + e,
                "kind": "wall",
                "id": "wall_east",
                "top": ROOM_HEIGHT,
                "color": WALL,
                "face": "x1",
            }
        )
        for box in room.boxes:
            slabs.append(
                {
                    "x0": box.x_min(),
                    "x1": box.x_max(),
                    "z0": box.z_min(),
                    "z1": box.z_max(),
                    "kind": "box",
                    "id": box.id,
                    "top": box.height,
                    "color": box.color,
                    "face": "box",
                    "box": box,
                }
            )
        return slabs

    def _cast(self, ox: float, oz: float, dx: float, dz: float, slabs: list[dict]):
        """Nearest 2D hit of the ray against all slabs. Returns (dist, surface)."""
        best = math.inf
        best_surf = None
        for s in slabs:
            t = _ray_aabb(ox, oz, dx, dz, s["x0"], s["x1"], s["z0"], s["z1"])
            if t is not None and t < best:
                best = t
                best_surf = s
        if best_surf is None:
            return best, None
        hit_x = ox + dx * best
        hit_z = oz + dz * best
        face = best_surf["face"]
        if face == "z0" or face == "z1":
            s_coord = hit_x
        elif face == "x0" or face == "x1":
            s_coord = hit_z
        else:
            # Box: pick the face by comparing t to each plane.
            s_coord = (
                hit_x if best_surf["box"].half_w <= best_surf["box"].half_d else hit_z
            )
        return best, {"surf": best_surf, "s": s_coord, "hit_x": hit_x, "hit_z": hit_z}

    def _row_of(self, world_h: float, t: float, eye: float, tan_pitch: float) -> float:
        """Screen row (float) where a height `world_h` at distance `t` lands."""
        if math.isinf(t):
            return 0.0 if world_h > eye else float(self.height)
        vv = (world_h - eye) / t - tan_pitch
        return (1.0 - vv / self.tan_fy) / 2.0 * self.height

    def _shade_wall(
        self, surf: dict | None, dist: float, s: float
    ) -> tuple[int, int, int]:
        if surf is None:
            return SKY
        s_def = surf["surf"]
        color = s_def["color"]
        face = s_def["face"]
        shade = 1.0
        if face == "x0":
            shade = 0.9
        elif face == "x1":
            shade = 1.08
        elif face == "z1":
            shade = 1.0
        elif face == "z0":
            shade = 0.95
        # Concrete panel grid: subtle stripes every 1 m.
        if int(abs(s) * 1.0) % 2 == 0:
            shade *= 1.06
        col = (
            min(255, int(color[0] * shade)),
            min(255, int(color[1] * shade)),
            min(255, int(color[2] * shade)),
        )
        return _fog(col, dist)

    def _shade_floor(self, wx: float, wz: float, t: float) -> tuple[int, int, int]:
        cell = math.floor(wx) + math.floor(wz)
        col = FLOOR_A if cell % 2 == 0 else FLOOR_B
        fx = wx - math.floor(wx)
        fz = wz - math.floor(wz)
        if min(fx, 1.0 - fx, fz, 1.0 - fz) < 0.07:
            col = FLOOR_LINE
        return _fog(col, t)

    def _draw_beacon(
        self,
        room: Room,
        pixels: bytearray,
        depth: list[list[float]],
        fx: float,
        fz: float,
        rx: float,
        rz: float,
        eye: float,
        tan_pitch: float,
    ) -> None:
        """Billboard the beacon with a soft glow, occluded by the depth buffer."""
        w, h = self.width, self.height
        b = room.beacon
        rel_x = b.x - room.capsule.x
        rel_z = b.z - room.capsule.z
        along = rel_x * fx + rel_z * fz
        if along < 0.25:
            return
        right = rel_x * rx + rel_z * rz
        col_c = (right / along / self.tan_fx + 1.0) / 2.0 * w
        vv = (b.height - eye) / along - tan_pitch
        row_c = (1.0 - vv / self.tan_fy) / 2.0 * h
        r_px = b.radius / along / self.tan_fx * w / 2.0
        r_px = max(r_px, 1.2)
        glow = r_px * 3.4

        pulse = 1.0
        if b.active:
            pulse = 0.9 + 0.1 * math.sin(room.tick * 0.6)
        core, mid, outer = BEACON_ON if b.active else BEACON_OFF
        core = tuple(pulse * v for v in core)
        mid = tuple(pulse * v for v in mid)
        outer = tuple(pulse * v for v in outer)

        c0 = max(0, int(col_c - glow))
        c1 = min(w, int(col_c + glow) + 1)
        r0 = max(0, int(row_c - glow))
        r1 = min(h, int(row_c + glow) + 1)
        inv_r2 = 1.0 / (r_px * r_px)
        for r in range(r0, r1):
            for c in range(c0, c1):
                if depth[r][c] <= along - 0.06:
                    continue  # occluded by a nearer surface
                dc = c - col_c
                dr = r - row_c
                d2 = (dc * dc + dr * dr) * inv_r2
                if d2 > glow * glow * inv_r2:
                    continue
                if d2 <= 1.0:
                    a = 1.0 - d2 * 0.55
                    col = core
                elif d2 <= 4.0:
                    f = (d2 - 1.0) / 3.0
                    a = 0.85 * (1.0 - f)
                    col = tuple(core[i] + (mid[i] - core[i]) * f for i in range(3))
                else:
                    f = (d2 - 4.0) / (glow * glow * inv_r2 - 4.0)
                    a = 0.5 * (1.0 - f)
                    col = tuple(mid[i] + (outer[i] - mid[i]) * f for i in range(3))
                o = (r * w + c) * 3
                base = (pixels[o], pixels[o + 1], pixels[o + 2])
                out = (
                    int(base[0] * (1 - a) + col[0] * 255 * a),
                    int(base[1] * (1 - a) + col[1] * 255 * a),
                    int(base[2] * (1 - a) + col[2] * 255 * a),
                )
                pixels[o] = min(255, out[0])
                pixels[o + 1] = min(255, out[1])
                pixels[o + 2] = min(255, out[2])

    def _depth_grid(
        self, depth: list[list[float]], max_depth: float
    ) -> list[list[float]]:
        """Downsample the per-pixel depth to 1/4 resolution (40x30 at 160x120)."""
        w, h = self.width, self.height
        gw, gh = w // 4, h // 4
        grid = [[0.0] * gw for _ in range(gh)]
        for r in range(gh):
            for c in range(gw):
                total = 0.0
                n = 0
                for rr in range(r * 4, r * 4 + 4):
                    row = depth[rr]
                    for cc in range(c * 4, c * 4 + 4):
                        t = row[cc]
                        if t >= MAX_VIEW or t > max_depth:
                            continue
                        total += t
                        n += 1
                grid[r][c] = round(total / n, 2) if n else 0.0
        return grid


def render_topdown(room: Room, size: int = 400) -> Image.Image:
    """Bird's-eye view of the whole room (debug helper; uses world state).

    Floor grid, crates, beacon (cyan / warm yellow when active) and the
    capsule as a circle with a heading arrow.
    """
    scale = size / ROOM_SIZE
    img = Image.new("RGB", (size, size), (52, 52, 58))
    draw = ImageDraw.Draw(img)

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

    for i in range(int(ROOM_SIZE) + 1):
        c = 60 if i % 2 == 0 else 54
        draw.line([xy(i, 0), xy(i, ROOM_SIZE)], fill=(c, c, c + 6), width=1)
        draw.line([xy(0, i), xy(ROOM_SIZE, i)], fill=(c, c, c + 6), width=1)

    for box in room.boxes:
        xa, za = xy(box.x_min(), box.z_min())
        xb, zb = xy(box.x_max(), box.z_max())
        draw.rectangle(
            [min(xa, xb), min(za, zb), max(xa, xb), max(za, zb)],
            fill=box.color,
            outline=(20, 20, 26),
            width=2,
        )

    b = room.beacon
    bx, bz = xy(b.x, b.z)
    color = (255, 200, 90) if b.active else (120, 210, 235)
    glow = (255, 240, 180) if b.active else (160, 235, 250)
    r = b.radius * scale * 2.2
    draw.ellipse([bx - r * 2.4, bz - r * 2.4, bx + r * 2.4, bz + r * 2.4], fill=glow)
    draw.ellipse(
        [bx - r, bz - r, bx + r, bz + r], fill=color, outline=(20, 20, 26), width=2
    )

    c = room.capsule
    cx, cz = xy(c.x, c.z)
    rad = c.radius * scale
    draw.ellipse(
        [cx - rad, cz - rad, cx + rad, cz + rad],
        fill=(240, 240, 250),
        outline=(20, 20, 26),
        width=2,
    )
    fx, fz = c.forward()
    tip = xy(c.x + fx * 0.7, c.z + fz * 0.7)
    draw.line([(cx, cz), tip], fill=(30, 30, 40), width=3)
    return img


def _ray_aabb(
    ox: float,
    oz: float,
    dx: float,
    dz: float,
    x0: float,
    x1: float,
    z0: float,
    z1: float,
) -> float | None:
    """Slab test in 2D. Returns entry distance or None."""
    if abs(dx) < _EPS:
        if ox < x0 - _EPS or ox > x1 + _EPS:
            return None
        tx0, tx1 = -math.inf, math.inf
    else:
        tx0, tx1 = (x0 - ox) / dx, (x1 - ox) / dx
        if tx0 > tx1:
            tx0, tx1 = tx1, tx0
    if abs(dz) < _EPS:
        if oz < z0 - _EPS or oz > z1 + _EPS:
            return None
        tz0, tz1 = -math.inf, math.inf
    else:
        tz0, tz1 = (z0 - oz) / dz, (z1 - oz) / dz
        if tz0 > tz1:
            tz0, tz1 = tz1, tz0
    tin = max(tx0, tz0)
    tout = min(tx1, tz1)
    if tout < 0.0 or tin > tout:
        return None
    if tin > _EPS:
        return tin
    return None if tout <= _EPS else tout


def _fog(color: tuple[int, int, int], t: float) -> tuple[int, int, int]:
    k = min(1.0, t / MAX_VIEW)
    k = k**1.4
    return (
        int(color[0] + (FOG[0] - color[0]) * k),
        int(color[1] + (FOG[1] - color[1]) * k),
        int(color[2] + (FOG[2] - color[2]) * k),
    )


def _lerp(
    a: tuple[int, int, int], b: tuple[int, int, int], k: float
) -> tuple[int, int, int]:
    return (
        int(a[0] + (b[0] - a[0]) * k),
        int(a[1] + (b[1] - a[1]) * k),
        int(a[2] + (b[2] - a[2]) * k),
    )


def render_preview(
    room: Room, path: str, width: int = 320, height: int = 240
) -> Image.Image:
    """Debug helper: render the current view and save it to a PNG file."""
    img, _ = Raycaster(width, height).render(room)
    img.save(path)
    return img