# SPDX-License-Identifier: Apache-2.0
# @author Arnaud Morin <arnaud.gfpv@mailops.fr>

"""Headless 3D render helper -- shared by every part module and the assembly.

    from render import render
    render(shape_or_assembly, "build/foo.png", view="iso")

Views: iso, top, front, right, or an (elev, azim) tuple.
"""

import matplotlib
matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

import cadquery as cq

VIEWS = {
    "iso": (26, -55),
    "iso2": (20, 35),
    "top": (89, -90),
    "front": (0, -90),
    "right": (0, 0),
}

DEFAULT_COLOR = (0.20, 0.20, 0.22)


def _iter_parts(assy, parent_loc=None):
    """Yield (located_shape, rgba|None) for every solid in an Assembly."""
    loc = assy.loc if parent_loc is None else parent_loc * assy.loc
    if assy.obj is not None:
        shape = assy.obj.val() if isinstance(assy.obj, cq.Workplane) else assy.obj
        col = assy.color.toTuple() if assy.color is not None else None
        yield shape.moved(loc), col
    for child in assy.children:
        yield from _iter_parts(child, loc)


def _shapes(obj):
    if isinstance(obj, cq.Assembly):
        return list(_iter_parts(obj))
    if isinstance(obj, cq.Workplane):
        return [(obj.val(), None)]
    return [(obj, None)]


def render(obj, path, view="iso", tol=0.08, figsize=(9, 9), dpi=110,
           edges=True, title=None):
    parts = _shapes(obj)

    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(111, projection="3d")

    lo = np.array([+1e9, +1e9, +1e9])
    hi = np.array([-1e9, -1e9, -1e9])

    for shape, col in parts:
        verts, tris = shape.tessellate(tol)
        if not tris:
            continue
        v = np.array([[p.x, p.y, p.z] for p in verts])
        lo = np.minimum(lo, v.min(axis=0))
        hi = np.maximum(hi, v.max(axis=0))
        face = col[:3] if col else DEFAULT_COLOR
        pc = Poly3DCollection(
            v[np.array(tris)],
            facecolor=face,
            edgecolor=(0, 0, 0, 0.25) if edges else "none",
            linewidth=0.15 if edges else 0,
        )
        pc.set_alpha(col[3] if col else 1.0)
        ax.add_collection3d(pc)

    ctr = (lo + hi) / 2.0
    span = float((hi - lo).max()) * 0.55 or 1.0
    ax.set_xlim(ctr[0] - span, ctr[0] + span)
    ax.set_ylim(ctr[1] - span, ctr[1] + span)
    ax.set_zlim(ctr[2] - span, ctr[2] + span)
    try:
        ax.set_box_aspect((1, 1, 1))
    except Exception:
        pass

    elev, azim = VIEWS.get(view, view) if isinstance(view, str) else view
    ax.view_init(elev=elev, azim=azim)
    ax.set_axis_off()
    if title:
        ax.set_title(title)
    fig.savefig(path, dpi=dpi, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print(f"wrote {path}")
    return path


def render_grid(obj, path, views=("iso", "top", "front", "right"), tol=0.08, dpi=100):
    """2x2 contact sheet -- the quickest way to eyeball a part or assembly."""
    parts = _shapes(obj)
    fig = plt.figure(figsize=(12, 12))
    for i, view in enumerate(views, start=1):
        ax = fig.add_subplot(2, 2, i, projection="3d")
        lo = np.array([+1e9] * 3)
        hi = np.array([-1e9] * 3)
        for shape, col in parts:
            verts, tris = shape.tessellate(tol)
            if not tris:
                continue
            v = np.array([[p.x, p.y, p.z] for p in verts])
            lo = np.minimum(lo, v.min(axis=0))
            hi = np.maximum(hi, v.max(axis=0))
            pc = Poly3DCollection(
                v[np.array(tris)],
                facecolor=col[:3] if col else DEFAULT_COLOR,
                edgecolor=(0, 0, 0, 0.25),
                linewidth=0.15,
            )
            pc.set_alpha(col[3] if col else 1.0)
            ax.add_collection3d(pc)
        ctr = (lo + hi) / 2.0
        span = float((hi - lo).max()) * 0.55 or 1.0
        ax.set_xlim(ctr[0] - span, ctr[0] + span)
        ax.set_ylim(ctr[1] - span, ctr[1] + span)
        ax.set_zlim(ctr[2] - span, ctr[2] + span)
        try:
            ax.set_box_aspect((1, 1, 1))
        except Exception:
            pass
        elev, azim = VIEWS.get(view, view) if isinstance(view, str) else view
        ax.view_init(elev=elev, azim=azim)
        ax.set_axis_off()
        ax.set_title(str(view))
    fig.savefig(path, dpi=dpi, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print(f"wrote {path}")
    return path