"""Embedded Open3D mesh renderer for Qt dashboards (offscreen or hidden-window capture)."""

from __future__ import annotations

import sys

import numpy as np

from orisys.finger.mesh import UVSurfaceMesh
from orisys.finger.viewer import flow_magnitude_vertex_colors, scalar_map_vertex_colors

from .contact_arrow import ContactArrowOverlay
from .mesh_camera import (
    apply_filament_sun_light,
    apply_shade_factors,
    apply_view_from_world_to_cam,
    camera_look_at_from_world_to_cam,
    precompute_shade_factors,
)


def _numpy_o3d_image_to_rgb(image) -> np.ndarray:
    arr = np.asarray(image)
    if arr.ndim == 3 and arr.shape[2] >= 3:
        if arr.dtype == np.float32 or arr.dtype == np.float64:
            return (np.clip(arr[:, :, :3], 0.0, 1.0) * 255.0).astype(np.uint8)
        return arr[:, :, :3].astype(np.uint8)
    return arr.astype(np.uint8)


def _make_lit_material(o3d):
    material = o3d.visualization.rendering.MaterialRecord()
    material.shader = "defaultLit"
    material.base_color = (1.0, 1.0, 1.0, 1.0)
    return material


def _shade_vertex_colors(mesh, colors: np.ndarray, shade_factors) -> np.ndarray:
    return apply_shade_factors(colors, shade_factors)


def _mesh_to_o3d(o3d, vertices: np.ndarray, triangles: np.ndarray, colors: np.ndarray):
    mesh = o3d.geometry.TriangleMesh()
    mesh.vertices = o3d.utility.Vector3dVector(vertices.astype(np.float64))
    mesh.triangles = o3d.utility.Vector3iVector(triangles.astype(np.int32))
    mesh.vertex_colors = o3d.utility.Vector3dVector(colors.astype(np.float64))
    mesh.compute_vertex_normals()
    return mesh


def _vertex_colors(mesh, *, color_mode, threshold, value_min, value_max, flow=None, scalar_map=None):
    if color_mode == "depth":
        return scalar_map_vertex_colors(
            scalar_map,
            mesh,
            value_min=value_min,
            value_max=value_max,
            threshold=threshold,
        )
    return flow_magnitude_vertex_colors(
        flow,
        mesh,
        threshold=threshold,
    )


class _FilamentOffscreenBackend:
    """Filament OffscreenRenderer (Linux headless / platforms with EGL)."""

    def __init__(
        self,
        mesh: UVSurfaceMesh,
        *,
        width: int,
        height: int,
        threshold: float,
        color_mode: str,
        value_min,
        value_max,
        world_to_cam: np.ndarray | None,
        lookat: np.ndarray,
        show_contact_arrow: bool,
        mesh_lighting=None,
        contact_arrow_style=None,
    ):
        import open3d as o3d

        self._o3d = o3d
        self.mesh = mesh
        self.width = width
        self.height = height
        self.threshold = threshold
        self.color_mode = color_mode
        self.value_min = value_min
        self.value_max = value_max
        self.world_to_cam = world_to_cam
        self.lookat = lookat
        self.mesh_lighting = mesh_lighting
        self._shade_factors = (
            precompute_shade_factors(mesh.rest_vertices, mesh.triangles, mesh_lighting)
            if mesh_lighting is not None
            else None
        )
        self._contact_arrow = ContactArrowOverlay(
            mesh,
            enabled=show_contact_arrow,
            style=contact_arrow_style,
        )
        self._renderer = o3d.visualization.rendering.OffscreenRenderer(self.width, self.height)
        self._renderer.scene.set_background([0.05, 0.05, 0.08, 1.0])
        if mesh_lighting is not None:
            apply_filament_sun_light(
                self._renderer,
                direction=mesh_lighting.direction,
                intensity=mesh_lighting.intensity,
            )
        self._material = _make_lit_material(o3d)
        self._mesh_added = False
        self._arrow_added = False
        self._apply_camera()

    def _apply_camera(self) -> None:
        if self.world_to_cam is not None:
            center, eye, up = camera_look_at_from_world_to_cam(self.world_to_cam, lookat=self.lookat)
        else:
            center = self.lookat.astype(np.float64)
            extent = np.ptp(self.mesh.rest_vertices, axis=0)
            span = float(np.max(extent))
            eye = center + np.array([0.0, 0.0, span * 2.0], dtype=np.float64)
            up = np.array([0.0, 1.0, 0.0], dtype=np.float64)
        self._renderer.scene.camera.look_at(center, eye, up)

    def resize(self, width: int, height: int) -> None:
        width = max(int(width), 64)
        height = max(int(height), 64)
        if width == self.width and height == self.height:
            return
        self.width = width
        self.height = height
        self._renderer = self._o3d.visualization.rendering.OffscreenRenderer(self.width, self.height)
        self._renderer.scene.set_background([0.05, 0.05, 0.08, 1.0])
        if self.mesh_lighting is not None:
            apply_filament_sun_light(
                self._renderer,
                direction=self.mesh_lighting.direction,
                intensity=self.mesh_lighting.intensity,
            )
        self._material = _make_lit_material(self._o3d)
        self._mesh_added = False
        self._arrow_added = False
        self._apply_camera()

    def render(self, *, flow=None, scalar_map=None, centroid=None, centroid_found=False, fnormal=0.0):
        o3d = self._o3d
        colors = _vertex_colors(
            self.mesh,
            color_mode=self.color_mode,
            threshold=self.threshold,
            value_min=self.value_min,
            value_max=self.value_max,
            flow=flow,
            scalar_map=scalar_map,
        )
        body_mesh = _mesh_to_o3d(o3d, self.mesh.rest_vertices, self.mesh.triangles, colors)

        if not self._mesh_added:
            self._renderer.scene.add_geometry("finger_mesh", body_mesh, self._material)
            self._mesh_added = True
        else:
            self._renderer.scene.remove_geometry("finger_mesh")
            self._renderer.scene.add_geometry("finger_mesh", body_mesh, self._material)

        arrow_result = self._contact_arrow.compute(centroid, centroid_found, fnormal)
        if arrow_result is not None and self._contact_arrow.arrow_mesh is not None:
            arrow_vertices, arrow_triangles, arrow_colors = arrow_result
            if np.any(arrow_vertices):
                arrow_mesh = _mesh_to_o3d(o3d, arrow_vertices, arrow_triangles, arrow_colors)
                if not self._arrow_added:
                    self._renderer.scene.add_geometry("contact_arrow", arrow_mesh, self._material)
                    self._arrow_added = True
                else:
                    self._renderer.scene.remove_geometry("contact_arrow")
                    self._renderer.scene.add_geometry("contact_arrow", arrow_mesh, self._material)
            elif self._arrow_added:
                self._renderer.scene.remove_geometry("contact_arrow")
                self._arrow_added = False

        return _numpy_o3d_image_to_rgb(self._renderer.render_to_image())

    def close(self) -> None:
        self._renderer = None


