"""
CUDASlicer: High-Performance Slicing & Polygon Offsetting Library
================================================================
Provides ultra-fast GPU distance-field (SDF/JFA) & multi-threaded CPU
polygon offsetting, binary mask transformations, and 3D mesh slicing
with automatic fallback and dynamic backend switching.
"""

from typing import List, Dict, Union, Any, Optional, Sequence, Tuple
import numpy as np

__version__ = "1.0.0"

# Try importing the compiled C++/CUDA pybind11 module
try:
    from . import _cudaslicer_backend as _backend
    _HAS_BACKEND = True
except ImportError:
    _backend = None
    _HAS_BACKEND = False


class LayerResult(dict):
    """
    Dictionary holding slicing and offsetting results for a single layer.
    Supports dictionary access, attribute access, and tuple unpacking:
    `perimeters, infill_mask = layer`
    """
    def __getattr__(self, key: str) -> Any:
        try:
            return self[key]
        except KeyError:
            raise AttributeError(f"'LayerResult' object has no attribute '{key}'")

    def __setattr__(self, key: str, value: Any) -> None:
        self[key] = value

    def __iter__(self):
        # Enables direct unpacking: perimeters, infill_mask = layer
        yield self.get("perimeters", [])
        yield self.get("infill_mask", None)


def has_cuda() -> bool:
    """
    Check if a functional CUDA GPU device is available on this system.

    Returns
    -------
    bool
        True if CUDA-capable GPU is available, False otherwise.
    """
    if not _HAS_BACKEND or _backend is None:
        return False
    try:
        return bool(_backend.has_cuda())
    except Exception:
        return False


def is_cuda_available() -> bool:
    """Alias for has_cuda(). Returns True if CUDA GPU is functional."""
    return has_cuda()


def is_available() -> bool:
    """
    Check if the CUDASlicer backend (CUDA or CPU fallback) is available.

    Returns
    -------
    bool
        True if the native backend is functional.
    """
    if not _HAS_BACKEND or _backend is None:
        return False
    try:
        ret = _backend.init(0)
        return ret == 0
    except Exception:
        return False


def get_cuda_device_count() -> int:
    """Return the number of available CUDA GPUs."""
    return 1 if has_cuda() else 0


def set_backend(backend: str = "auto") -> None:
    """
    Set the global default compute backend.

    Parameters
    ----------
    backend : str, default='auto'
        Execution backend: 'auto' (use CUDA if available, fallback to CPU),
        'cuda' (require CUDA; throws error if unavailable), or 'cpu' (compatibility mode).
    """
    if not _HAS_BACKEND or _backend is None:
        raise RuntimeError("CUDASlicer native backend extension is not loaded.")
    _backend.set_backend(str(backend).lower())


def get_backend() -> str:
    """
    Get the active global compute backend ('auto', 'cuda', or 'cpu').
    """
    if not _HAS_BACKEND or _backend is None:
        return "none"
    return _backend.get_backend()


