"""Gmsh polyprism definitions."""
from __future__ import annotations

from typing import TYPE_CHECKING

import gmsh
from shapely.geometry import MultiPolygon, Polygon

from meshwell.cad import CAD
from meshwell.geometry_entity import GeometryEntity
from meshwell.validation import format_physical_name

if TYPE_CHECKING:
    from OCP.TopoDS import TopoDS_Shape


class PolyPrism(GeometryEntity):
    """Creates a bottom-up GMSH "prism" formed by a polygon associated with (optional) z-dependent grow/shrink operations.

    Attributes:
        polygons: list of shapely (Multi)Polygon
        buffers: dict of {z: buffer} used to shrink/grow base polygons at specified z-values
        physical_name: name of the physical this entity will belong to
        mesh_order: priority of the entity if it overlaps with others (lower numbers override higher numbers)
        mesh_bool: if True, entity will be meshed; if not, will not be meshed (useful to tag boundaries)

    """

    def __init__(
        self,
        polygons: Polygon | list[Polygon] | MultiPolygon | list[MultiPolygon],
        buffers: dict[float, float],
        physical_name: str | tuple[str, ...] | None = None,
        mesh_order: float | None = None,
        mesh_bool: bool = True,
        additive: bool = False,
        subdivision: tuple[int, int, int] | None = None,
        point_tolerance: float = 1e-3,
        identify_arcs: bool = False,
        min_arc_points: int = 4,
        arc_tolerance: float = 1e-3,
        translation: tuple[float, float, float] | None = None,
        rotation_axis: tuple[float, float, float] | None = None,
        rotation_point: tuple[float, float, float] | None = None,
        rotation_angle: float = 0.0,
    ):
        # Initialize parent class with point tracking and transformation parameters
        super().__init__(
            point_tolerance=point_tolerance,
            translation=translation,
            rotation_axis=rotation_axis,
            rotation_point=rotation_point,
            rotation_angle=rotation_angle,
        )

        # Parse buffers or prepare extrusion
        self.polygons = polygons
        if all(buffer == 0 for buffer in buffers.values()):
            self.extrude = True
            self.zmin, self.zmax = min(buffers.keys()), max(buffers.keys())
        else:
            self.extrude = False
            self.buffered_polygons: list[
                tuple[float, Polygon]
            ] = self._get_buffered_polygons(polygons, buffers)

        # Store other attributes
        self.buffers = buffers
        self.mesh_order = mesh_order
        self.additive = additive
        self.dimension = 3
        self.subdivision = subdivision
        self.identify_arcs = identify_arcs
        self.min_arc_points = min_arc_points
        self.arc_tolerance = arc_tolerance

        if self.identify_arcs and not self.extrude:
            raise NotImplementedError(
                "Arc identification is currently only supported for PolyPrism when extrude=True."
            )

        # Format physical name
        self.physical_name = format_physical_name(physical_name)
        self.mesh_bool = mesh_bool
        self.additive = additive

    def _create_volumes_directly(self) -> list[int]:
        """Create GMSH volumes directly without using CAD class methods."""
        if self.extrude:
            surfaces = self._create_surfaces_with_holes_at_z(self.polygons, self.zmin)
            surface_dimtags = [(2, surface) for surface in surfaces]
            entities = gmsh.model.occ.extrude(
                surface_dimtags, 0, 0, self.zmax - self.zmin
            )
            volumes = [tag for dim, tag in entities if dim == 3]
        else:
            volumes = [
                tag
                for entry in self.buffered_polygons
                if (tag := self._create_volume_with_holes_directly(entry)) != 0
            ]

        # Fuse multiple volumes if needed
        if len(volumes) <= 1:
            return volumes

        fused_dimtags = gmsh.model.occ.fuse(
            [(3, volumes[0])],
            [(3, tag) for tag in volumes[1:]],
            removeObject=True,
            removeTool=True,
        )[0]
        # Clear caches after boolean operations that may invalidate geometry IDs
        self._clear_caches()
        return [tag for dim, tag in fused_dimtags]

    def _get_buffered_polygons(
        self, polygons: list[Polygon], buffers: dict[float, float]
    ) -> list[tuple[float, Polygon]]:
        """Break up polygons on each layer into lists of (z,polygon) tuples according to buffer entries.

        Arguments (implicit):
            polygons: list of (Multi)Polygons to bufferize
            buffers: {z: buffer} values to apply to the polygons

        Returns:
            buffered_polygons: list of (z, buffered_polygons)

        """
        all_polygons_list = []
        for polygon in polygons.geoms if hasattr(polygons, "geoms") else [polygons]:
            current_buffers = []
            for z, width_buffer in buffers.items():
                current_buffers.append((z, polygon.buffer(width_buffer, join_style=2)))
            all_polygons_list.append(current_buffers)

        return all_polygons_list

    def _create_volume_directly(
        self,
        entry: list[tuple[float, Polygon]],
        exterior: bool = True,
        interior_index: int = 0,
    ) -> int:
        """Create volume directly using GMSH calls."""
        curve_loops = []
        for z, polygon in entry:
            vertices = self.xy_surface_vertices(
                polygon=polygon,
                polygon_z=z,
                exterior=exterior,
                interior_index=interior_index,
            )
            # Create a curve loop for this polygon at this z
            points = self._create_points_from_vertices(vertices)
            lines = []
            for i in range(len(points) - 1):
                line_id = self._add_line_with_cache(points[i], points[i + 1])
                if line_id != 0:
                    lines.append(line_id)
            if lines:
                loop_id = gmsh.model.occ.addCurveLoop(lines)
                curve_loops.append(loop_id)

        if not curve_loops:
            return 0

        # If only one curve loop, we can't build a volume
        if len(curve_loops) < 2:
            return 0

        # Return volume from thru-sections
        try:
            volume_dimtags = gmsh.model.occ.addThruSections(
                curve_loops, makeSolid=True, makeRuled=True
            )
            gmsh.model.occ.synchronize()
            if volume_dimtags and volume_dimtags[0][0] == 3:
                return volume_dimtags[0][1]
        except Exception:
            return 0

        return 0

    def xy_surface_vertices(
        self,
        polygon: Polygon,
        polygon_z: float,
        exterior: bool,
        interior_index: int,
    ) -> list[tuple[float, float, float]]:
        """Draw xy surface."""
        return (
            [(x, y, polygon_z) for x, y in polygon.exterior.coords]
            if exterior
            else [
                (x, y, polygon_z) for x, y in polygon.interiors[interior_index].coords
            ]
        )

    def _create_volume_with_holes_directly(
        self, entry: list[tuple[float, Polygon]]
    ) -> int:
        """Create volume with holes directly using GMSH calls."""
        exterior = self._create_volume_directly(entry, exterior=True)
        if exterior == 0:
            return 0
        interiors = [
            tag
            for interior_index in range(len(entry[0][1].interiors))
            if (
                tag := self._create_volume_directly(
                    entry,
                    exterior=False,
                    interior_index=interior_index,
                )
            )
            != 0
        ]
        if interiors:
            cut_result = gmsh.model.occ.cut(
                [(3, exterior)],
                [(3, interior) for interior in interiors],
                removeObject=True,
                removeTool=True,
            )
            gmsh.model.occ.synchronize()
            # Handle cut result - if cut succeeded, use the result, otherwise check if original still exists
            if cut_result and cut_result[0]:
                exterior = cut_result[0][0][1]  # Parse `outDimTags', `outDimTagsMap'
            else:
                # If cut failed, original might have been removed due to removeObject=True
                existing_3d = {tag for dim, tag in gmsh.model.getEntities(3)}
                if exterior not in existing_3d:
                    exterior = 0
            # Clear caches after boolean operations that may invalidate geometry IDs
            self._clear_caches()
        return exterior

    def _create_surfaces_with_holes_at_z(self, polygons, z) -> list[int]:
        """Create surfaces with holes at given z level directly using GMSH calls."""
        surfaces = []
        for polygon in polygons.geoms if hasattr(polygons, "geoms") else [polygons]:
            # Create outer surface
            exterior_vertices = [(x, y, z) for x, y in polygon.exterior.coords]
            exterior = self._create_surface_from_vertices(
                exterior_vertices,
                identify_arcs=self.identify_arcs,
                min_arc_points=self.min_arc_points,
                arc_tolerance=self.arc_tolerance,
            )
            if exterior == 0:
                continue

            # Create interior surfaces (holes)
            interior_surfaces = []
            for interior in polygon.interiors:
                interior_vertices = [(x, y, z) for x, y in interior.coords]
                interior_surface = self._create_surface_from_vertices(
                    interior_vertices,
                    identify_arcs=self.identify_arcs,
                    min_arc_points=self.min_arc_points,
                    arc_tolerance=self.arc_tolerance,
                )
                if interior_surface != 0:
                    interior_surfaces.append(interior_surface)

            # Cut holes from exterior surface
            for interior_surface in interior_surfaces:
                cut_result = gmsh.model.occ.cut(
                    [(2, exterior)],
                    [(2, interior_surface)],
                    removeObject=True,
                    removeTool=True,
                )
                gmsh.model.occ.synchronize()
                # Handle cut result - if cut succeeded, use the result, otherwise keep original
                if cut_result and cut_result[0]:
                    exterior = cut_result[0][0][1]
                # Clear caches after boolean operations that may invalidate geometry IDs
                self._clear_caches()

            surfaces.append(exterior)
        return surfaces

    def subdivide(self, model, prisms, subdivision):
        """Split the prisms into subprisms according to subdivision."""
        subdivided_prisms = []
        import numpy as np

        global_xmin = np.inf
        global_ymin = np.inf
        global_zmin = np.inf
        global_xmax = -np.inf
        global_ymax = -np.inf
        global_zmax = -np.inf
        for prism in prisms:
            xmin, ymin, zmin, xmax, ymax, zmax = model.occ.getBoundingBox(3, prism)
            if xmin < global_xmin:
                global_xmin = xmin
            if ymin < global_ymin:
                global_ymin = ymin
            if zmin < global_zmin:
                global_zmin = zmin
            if xmax > global_xmax:
                global_xmax = xmax
            if ymax > global_ymax:
                global_ymax = ymax
            if zmax > global_zmax:
                global_zmax = zmax
        dx = (global_xmax - global_xmin) / subdivision[0]
        dy = (global_ymax - global_ymin) / subdivision[1]
        dz = (global_zmax - global_zmin) / subdivision[2]
        prisms_dimtags = {(3, prism) for prism in prisms}
        for x_index in range(subdivision[0]):
            for y_index in range(subdivision[1]):
                for z_index in range(subdivision[2]):
                    tool = model.occ.add_box(
                        global_xmin + x_index * dx,
                        global_ymin + y_index * dy,
                        global_zmin + z_index * dz,
                        dx,
                        dy,
                        dz,
                    )
                    intersection, intersection_map = model.occ.intersect(
                        list(prisms_dimtags),
                        [(3, tool)],
                        removeObject=False,
                        removeTool=True,
                    )
                    prisms_dimtags -= set(intersection_map[0])
                    subdivided_prisms.extend(intersection_map[0] + intersection)
        model.occ.remove(list(prisms_dimtags))
        return subdivided_prisms

    def instanciate(self, cad_model: CAD) -> list[tuple[int, int]]:
        """Create GMSH volumes directly without using CAD class methods."""
        prisms = self._create_volumes_directly()
        if self.subdivision is not None:
            prisms = self.subdivide(
                cad_model.model_manager.model, prisms, self.subdivision
            )

        dimtags = [(3, prism) for prism in prisms]
        rotation_point = self._get_rotation_point(self.polygons)
        dimtags = self._apply_transformation_gmsh(dimtags, rotation_point)
        gmsh.model.occ.synchronize()
        return dimtags

    def _create_occ_volume(
        self,
        entry: list[tuple[float, Polygon]],
        exterior: bool = True,
        interior_index: int = 0,
    ) -> TopoDS_Shape:
        """Create OCC volume directly using OCP."""
        from OCP.BRep import BRep_Builder
        from OCP.BRepBuilderAPI import (
            BRepBuilderAPI_MakeSolid,
        )
        from OCP.TopoDS import TopoDS_Shell

        # Draw bottom surface
        bottom_polygon = entry[0][1]
        bottom_z = entry[0][0]
        bottom_polygon_vertices = self.xy_surface_vertices(
            polygon=bottom_polygon,
            polygon_z=bottom_z,
            exterior=exterior,
            interior_index=interior_index,
        )
        faces = [
            self._make_occ_face_from_vertices(
                bottom_polygon_vertices,
                identify_arcs=self.identify_arcs,
                min_arc_points=self.min_arc_points,
                arc_tolerance=self.arc_tolerance,
            )
        ]

        # Draw top surface
        top_polygon = entry[-1][1]
        top_z = entry[-1][0]
        top_polygon_vertices = self.xy_surface_vertices(
            polygon=top_polygon,
            polygon_z=top_z,
            exterior=exterior,
            interior_index=interior_index,
        )
        faces.append(
            self._make_occ_face_from_vertices(
                top_polygon_vertices,
                identify_arcs=self.identify_arcs,
                min_arc_points=self.min_arc_points,
                arc_tolerance=self.arc_tolerance,
            )
        )

        # Draw vertical surfaces
        # Note: Vertical surfaces remain straight lines between layers for now
        # unless we want to also identify arcs in the vertical direction.
        # But usually PolyPrism arcs are in the XY plane.
        for pair_index in range(len(entry) - 1):
            if exterior:
                bottom_coords = entry[pair_index][1].exterior.coords
                top_coords = entry[pair_index + 1][1].exterior.coords
            else:
                bottom_coords = entry[pair_index][1].interiors[interior_index].coords
                top_coords = entry[pair_index + 1][1].interiors[interior_index].coords

            bottom_z = entry[pair_index][0]
            top_z = entry[pair_index + 1][0]

            for facet_pt_ind in range(len(bottom_coords) - 1):
                p1 = (
                    bottom_coords[facet_pt_ind][0],
                    bottom_coords[facet_pt_ind][1],
                    bottom_z,
                )
                p2 = (
                    bottom_coords[facet_pt_ind + 1][0],
                    bottom_coords[facet_pt_ind + 1][1],
                    bottom_z,
                )
                p3 = (
                    top_coords[facet_pt_ind + 1][0],
                    top_coords[facet_pt_ind + 1][1],
                    top_z,
                )
                p4 = (top_coords[facet_pt_ind][0], top_coords[facet_pt_ind][1], top_z)

                facet_vertices = [p1, p2, p3, p4, p1]
                # For vertical facets, we use simple linear interpolation
                faces.append(self._make_occ_face_from_vertices(facet_vertices))

        # Build shell and solid
        shell_builder = BRep_Builder()
        shell = TopoDS_Shell()
        shell_builder.MakeShell(shell)
        for face in faces:
            shell_builder.Add(shell, face)

        solid_builder = BRepBuilderAPI_MakeSolid()
        solid_builder.Add(shell)
        return solid_builder.Solid()

    def plot_decomposition(
        self,
        ax=None,
        line_color: str = "blue",
        arc_color: str = "red",
        show_centers: bool = True,
        **kwargs,
    ):
        """Visualize the decomposition of the base cross-section."""
        # For PolyPrism, we plot the base polygon (buffered_polygons[0] or polygons)
        if self.extrude:
            polygons = (
                self.polygons.geoms
                if hasattr(self.polygons, "geoms")
                else [self.polygons]
            )
            for polygon in polygons:
                vertices = [
                    self._parse_coords(coords) for coords in polygon.exterior.coords
                ]
                ax = super().plot_decomposition(
                    vertices,
                    ax=ax,
                    line_color=line_color,
                    arc_color=arc_color,
                    show_centers=show_centers,
                    identify_arcs=self.identify_arcs,
                    min_arc_points=self.min_arc_points,
                    arc_tolerance=self.arc_tolerance,
                    **kwargs,
                )
                for interior in polygon.interiors:
                    vertices = [
                        self._parse_coords(coords) for coords in interior.coords
                    ]
                    ax = super().plot_decomposition(
                        vertices,
                        ax=ax,
                        line_color=line_color,
                        arc_color=arc_color,
                        show_centers=show_centers,
                        identify_arcs=self.identify_arcs,
                        min_arc_points=self.min_arc_points,
                        arc_tolerance=self.arc_tolerance,
                        **kwargs,
                    )
        else:
            # For buffered polygons, we plot the first layer
            for entry in self.buffered_polygons:
                # entry is list of (z, polygon)
                z, polygon = entry[0]
                vertices = [(x, y, z) for x, y in polygon.exterior.coords]
                ax = super().plot_decomposition(
                    vertices,
                    ax=ax,
                    line_color=line_color,
                    arc_color=arc_color,
                    show_centers=show_centers,
                    identify_arcs=self.identify_arcs,
                    min_arc_points=self.min_arc_points,
                    arc_tolerance=self.arc_tolerance,
                    **kwargs,
                )
                for interior in polygon.interiors:
                    vertices = [(x, y, z) for x, y in interior.coords]
                    ax = super().plot_decomposition(
                        vertices,
                        ax=ax,
                        line_color=line_color,
                        arc_color=arc_color,
                        show_centers=show_centers,
                        identify_arcs=self.identify_arcs,
                        min_arc_points=self.min_arc_points,
                        arc_tolerance=self.arc_tolerance,
                        **kwargs,
                    )
        return ax

    def _create_occ_volume_with_holes(
        self, entry: list[tuple[float, Polygon]]
    ) -> TopoDS_Shape:
        """Create OCC volume with holes directly using OCP."""
        from OCP.BRepAlgoAPI import BRepAlgoAPI_Cut

        exterior = self._create_occ_volume(entry, exterior=True)
        # entry[0][1] is a Polygon
        interior_count = len(entry[0][1].interiors)

        for i in range(interior_count):
            interior = self._create_occ_volume(entry, exterior=False, interior_index=i)
            cut_api = BRepAlgoAPI_Cut(exterior, interior)
            cut_api.Build()
            exterior = cut_api.Shape()

        return exterior

    def instanciate_occ(self) -> TopoDS_Shape:
        """Create OCC volumes directly using OCP."""
        from OCP.BRepAlgoAPI import BRepAlgoAPI_Cut, BRepAlgoAPI_Fuse
        from OCP.BRepPrimAPI import BRepPrimAPI_MakePrism
        from OCP.gp import gp_Vec

        volumes = []
        if self.extrude:
            polys = (
                self.polygons.geoms
                if hasattr(self.polygons, "geoms")
                else [self.polygons]
            )
            for poly in polys:
                # Create face (with holes) at zmin
                exterior_vertices = [(x, y, self.zmin) for x, y in poly.exterior.coords]
                face = self._make_occ_face_from_vertices(
                    exterior_vertices,
                    identify_arcs=self.identify_arcs,
                    min_arc_points=self.min_arc_points,
                    arc_tolerance=self.arc_tolerance,
                )

                for interior in poly.interiors:
                    hole_vertices = [(x, y, self.zmin) for x, y in interior.coords]
                    hole_face = self._make_occ_face_from_vertices(
                        hole_vertices,
                        identify_arcs=self.identify_arcs,
                        min_arc_points=self.min_arc_points,
                        arc_tolerance=self.arc_tolerance,
                    )
                    cut_api = BRepAlgoAPI_Cut(face, hole_face)
                    cut_api.Build()
                    face = cut_api.Shape()

                # Extrude
                vec = gp_Vec(0, 0, self.zmax - self.zmin)
                prism_api = BRepPrimAPI_MakePrism(face, vec)
                # BRepPrimAPI classes often perform on Construction, or have Build/Perform.
                # Let's check prism_api. It usually has Build().
                prism_api.Build()
                volumes.append(prism_api.Shape())
        else:
            volumes.extend(
                [
                    self._create_occ_volume_with_holes(entry)
                    for entry in self.buffered_polygons
                ]
            )

        if not volumes:
            return None

        result = volumes[0]
        for v in volumes[1:]:
            fuse_api = BRepAlgoAPI_Fuse(result, v)
            fuse_api.Build()
            result = fuse_api.Shape()

        rotation_point = self._get_rotation_point(self.polygons)
        return self._apply_transformation_occ(result, rotation_point)

    def to_dict(self) -> dict:
        """Convert entity to dictionary representation.

        Returns:
            Dictionary containing serializable entity data
        """
        import shapely.wkt
        from shapely.geometry import MultiPolygon

        if isinstance(self.polygons, MultiPolygon):
            polygons_wkt = [
                shapely.wkt.dumps(p, rounding_precision=12) for p in self.polygons.geoms
            ]
        elif isinstance(self.polygons, list):
            polygons_wkt = [
                shapely.wkt.dumps(p, rounding_precision=12) for p in self.polygons
            ]
        else:
            polygons_wkt = [shapely.wkt.dumps(self.polygons, rounding_precision=12)]

        return {
            "type": "PolyPrism",
            "polygons_wkt": polygons_wkt,
            "buffers": {str(k): v for k, v in self.buffers.items()},
            "physical_name": self.physical_name,
            "mesh_order": self.mesh_order,
            "mesh_bool": self.mesh_bool,
            "additive": self.additive,
            "point_tolerance": self.point_tolerance,
            "identify_arcs": self.identify_arcs,
            "min_arc_points": self.min_arc_points,
            "arc_tolerance": self.arc_tolerance,
            "subdivision": list(self.subdivision) if self.subdivision else None,
            "translation": self.translation,
            "rotation_axis": self.rotation_axis,
            "rotation_point": self.rotation_point,
            "rotation_angle": self.rotation_angle,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "PolyPrism":
        """Reconstruct entity from dictionary representation.

        Args:
            data: Dictionary containing entity data

        Returns:
            PolyPrism instance
        """
        import shapely.wkt
        from shapely.geometry import MultiPolygon

        polygons = [shapely.wkt.loads(wkt) for wkt in data["polygons_wkt"]]
        polygons = MultiPolygon(polygons) if len(polygons) > 1 else polygons[0]

        buffers = {float(k): v for k, v in data["buffers"].items()}
        subdivision = tuple(data["subdivision"]) if data["subdivision"] else None

        return cls(
            polygons=polygons,
            buffers=buffers,
            physical_name=data["physical_name"],
            mesh_order=data["mesh_order"],
            mesh_bool=data["mesh_bool"],
            additive=data["additive"],
            point_tolerance=data["point_tolerance"],
            identify_arcs=data["identify_arcs"],
            min_arc_points=data["min_arc_points"],
            arc_tolerance=data["arc_tolerance"],
            subdivision=subdivision,
            translation=data.get("translation"),
            rotation_axis=data.get("rotation_axis"),
            rotation_point=data.get("rotation_point"),
            rotation_angle=data.get("rotation_angle", 0.0),
        )

    def _validate_polygon_buffers(self) -> bool:
        """Check if any buffering operation changes the topology of the polygon."""
        # Get first polygon or multipolygon
        first_geom = self.buffered_polygons[0][0][1]

        # Handle both single polygons and multipolygons
        first_polygons = (
            first_geom.geoms if hasattr(first_geom, "geoms") else [first_geom]
        )

        # Get reference counts from first polygon(s)
        reference_counts = []
        for polygon in first_polygons:
            # Store exterior vertex count and interior vertex counts for this polygon
            polygon_counts = {
                "exterior": len(polygon.exterior.coords),
                "interiors": [len(interior.coords) for interior in polygon.interiors],
            }
            reference_counts.append(polygon_counts)

        # Check each buffered polygon matches reference counts
        for buffered_polygon in self.buffered_polygons[0][1:]:
            geom = buffered_polygon[1]
            polygons = geom.geoms if hasattr(geom, "geoms") else [geom]

            if len(polygons) != len(reference_counts):
                return False

            for polygon, ref_counts in zip(polygons, reference_counts):
                # Check exterior vertices match
                if len(polygon.exterior.coords) != ref_counts["exterior"]:
                    return False

                # Check interior vertices match
                polygon_interior_counts = [
                    len(interior.coords) for interior in polygon.interiors
                ]
                if len(polygon_interior_counts) != len(ref_counts["interiors"]):
                    return False

                for count, ref_count in zip(
                    polygon_interior_counts, ref_counts["interiors"]
                ):
                    if count != ref_count:
                        return False

        return True