class _HiddenVisualizerBackend:
    """Invisible Open3D Visualizer + framebuffer capture (Windows-friendly)."""

    def __init__(
        self,
        mesh: UVSurfaceMesh,
        *,
        width: int,
        height: int,
        threshold: float,
        color_mode: str,
        value_min,
        value_max,
        world_to_cam: np.ndarray | None,
        lookat: np.ndarray,
        show_contact_arrow: bool,
        mesh_lighting=None,
        contact_arrow_style=None,
    ):
        import open3d as o3d

        self._o3d = o3d
        self.mesh = mesh
        self.width = width
        self.height = height
        self.threshold = threshold
        self.color_mode = color_mode
        self.value_min = value_min
        self.value_max = value_max
        self.world_to_cam = world_to_cam
        self.lookat = lookat
        self.mesh_lighting = mesh_lighting
        self._shade_factors = (
            precompute_shade_factors(mesh.rest_vertices, mesh.triangles, mesh_lighting)
            if mesh_lighting is not None
            else None
        )
        self._contact_arrow = ContactArrowOverlay(
            mesh,
            enabled=show_contact_arrow,
            style=contact_arrow_style,
        )
        self.vis = None
        self.tri_mesh = None
        self._create_window()

    def _create_window(self) -> None:
        o3d = self._o3d
        if self.vis is not None:
            self.vis.destroy_window()

        colors = _vertex_colors(
            self.mesh,
            color_mode=self.color_mode,
            threshold=self.threshold,
            value_min=self.value_min,
            value_max=self.value_max,
        )
        colors = _shade_vertex_colors(self.mesh, colors, self._shade_factors)
        self.vis = o3d.visualization.Visualizer()
        self.vis.create_window(
            window_name="Orisys Finger Mesh (embedded)",
            width=self.width,
            height=self.height,
            visible=False,
        )

        self.tri_mesh = o3d.geometry.TriangleMesh()
        self.tri_mesh.vertices = o3d.utility.Vector3dVector(self.mesh.rest_vertices.astype(np.float64))
        self.tri_mesh.triangles = o3d.utility.Vector3iVector(self.mesh.triangles.astype(np.int32))
        self.tri_mesh.vertex_colors = o3d.utility.Vector3dVector(colors.astype(np.float64))
        self.tri_mesh.compute_vertex_normals()
        self.vis.add_geometry(self.tri_mesh)

        if self._contact_arrow.arrow_mesh is not None:
            self.vis.add_geometry(self._contact_arrow.arrow_mesh)

        render_option = self.vis.get_render_option()
        render_option.mesh_show_back_face = True
        render_option.background_color = np.array([0.05, 0.05, 0.08])
        if hasattr(render_option, "light_on"):
            render_option.light_on = False

        if self.world_to_cam is not None:
            apply_view_from_world_to_cam(self.vis, self.world_to_cam, lookat=self.lookat)
        else:
            view_control = self.vis.get_view_control()
            center = self.lookat.astype(np.float64)
            extent = np.ptp(self.mesh.rest_vertices, axis=0)
            span = float(np.max(extent))
            view_control.set_lookat(center)
            view_control.set_front(np.array([0.0, 0.0, -1.0]))
            view_control.set_up(np.array([0.0, 1.0, 0.0]))
            view_control.set_zoom(0.5 if span > 0 else 0.7)

        self.vis.poll_events()
        self.vis.update_renderer()

    def resize(self, width: int, height: int) -> None:
        width = max(int(width), 64)
        height = max(int(height), 64)
        if width == self.width and height == self.height:
            return
        self.width = width
        self.height = height
        self._create_window()

    def render(self, *, flow=None, scalar_map=None, centroid=None, centroid_found=False, fnormal=0.0):
        o3d = self._o3d
        colors = _vertex_colors(
            self.mesh,
            color_mode=self.color_mode,
            threshold=self.threshold,
            value_min=self.value_min,
            value_max=self.value_max,
            flow=flow,
            scalar_map=scalar_map,
        )
        colors = _shade_vertex_colors(self.mesh, colors, self._shade_factors)

        self.tri_mesh.vertex_colors = o3d.utility.Vector3dVector(colors.astype(np.float64))
        self.vis.update_geometry(self.tri_mesh)

        arrow_result = self._contact_arrow.compute(centroid, centroid_found, fnormal)
        if arrow_result is not None and self._contact_arrow.arrow_mesh is not None:
            arrow_vertices, _, arrow_colors = arrow_result
            self._contact_arrow.apply(arrow_vertices, arrow_colors)
            self.vis.update_geometry(self._contact_arrow.arrow_mesh)

        self.vis.poll_events()
        self.vis.update_renderer()
        return _numpy_o3d_image_to_rgb(self.vis.capture_screen_float_buffer(do_render=False))

    def close(self) -> None:
        if self.vis is not None:
            self.vis.destroy_window()
            self.vis = None