def offset_masks(
    masks: np.ndarray,
    perimeters_count: int = 2,
    layer_width: float = 0.45,
    pitch_mm: float = 0.10,
    offset_distances: Optional[Sequence[float]] = None,
    backend: Optional[str] = "auto"
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Compute Euclidean offset perimeter masks and infill mask (CUDA or CPU).

    Parameters
    ----------
    masks : np.ndarray
        2D (H, W) or 3D (N, H, W) binary occupancy grid (uint8 or bool).
    perimeters_count : int, default=2
        Number of perimeter passes (Perimeter 1 at -0.5*width, Perimeter 2 at -1.5*width, etc.).
    layer_width : float, default=0.45
        Extrusion line width in millimeters.
    pitch_mm : float, default=0.10
        Grid pixel pitch in millimeters.
    offset_distances : sequence of float, optional
        Custom offset distances. If provided, overrides perimeters_count and layer_width.
    backend : str, optional, default='auto'
        Compute backend: 'auto', 'cuda', or 'cpu'.

    Returns
    -------
    perimeters_masks : np.ndarray
        Shape (perimeters_count, H, W) or (perimeters_count, N, H, W).
    infill_mask : np.ndarray
        Shape (H, W) or (N, H, W) at -perimeters_count * layer_width.
    """
    if not _HAS_BACKEND or _backend is None:
        raise RuntimeError("CUDASlicer backend extension (_cudaslicer_backend) is not available.")

    masks_uint8 = np.ascontiguousarray(masks, dtype=np.uint8)
    offsets_float = [float(d) for d in offset_distances] if offset_distances is not None else []
    backend_str = str(backend) if backend is not None else "auto"

    out_masks = _backend.offset_masks(
        masks_uint8,
        int(perimeters_count),
        float(layer_width),
        float(pitch_mm),
        offsets_float,
        backend_str
    )

    if offset_distances is not None:
        return out_masks

    # Separate perimeters and infill mask
    perimeters_masks = out_masks[:perimeters_count]
    infill_mask = out_masks[perimeters_count]
    return perimeters_masks, infill_mask


def offset_polygons(
    layers: Union[List[np.ndarray], List[Dict[str, Any]]],
    perimeters_count: int = 2,
    layer_width: float = 0.45,
    pitch_mm: float = 0.10,
    pad_mm: float = 1.5,
    return_infill_mask: bool = True,
    offset_distances: Optional[Sequence[float]] = None,
    backend: Optional[str] = "auto"
) -> Union[LayerResult, List[LayerResult]]:
    """
    Offset 2D polygon loops across one or more layers (CUDA or CPU fallback).

    Parameters
    ----------
    layers : list of np.ndarray OR list of dict
        Input layer data. Can be:
        1. Single layer with a list of (K, 2) NumPy polygon contours:
           `[poly1, poly2, ...]`
        2. Multi-layer list of layer dicts:
           `[{"z": 0.2, "contours": [poly1, ...]}, {"z": 0.4, "contours": [...]}]`
    perimeters_count : int, default=2
        Number of perimeter passes (Perimeter 1: -0.5*width, Perimeter 2: -1.5*width, etc.).
    layer_width : float, default=0.45
        Extrusion line width in millimeters.
    pitch_mm : float, default=0.10
        Pixel raster pitch in millimeters.
    pad_mm : float, default=1.5
        Bounding box padding around the contours in millimeters.
    return_infill_mask : bool, default=True
        Whether to export 2D binary infill occupancy grid.
    offset_distances : sequence of float, optional
        Custom offset distances.
    backend : str, optional, default='auto'
        Compute backend: 'auto', 'cuda', or 'cpu'.

    Returns
    -------
    LayerResult or list of LayerResult
        If single layer: returns single LayerResult object.
        If multi-layer: returns list of LayerResult objects.
    """
    if not _HAS_BACKEND or _backend is None:
        raise RuntimeError("CUDASlicer backend extension (_cudaslicer_backend) is not available.")

    if not layers:
        return []

    is_single_layer_format = False
    if isinstance(layers[0], (np.ndarray, list, tuple)) and not isinstance(layers[0], dict):
        is_single_layer_format = True
        formatted_layers = [{"z": 0.0, "contours": [np.ascontiguousarray(c, dtype=np.float32) for c in layers]}]
    else:
        formatted_layers = []
        for l in layers:
            z_val = float(l.get("z", l.get("z_height", 0.0)))
            cnts = [np.ascontiguousarray(c, dtype=np.float32) for c in l.get("contours", [])]
            formatted_layers.append({"z": z_val, "contours": cnts})

    all_pts = []
    for l in formatted_layers:
        for c in l["contours"]:
            if len(c) >= 3:
                all_pts.append(c)

    if not all_pts:
        empty_res = LayerResult({"z": 0.0, "perimeters": [[] for _ in range(perimeters_count)], "infill_contours": [], "infill_mask": None})
        return empty_res if is_single_layer_format else []

    stacked = np.concatenate(all_pts, axis=0)
    min_x, min_y = float(stacked[:, 0].min()), float(stacked[:, 1].min())
    max_x, max_y = float(stacked[:, 0].max()), float(stacked[:, 1].max())

    offsets_float = [float(d) for d in offset_distances] if offset_distances is not None else []
    backend_str = str(backend) if backend is not None else "auto"

    raw_results = _backend.offset_polygon_layers(
        formatted_layers, min_x, min_y, max_x, max_y,
        int(perimeters_count), float(layer_width),
        float(pitch_mm), float(pad_mm),
        bool(return_infill_mask), offsets_float,
        backend_str
    )

    layer_results = [LayerResult(r) for r in raw_results]
    if is_single_layer_format and layer_results:
        return layer_results[0]

    return layer_results


def slice_mesh_and_offset(
    mesh_or_triangles: Any,
    layer_height: float = 0.20,
    first_layer_height: Optional[float] = None,
    perimeters_count: int = 2,
    layer_width: float = 0.45,
    pitch_mm: float = 0.10,
    pad_mm: float = 1.5,
    return_infill_mask: bool = True,
    offset_distances: Optional[Sequence[float]] = None,
    backend: Optional[str] = "auto",
    simplify_tolerance_mm: float = 0.0,
    subpixel_lerp: bool = True
) -> List[LayerResult]:
    """
    3D mesh slicing, rasterization, distance field, and perimeter offsetting (CUDA or CPU).

    Parameters
    ----------
    mesh_or_triangles : trimesh.Trimesh OR np.ndarray
        Either a trimesh.Trimesh object, or a NumPy array of shape (N, 3, 3) or (N, 9) of triangle vertices.
    layer_height : float, default=0.20
        Z slicing layer height in millimeters.
    first_layer_height : float, optional
        First layer height in millimeters. Defaults to `layer_height` if None or <= 0.
    perimeters_count : int, default=2
        Number of perimeter passes (Perimeter 1: -0.5*width, Perimeter 2: -1.5*width, etc.).
    layer_width : float, default=0.45
        Extrusion line width in millimeters.
    pitch_mm : float, default=0.10
        Grid pixel pitch in millimeters.
    pad_mm : float, default=1.5
        Bounding box padding margin in millimeters.
    return_infill_mask : bool, default=True
        Whether to generate the 2D binary infill occupancy grid.
    offset_distances : sequence of float, optional
        Custom offset distances.
    backend : str, optional, default='auto'
        Compute backend: 'auto' (CUDA if available, else CPU), 'cuda' (require CUDA), or 'cpu'.
    simplify_tolerance_mm : float, default=0.0
        Ramer-Douglas-Peucker (RDP) polygon simplification tolerance in mm (0.0 to disable).
    subpixel_lerp : bool, default=True
        Continuous distance-field sub-pixel linear interpolation for smooth arbitrary angles.

    Returns
    -------
    list of LayerResult
        List of LayerResult dictionaries containing:
        - 'z': Layer Z height
        - 'perimeters': List of perimeters (each containing list of (K, 2) contour arrays)
        - 'infill_contours': Infill boundary contour loops
        - 'infill_mask': 2D binary NumPy array (H, W) of the infill region
    """
    if not _HAS_BACKEND or _backend is None:
        raise RuntimeError("CUDASlicer backend extension (_cudaslicer_backend) is not available.")

    # Handle Trimesh object
    if hasattr(mesh_or_triangles, "triangles"):
        tris = np.ascontiguousarray(mesh_or_triangles.triangles, dtype=np.float32)
        bounds = mesh_or_triangles.bounds
        min_x, min_y, min_z = float(bounds[0][0]), float(bounds[0][1]), float(bounds[0][2])
        max_x, max_y, max_z = float(bounds[1][0]), float(bounds[1][1]), float(bounds[1][2])
    else:
        tris = np.ascontiguousarray(mesh_or_triangles, dtype=np.float32)
        if tris.ndim == 2 and tris.shape[1] == 9:
            tris_3d = tris.reshape(-1, 3, 3)
        else:
            tris_3d = tris
        min_x = float(tris_3d[:, :, 0].min())
        max_x = float(tris_3d[:, :, 0].max())
        min_y = float(tris_3d[:, :, 1].min())
        max_y = float(tris_3d[:, :, 1].max())
        min_z = float(tris_3d[:, :, 2].min())
        max_z = float(tris_3d[:, :, 2].max())

    h_first = float(first_layer_height) if first_layer_height is not None and first_layer_height > 0 else float(layer_height)
    offsets_float = [float(d) for d in offset_distances] if offset_distances is not None else []
    backend_str = str(backend) if backend is not None else "auto"

    raw_results = _backend.slice_mesh_and_offset(
        tris, min_x, min_y, min_z, max_x, max_y, max_z,
        float(layer_height), float(h_first),
        int(perimeters_count), float(layer_width),
        float(pitch_mm), float(pad_mm),
        bool(return_infill_mask), offsets_float,
        backend_str,
        float(simplify_tolerance_mm),
        bool(subpixel_lerp)
    )

    return [LayerResult(r) for r in raw_results]


# Alias for convenience
slice_mesh = slice_mesh_and_offset


class SdfStack:
    """
    GPU-resident Signed Distance Field (SDF) stack representing a 3D model or modifier volume.
    Enables chained GPU-accelerated Boolean CSG operations (union, intersection, difference,
    smooth blend, offset, shelling, procedural 3D infill) without host-to-device data transfers.
    """
    def __init__(self, backend_stack: Any):
        self._stack = backend_stack

    @property
    def metadata(self) -> Dict[str, Any]:
        """Dictionary containing spatial bounds, grid dimensions, and layer count."""
        return self._stack.get_metadata()

    @property
    def layer_z_um(self) -> np.ndarray:
        """1D array of layer heights in micrometers (µm)."""
        return self._stack.get_layer_z_um()

    @property
    def layer_z_mm(self) -> np.ndarray:
        """1D array of layer heights in millimeters (mm)."""
        return self._stack.get_layer_z_mm()

    @property
    def shape(self) -> Tuple[int, int, int]:
        """Grid shape (num_layers, height, width)."""
        m = self.metadata
        return (m["num_layers"], m["height"], m["width"])

    def numpy(self) -> np.ndarray:
        """Download raw 3D distance field array (num_layers, H, W) to CPU NumPy float32."""
        return self._stack.to_numpy()

    def clone(self) -> "SdfStack":
        """Duplicate this SDF stack in GPU VRAM."""
        return SdfStack(self._stack.clone())

    def union(self, other: "SdfStack", blend_radius_mm: float = 0.0) -> "SdfStack":
        """Compute Boolean Union (A ∪ B) on GPU. Optional blend_radius_mm for smooth fillet union."""
        if not isinstance(other, SdfStack):
            raise TypeError(f"Operand must be an SdfStack, got {type(other)}")
        return SdfStack(self._stack.op_union(other._stack, float(blend_radius_mm)))

    def intersection(self, other: "SdfStack", blend_radius_mm: float = 0.0) -> "SdfStack":
        """Compute Boolean Intersection (A ∩ B) on GPU. Optional blend_radius_mm for smooth intersection."""
        if not isinstance(other, SdfStack):
            raise TypeError(f"Operand must be an SdfStack, got {type(other)}")
        return SdfStack(self._stack.op_intersection(other._stack, float(blend_radius_mm)))

    def difference(self, other: "SdfStack", blend_radius_mm: float = 0.0) -> "SdfStack":
        """Compute Boolean Difference / Subtraction (A \\ B) on GPU. Optional blend_radius_mm for smooth difference."""
        if not isinstance(other, SdfStack):
            raise TypeError(f"Operand must be an SdfStack, got {type(other)}")
        return SdfStack(self._stack.op_difference(other._stack, float(blend_radius_mm)))

    def offset(self, delta_mm: float) -> "SdfStack":
        """Constant-distance offset/dilation (delta > 0) or inset/erosion (delta < 0) on GPU."""
        return SdfStack(self._stack.offset(float(delta_mm)))

    def shell(self, thickness_mm: float) -> "SdfStack":
        """Hollow / shell the volume with specified wall thickness on GPU."""
        return SdfStack(self._stack.shell(float(thickness_mm)))



    def vectorize(
        self,
        perimeters_count: int = 2,
        layer_width: float = 0.45,
        export_infill_mask: bool = True,
        simplify_tolerance_mm: float = 0.02,
        subpixel_lerp: bool = True,
        offset_distances: Optional[Sequence[float]] = None
    ) -> List[LayerResult]:
        """
        Vectorize GPU SDF stack directly to smooth polygon perimeter contours and infill masks.

        Parameters
        ----------
        perimeters_count : int, default=2
            Number of perimeter passes.
        layer_width : float, default=0.45
            Extrusion line width in mm.
        export_infill_mask : bool, default=True
            Whether to export binary 2D infill mask.
        simplify_tolerance_mm : float, default=0.02
            Ramer-Douglas-Peucker (RDP) polygon simplification tolerance in mm (eliminates micro-steps).
        subpixel_lerp : bool, default=True
            Continuous SDF linear interpolation for sub-pixel boundary placement at arbitrary angles.
        offset_distances : sequence of float, optional
            Custom offset distances.
        """
        offsets_float = [float(d) for d in offset_distances] if offset_distances is not None else []
        raw_res = self._stack.vectorize(
            int(perimeters_count),
            float(layer_width),
            bool(export_infill_mask),
            float(simplify_tolerance_mm),
            bool(subpixel_lerp),
            offsets_float
        )
        return [LayerResult(r) for r in raw_res]

    # Operator Overloads for fluent Python syntax
    def __or__(self, other: "SdfStack") -> "SdfStack":
        return self.union(other)

    def __and__(self, other: "SdfStack") -> "SdfStack":
        return self.intersection(other)

    def __sub__(self, other: "SdfStack") -> "SdfStack":
        return self.difference(other)

    def __add__(self, delta: float) -> "SdfStack":
        if isinstance(delta, (int, float)):
            return self.offset(float(delta))
        raise TypeError(f"Unsupported operand type for +: 'SdfStack' and '{type(delta)}'")

    def __repr__(self) -> str:
        m = self.metadata
        return f"<SdfStack layers={m['num_layers']} shape=({m['height']}, {m['width']}) pitch={m['pitch_mm']:.2f}mm>"


def slice_mesh_to_sdf(
    mesh_or_triangles: Any,
    layer_height: float = 0.20,
    first_layer_height: Optional[float] = None,
    pitch_mm: float = 0.10,
    pad_mm: float = 1.5
) -> SdfStack:
    """
    Slice a 3D mesh directly to a GPU-resident Signed Distance Field (SDF) stack.

    Parameters
    ----------
    mesh_or_triangles : trimesh.Trimesh OR np.ndarray
        Mesh or triangle array of shape (N, 3, 3) or (N, 9).
    layer_height : float, default=0.20
        Z slicing layer height in millimeters.
    first_layer_height : float, optional
        First layer height in millimeters. Defaults to `layer_height` if None.
    pitch_mm : float, default=0.10
        Pixel raster pitch in millimeters.
    pad_mm : float, default=1.5
        Bounding box padding in millimeters.

    Returns
    -------
    SdfStack
        GPU-resident SDF stack object ready for CSG operations or vectorization.
    """
    if not _HAS_BACKEND or _backend is None:
        raise RuntimeError("CUDASlicer backend extension (_cudaslicer_backend) is not available.")

    if hasattr(mesh_or_triangles, "triangles"):
        tris = np.ascontiguousarray(mesh_or_triangles.triangles, dtype=np.float32)
        bounds = mesh_or_triangles.bounds
        min_x, min_y, min_z = float(bounds[0][0]), float(bounds[0][1]), float(bounds[0][2])
        max_x, max_y, max_z = float(bounds[1][0]), float(bounds[1][1]), float(bounds[1][2])
    else:
        tris = np.ascontiguousarray(mesh_or_triangles, dtype=np.float32)
        if tris.ndim == 2 and tris.shape[1] == 9:
            tris_3d = tris.reshape(-1, 3, 3)
        else:
            tris_3d = tris
        min_x = float(tris_3d[:, :, 0].min())
        max_x = float(tris_3d[:, :, 0].max())
        min_y = float(tris_3d[:, :, 1].min())
        max_y = float(tris_3d[:, :, 1].max())
        min_z = float(tris_3d[:, :, 2].min())
        max_z = float(tris_3d[:, :, 2].max())

    h_first = float(first_layer_height) if first_layer_height is not None and first_layer_height > 0 else float(layer_height)

    backend_sdf = _backend.slice_mesh_to_sdf(
        tris, min_x, min_y, min_z, max_x, max_y, max_z,
        float(layer_height), float(h_first), float(pitch_mm), float(pad_mm)
    )
    return SdfStack(backend_sdf)


def masks_to_sdf(
    masks: np.ndarray,
    layer_z_um: Optional[Sequence[int]] = None,
    min_x_mm: float = 0.0,
    min_y_mm: float = 0.0,
    pitch_mm: float = 0.10
) -> SdfStack:
    """
    Compute a GPU-resident Signed Distance Field from 2D (H, W) or 3D (N, H, W) binary masks.
    """
    if not _HAS_BACKEND or _backend is None:
        raise RuntimeError("CUDASlicer backend extension (_cudaslicer_backend) is not available.")

    masks_u8 = np.ascontiguousarray(masks, dtype=np.uint8)
    z_arr = np.ascontiguousarray(layer_z_um, dtype=np.int32) if layer_z_um is not None else np.array([], dtype=np.int32)

    backend_sdf = _backend.masks_to_sdf(masks_u8, z_arr, float(min_x_mm), float(min_y_mm), float(pitch_mm))
    return SdfStack(backend_sdf)


__all__ = [
    "offset_masks",
    "offset_polygons",
    "slice_mesh_and_offset",
    "slice_mesh",
    "slice_mesh_to_sdf",
    "masks_to_sdf",
    "SdfStack",
    "LayerResult",
    "is_available",
    "has_cuda",
    "is_cuda_available",
    "set_backend",
    "get_backend",
    "get_cuda_device_count",
    "__version__",
]