from __future__ import annotations

import cv2
import numpy as np

_FRONT = (0, 1, 2, 3)
_BACK = (5, 4, 7, 6)
_SIDES = (
    (4, 5, 1, 0),
    (3, 2, 6, 7),
    (4, 0, 3, 7),
    (1, 5, 6, 2),
)

_LIGHT = np.array([0.35, -0.55, -1.0])
_LIGHT /= np.linalg.norm(_LIGHT)

_CARD_BG = (40, 37, 34)
_CARD_BORDER = (120, 115, 110)
_EDGE_COLOR = (150, 145, 140)
_SIDE_COLOR = np.array([120, 116, 112], dtype=np.float32)


def _rotation(yaw: float, pitch: float, roll: float) -> np.ndarray:
    """Build rotation matrix."""
    cy, sy = np.cos(yaw), np.sin(yaw)
    cp, sp = np.cos(pitch), np.sin(pitch)
    cr, sr = np.cos(roll), np.sin(roll)
    ry = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]])
    rx = np.array([[1, 0, 0], [0, cp, -sp], [0, sp, cp]])
    rz = np.array([[cr, -sr, 0], [sr, cr, 0], [0, 0, 1]])
    return rz @ ry @ rx


class Object3D:
    """Textured 3D slab."""

    def __init__(self, drawing_bgra: np.ndarray, thickness_ratio: float = 0.05) -> None:
        self.texture = self._make_card(drawing_bgra)
        h, w = self.texture.shape[:2]
        self.aspect = w / max(h, 1)
        self.thickness_ratio = thickness_ratio
        self._corners = self._build_corners()

    @staticmethod
    def _make_card(drawing_bgra: np.ndarray) -> np.ndarray:
        """Flatten drawing to card."""
        if drawing_bgra.shape[2] == 4:
            rgb = drawing_bgra[:, :, :3].astype(np.float32)
            alpha = (drawing_bgra[:, :, 3].astype(np.float32) / 255.0)[:, :, None]
        else:
            rgb = drawing_bgra.astype(np.float32)
            alpha = np.ones(drawing_bgra.shape[:2] + (1,), dtype=np.float32)

        h, w = rgb.shape[:2]
        max_side = 900
        if max(h, w) > max_side:
            k = max_side / max(h, w)
            rgb = cv2.resize(rgb, (int(w * k), int(h * k)), interpolation=cv2.INTER_AREA)
            alpha = cv2.resize(alpha, (int(w * k), int(h * k)), interpolation=cv2.INTER_AREA)[:, :, None]

        bg = np.full(rgb.shape, _CARD_BG, dtype=np.float32)
        card = (bg * (1.0 - alpha) + rgb * alpha).astype(np.uint8)
        cv2.rectangle(card, (0, 0), (card.shape[1] - 1, card.shape[0] - 1), _CARD_BORDER, 2)
        return card

    def _build_corners(self) -> np.ndarray:
        """Eight slab vertices."""
        w = self.aspect * 0.5
        h = 0.5
        t = max(self.thickness_ratio, 0.005) * 0.5
        return np.array(
            [
                (-w, -h, t), (w, -h, t), (w, h, t), (-w, h, t),
                (-w, -h, -t), (w, -h, -t), (w, h, -t), (-w, h, -t),
            ],
            dtype=np.float64,
        )

    @staticmethod
    def _project(pts3d: np.ndarray, focal: float, distance: float, cx: float, cy: float):
        """Perspective project points."""
        z = pts3d[:, 2] + distance
        z = np.maximum(z, 1e-3)
        x = focal * pts3d[:, 0] / z + cx
        y = focal * pts3d[:, 1] / z + cy
        return np.stack([x, y], axis=1), z

    @staticmethod
    def _shade(face_cam: np.ndarray) -> float:
        """Face lighting factor."""
        n = np.cross(face_cam[1] - face_cam[0], face_cam[2] - face_cam[0])
        norm = np.linalg.norm(n)
        if norm < 1e-9:
            return 0.2
        n = n / norm
        if n[2] > 0:
            n = -n
        return float(np.clip(abs(np.dot(n, _LIGHT)), 0.22, 1.0))

    @staticmethod
    def _quad_area(quad: np.ndarray) -> float:
        """Screen area of quad."""
        x, y = quad[:, 0], quad[:, 1]
        return 0.5 * abs(np.dot(x, np.roll(y, -1)) - np.dot(y, np.roll(x, -1)))

    def _blit_texture(
        self, out: np.ndarray, texture: np.ndarray, quad: np.ndarray, shade: float
    ) -> None:
        """Warp texture onto face."""
        oh, ow = out.shape[:2]
        x0 = int(np.floor(quad[:, 0].min()))
        y0 = int(np.floor(quad[:, 1].min()))
        x1 = int(np.ceil(quad[:, 0].max())) + 1
        y1 = int(np.ceil(quad[:, 1].max())) + 1
        x0, y0 = max(0, x0), max(0, y0)
        x1, y1 = min(ow, x1), min(oh, y1)
        if x1 - x0 < 2 or y1 - y0 < 2:
            return

        th, tw = texture.shape[:2]
        src = np.array([[0, 0], [tw - 1, 0], [tw - 1, th - 1], [0, th - 1]], dtype=np.float32)
        dst = (quad - np.array([x0, y0], dtype=np.float64)).astype(np.float32)
        try:
            m = cv2.getPerspectiveTransform(src, dst)
        except cv2.error:
            return
        if not np.all(np.isfinite(m)):
            return

        size = (x1 - x0, y1 - y0)
        warped = cv2.warpPerspective(texture, m, size, flags=cv2.INTER_LINEAR)
        mask = cv2.warpPerspective(
            np.full((th, tw), 255, np.uint8), m, size, flags=cv2.INTER_NEAREST
        )
        warped = np.clip(warped.astype(np.float32) * (0.70 + 0.30 * shade), 0, 255).astype(np.uint8)

        roi = out[y0:y1, x0:x1]
        a = (mask.astype(np.float32) / 255.0)[:, :, None]
        out[y0:y1, x0:x1] = (roi * (1.0 - a) + warped * a).astype(np.uint8)

    def render(
        self,
        background: np.ndarray,
        yaw: float,
        pitch: float,
        roll: float = 0.0,
        scale: float = 1.0,
        distance: float = 3.0,
        center: tuple[float, float] | None = None,
        draw_edges: bool = True,
    ) -> np.ndarray:
        """Render object over background."""
        out = background.copy()
        oh, ow = out.shape[:2]
        cx, cy = center if center is not None else (ow * 0.5, oh * 0.5)
        focal = 1.7 * oh * max(scale, 0.05)

        rot = _rotation(yaw, pitch, roll)
        cam = self._corners @ rot.T
        proj, depth = self._project(cam, focal, distance, cx, cy)

        faces: list[tuple[float, tuple[int, ...], str]] = [
            (float(depth[list(_FRONT)].mean()), _FRONT, "front"),
            (float(depth[list(_BACK)].mean()), _BACK, "back"),
        ]
        faces += [(float(depth[list(f)].mean()), f, "side") for f in _SIDES]
        faces.sort(key=lambda item: item[0], reverse=True)

        back_tex = None
        for _, idx, kind in faces:
            centre = cam[list(idx)].mean(axis=0)
            view = centre + np.array([0.0, 0.0, distance])
            if float(np.dot(centre, view)) >= 0.0:
                continue

            quad = proj[list(idx)]
            if self._quad_area(quad) < 4.0:
                continue
            shade = self._shade(cam[list(idx)])

            if kind == "front":
                self._blit_texture(out, self.texture, quad, shade)
            elif kind == "back":
                if back_tex is None:
                    back_tex = cv2.flip(self.texture, 1)
                    back_tex = (back_tex.astype(np.float32) * 0.55).astype(np.uint8)
                self._blit_texture(out, back_tex, quad, shade)
            else:
                color = np.clip(_SIDE_COLOR * shade, 0, 255)
                cv2.fillConvexPoly(
                    out, quad.astype(np.int32), color.tolist(), lineType=cv2.LINE_AA
                )

            if draw_edges:
                cv2.polylines(
                    out, [quad.astype(np.int32)], True, _EDGE_COLOR, 1, cv2.LINE_AA
                )
        return out