"""Mitsuba 3 photorealistic render backend.

Requires: ``pip install mitsuba``
"""

from __future__ import annotations

import logging
import os
from typing import Callable, Optional

import numpy as np

from .render_backend import RenderBackend, RenderCamera, RenderMaterial, RenderSettings

logger = logging.getLogger(__name__)


class MitsubaBackend(RenderBackend):
    """Render backend using Mitsuba 3 path tracer."""

    def name(self) -> str:
        return "Mitsuba 3"

    def is_available(self) -> bool:
        import sys
        import io

        old_stderr = sys.stderr
        sys.stderr = io.StringIO()
        try:
            import mitsuba  # noqa: F401

            return True
        except ImportError:
            return False
        finally:
            sys.stderr = old_stderr

    # ── Scene construction ──────────────────────────────────────────

    def _base_scene_dict(
        self,
        camera: RenderCamera,
        settings: RenderSettings,
        first_mesh_path: Optional[str] = None,
    ) -> dict:
        """Return a scene dict with everything *except* the shape entries.

        When *first_mesh_path* is given the ground plane / backdrop is sized
        from its bounding box; otherwise a large default is used.
        """
        import mitsuba as mi

        lighting = settings.lighting
        ground = settings.ground_plane

        cam_to_world = mi.ScalarTransform4f.look_at(
            origin=list(camera.origin),
            target=list(camera.target),
            up=list(camera.up),
        )

        scene: dict = {
            "type": "scene",
            "integrator": {"type": "path", "max_depth": settings.max_depth},
            "sensor": {
                "type": "perspective",
                "fov": camera.fov,
                "to_world": cam_to_world,
                "film": {
                    "type": "hdrfilm",
                    "width": settings.width,
                    "height": settings.height,
                    "rfilter": {"type": "gaussian"},
                },
                "sampler": {
                    "type": "independent",
                    "sample_count": settings.spp,
                },
            },
            "emitter": {
                "type": "constant",
                "radiance": {
                    "type": "rgb",
                    "value": [
                        lighting.ambient_intensity,
                        lighting.ambient_intensity * 0.97,
                        lighting.ambient_intensity * 0.94,
                    ],
                },
            },
        }

        # ── 3-point lighting ──────────────────────────────────────────
        key_rgb = [c * lighting.key_intensity for c in lighting.key_color]
        scene["key_light"] = {
            "type": "directional",
            "to_world": mi.ScalarTransform4f.look_at(
                origin=[1.0, -0.8, 1.2],
                target=[0.0, 0.0, 0.0],
                up=[0.0, 0.0, 1.0],
            ),
            "irradiance": {"type": "rgb", "value": key_rgb},
        }

        fill_rgb = [c * lighting.fill_intensity for c in lighting.fill_color]
        scene["fill_light"] = {
            "type": "directional",
            "to_world": mi.ScalarTransform4f.look_at(
                origin=[-1.0, 0.6, 0.8],
                target=[0.0, 0.0, 0.0],
                up=[0.0, 0.0, 1.0],
            ),
            "irradiance": {"type": "rgb", "value": fill_rgb},
        }

        rim_rgb = [c * lighting.rim_intensity for c in lighting.rim_color]
        scene["rim_light"] = {
            "type": "directional",
            "to_world": mi.ScalarTransform4f.look_at(
                origin=[-0.3, 1.2, -0.8],
                target=[0.0, 0.0, 0.0],
                up=[0.0, 0.0, 1.0],
            ),
            "irradiance": {"type": "rgb", "value": rim_rgb},
        }

        # ── Ground plane / backdrop ───────────────────────────────────
        if ground.enabled:
            try:
                ext = os.path.splitext(first_mesh_path)[1].lower() if first_mesh_path else ""
                shape_type = "ply" if ext == ".ply" else "obj"
                mesh_shape = mi.load_dict({"type": shape_type, "filename": first_mesh_path})
                bbox = mesh_shape.bbox()
                bbox_min, bbox_max = bbox[0], bbox[1]
                model_height = bbox_max[2] - bbox_min[2]
                ground_z = bbox_min[2] - 0.001 * model_height
                dx = bbox_max[0] - bbox_min[0]
                dy = bbox_max[1] - bbox_min[1]
                dz = bbox_max[2] - bbox_min[2]
                diag = float((dx * dx + dy * dy + dz * dz) ** 0.5)
            except Exception:
                ground_z = -ground.distance_below
                diag = 1000.0

            bsdf_ground = {
                "type": "diffuse",
                "reflectance": {"type": "rgb", "value": list(ground.color)},
            }

            if ground.curved_backdrop:
                half_size = diag * 50.0
                radius = diag * 3.0
                cyl_height = diag * 20.0
                scene["ground_floor"] = {
                    "type": "rectangle",
                    "to_world": mi.ScalarTransform4f.translate([0.0, 0.0, ground_z])
                    @ mi.ScalarTransform4f.scale([half_size, half_size, 1.0]),
                    "bsdf": bsdf_ground,
                }
                scene["ground_backdrop"] = {
                    "type": "cylinder",
                    "radius": radius,
                    "p0": [-cyl_height / 2, -radius, ground_z],
                    "p1": [cyl_height / 2, -radius, ground_z],
                    "to_world": mi.ScalarTransform4f.rotate([1, 0, 0], 90)
                    @ mi.ScalarTransform4f.translate([0.0, 0.0, -radius]),
                    "bsdf": bsdf_ground,
                }
            else:
                half_size = diag * 50.0
                scene["ground_plane"] = {
                    "type": "rectangle",
                    "to_world": mi.ScalarTransform4f.translate([0.0, 0.0, ground_z])
                    @ mi.ScalarTransform4f.scale([half_size, half_size, 1.0]),
                    "bsdf": bsdf_ground,
                }

        return scene

    def _build_scene_dict(
        self,
        mesh_path: str,
        material: RenderMaterial,
        camera: RenderCamera,
        settings: RenderSettings,
    ) -> dict:
        """Build a single-shape Mitsuba scene dictionary."""
        ext = os.path.splitext(mesh_path)[1].lower()
        shape_type = "ply" if ext == ".ply" else "obj"

        scene = self._base_scene_dict(camera, settings, first_mesh_path=mesh_path)
        scene["shape"] = {
            "type": shape_type,
            "filename": mesh_path,
            "bsdf": self._make_bsdf(material),
        }
        return scene

    def _build_assembly_scene_dict(
        self,
        parts: list,
        camera: RenderCamera,
        settings: RenderSettings,
    ) -> dict:
        """Build a multi-shape Mitsuba scene dictionary.

        *parts* is a list of ``(mesh_path, RenderMaterial)`` tuples.
        """
        first_path = parts[0][0] if parts else None
        scene = self._base_scene_dict(camera, settings, first_mesh_path=first_path)

        for i, (mesh_path, material) in enumerate(parts):
            ext = os.path.splitext(mesh_path)[1].lower()
            shape_type = "ply" if ext == ".ply" else "obj"
            scene[f"shape_{i}"] = {
                "type": shape_type,
                "filename": mesh_path,
                "bsdf": self._make_bsdf(material),
            }
        return scene

    def _make_bsdf(self, material: RenderMaterial) -> dict:
        """Convert a RenderMaterial to a Mitsuba BSDF dict."""
        mt = material.bsdf_type

        if mt == "roughconductor":
            # Use material="none" with specular_reflectance set to the
            # metal color.  The pip-installed Mitsuba doesn't include
            # spectral metal data files (iron.spd, copper.spd, etc.).
            return {
                "type": "roughconductor",
                "material": "none",
                "alpha": max(material.roughness, 0.01),
                "specular_reflectance": {
                    "type": "rgb",
                    "value": list(material.color),
                },
            }

        if mt == "roughdielectric":
            return {
                "type": "roughdielectric",
                "int_ior": material.int_ior,
                "ext_ior": 1.0,
                "alpha": max(material.roughness, 0.01),
            }

        if mt == "dielectric":
            return {
                "type": "dielectric",
                "int_ior": material.int_ior,
                "ext_ior": 1.0,
            }

        if mt == "plastic":
            return {
                "type": "plastic",
                "diffuse_reflectance": {
                    "type": "rgb",
                    "value": list(material.color),
                },
                "int_ior": material.int_ior,
            }

        # Default: diffuse
        return {
            "type": "diffuse",
            "reflectance": {
                "type": "rgb",
                "value": list(material.color),
            },
        }

    # ── Rendering ───────────────────────────────────────────────────

    def render(
        self,
        mesh_path: str,
        material: RenderMaterial,
        camera: RenderCamera,
        settings: RenderSettings,
        progress_callback: Optional[Callable[[float], None]] = None,
    ) -> np.ndarray:
        """Render a mesh file and return (H, W, 3) float32 RGB array."""
        self._set_variant()
        import mitsuba as mi

        scene_dict = self._build_scene_dict(mesh_path, material, camera, settings)
        scene = mi.load_dict(scene_dict)

        logger.info(f"Rendering {settings.width}x{settings.height} @ {settings.spp} spp")

        try:
            image = mi.render(scene, spp=settings.spp, seed=int(settings.seed or 0))
        except Exception as e:
            logger.error(f"Mitsuba render failed: {e}")
            raise

        if progress_callback:
            progress_callback(1.0)

        arr = np.array(image, dtype=np.float32)
        arr = np.clip(arr, 0.0, None)
        arr = np.power(arr, 1.0 / 2.2)
        arr = np.clip(arr, 0.0, 1.0)
        return arr

    def render_assembly(
        self,
        parts: list,
        camera: RenderCamera,
        settings: RenderSettings,
        progress_callback: Optional[Callable[[float], None]] = None,
    ) -> np.ndarray:
        """Render multiple meshes with individual materials.

        *parts* is a list of ``(mesh_path, RenderMaterial)`` tuples.
        Returns (H, W, 3) float32 RGB array.
        """
        self._set_variant()
        import mitsuba as mi

        scene_dict = self._build_assembly_scene_dict(parts, camera, settings)
        scene = mi.load_dict(scene_dict)

        logger.info(
            f"Rendering assembly ({len(parts)} parts) "
            f"{settings.width}x{settings.height} @ {settings.spp} spp"
        )

        try:
            image = mi.render(scene, spp=settings.spp, seed=int(settings.seed or 0))
        except Exception as e:
            logger.error(f"Mitsuba assembly render failed: {e}")
            raise

        if progress_callback:
            progress_callback(1.0)

        arr = np.array(image, dtype=np.float32)
        arr = np.clip(arr, 0.0, None)
        arr = np.power(arr, 1.0 / 2.2)
        arr = np.clip(arr, 0.0, 1.0)
        return arr

    def render_preview(
        self,
        mesh_path: str,
        material: RenderMaterial,
        camera: RenderCamera,
        settings: RenderSettings,
    ) -> np.ndarray:
        """Quick low-quality preview (4x fewer spp)."""
        preview_settings = RenderSettings(
            width=settings.width // 2,
            height=settings.height // 2,
            spp=max(settings.spp // 4, 16),
            max_depth=min(settings.max_depth, 4),
            seed=settings.seed,
            lighting=settings.lighting,
            ground_plane=settings.ground_plane,
        )
        return self.render(mesh_path, material, camera, preview_settings)

    # ── Export ──────────────────────────────────────────────────────

    def export_image(self, image: np.ndarray, path: str) -> None:
        """Save a rendered image to PNG or EXR."""
        from PIL import Image

        ext = os.path.splitext(path)[1].lower()

        if ext == ".exr":
            # Save as EXR (HDR) — no tonemapping
            try:
                import OpenEXR  # type: ignore[import-not-found]
                import Imath  # type: ignore[import-not-found]

                h, w = image.shape[:2]
                header = OpenEXR.Header(w, h)
                header["channels"] = {
                    "R": Imath.PixelType(Imath.PixelType.FLOAT),
                    "G": Imath.PixelType(Imath.PixelType.FLOAT),
                    "B": Imath.PixelType(Imath.PixelType.FLOAT),
                }
                exr = OpenEXR.OutputFile(path, header)
                exr.write(
                    {
                        "R": image[:, :, 0].tobytes(),
                        "G": image[:, :, 1].tobytes(),
                        "B": image[:, :, 2].tobytes(),
                    }
                )
                exr.close()
            except ImportError:
                # Fallback: save as 16-bit PNG
                logger.warning("OpenEXR not available, saving as 16-bit PNG")
                img = Image.fromarray((image * 65535).astype(np.uint16), "RGB")
                img.save(path)
        else:
            # PNG / JPEG — already tonemapped
            img = Image.fromarray((image * 255).astype(np.uint8), "RGB")
            img.save(path)

        logger.info(f"Exported render to {path}")

    # ── Helpers ─────────────────────────────────────────────────────

    def _set_variant(self) -> None:
        """Set the Mitsuba variant (called once)."""
        import sys
        import io

        # Suppress the harmless "LLVM API initialization failed" warning
        # that drjit emits on macOS ARM when scalar variant is used.
        old_stderr = sys.stderr
        sys.stderr = io.StringIO()
        try:
            import mitsuba as mi

            mi.set_variant("scalar_rgb")
        finally:
            sys.stderr = old_stderr