def _create_backend(mesh: UVSurfaceMesh, **kwargs):
    # Filament OffscreenRenderer needs EGL headless (not available on Windows).
    if sys.platform == "win32":
        return _HiddenVisualizerBackend(mesh, **kwargs)
    try:
        return _FilamentOffscreenBackend(mesh, **kwargs)
    except RuntimeError:
        return _HiddenVisualizerBackend(mesh, **kwargs)

class OffscreenFingerMeshRenderer:
    """Render finger UV mesh + contact arrow to an RGB image for Qt embedding."""

    def __init__(
        self,
        mesh: UVSurfaceMesh,
        *,
        width: int = 512,
        height: int = 512,
        threshold: float = 0.0,
        color_mode: str = "depth",
        value_min=None,
        value_max=None,
        world_to_cam: np.ndarray | None = None,
        lookat: np.ndarray | None = None,
        show_contact_arrow: bool = True,
        mesh_lighting=None,
        contact_arrow_style=None,
    ):
        lookat = lookat if lookat is not None else mesh.rest_vertices.mean(axis=0)
        self._backend = _create_backend(
            mesh,
            width=int(width),
            height=int(height),
            threshold=threshold,
            color_mode=color_mode,
            value_min=value_min,
            value_max=value_max,
            world_to_cam=world_to_cam,
            lookat=lookat,
            show_contact_arrow=show_contact_arrow,
            mesh_lighting=mesh_lighting,
            contact_arrow_style=contact_arrow_style,
        )

    @property
    def width(self) -> int:
        return self._backend.width

    @property
    def height(self) -> int:
        return self._backend.height

    def resize(self, width: int, height: int) -> None:
        self._backend.resize(width, height)

    def render(
        self,
        *,
        flow=None,
        scalar_map=None,
        centroid=None,
        centroid_found: bool = False,
        fnormal: float = 0.0,
    ) -> np.ndarray:
        return self._backend.render(
            flow=flow,
            scalar_map=scalar_map,
            centroid=centroid,
            centroid_found=centroid_found,
            fnormal=fnormal,
        )

    def close(self) -> None:
        self._backend.close()