from __future__ import annotations

import logging
from typing import Iterable, List, Sequence, Tuple

import networkx as nx
from centerline.geometry import Centerline
from shapely import line_merge
from shapely.geometry import (
    LineString,
    MultiLineString,
    MultiPoint,
    MultiPolygon,
    Point,
    Polygon,
)
from shapely.ops import linemerge, voronoi_diagram
from shapely.validation import explain_validity

__all__ = ["split_polygon"]

TOLERANCE: float = 1e-8  # snap grid size when deduplicating vertices


def _round_pt(x: float, y: float, tolerance: float = TOLERANCE) -> Tuple[float, float]:
    """Snap *x*, *y* to a *tolerance*‑sized grid so nearly‑identical vertices collapse."""
    return (round(x / tolerance) * tolerance, round(y / tolerance) * tolerance)


def _prune_adjacent_leaves(graph: nx.Graph, weight: str = "weight") -> None:
    """Remove redundant leaves *in‑place*.

    For each node that is **not** itself a leaf, keep only **one** incident leaf –‑
    specifically the one whose connecting edge has the highest ``weight`` attribute.

    When a node is connected to exactly **one** leaf, that leaf is removed as well.
    This considerably simplifies the graph while preserving the main branches needed
    for a meaningful centreline.
    """
    leaves_to_drop: List = []

    for parent in graph.nodes:
        leaves: List = [
            nbr for nbr in graph.neighbors(parent) if graph.degree[nbr] == 1
        ]
        if not leaves:
            continue

        if len(leaves) == 1:
            leaves_to_drop.append(leaves[0])
            continue

        # more than one leaf – keep only the *best* one
        best_leaf = max(leaves, key=lambda leaf: graph[parent][leaf].get(weight, 0))
        leaves_to_drop.extend(leaf for leaf in leaves if leaf is not best_leaf)

    graph.remove_nodes_from(leaves_to_drop)


def _extract_centerline(
    polygon: Polygon, spacing: float
) -> LineString | MultiLineString:
    """Return a de‑branched centreline for *polygon*.

    The heavy‑lifting is done by *Centerline* from the ``centerline`` package, but the
    raw result is *cleaned* by merging contiguous segments and pruning minor twigs so
    that downstream point‑generation behaves predictably.
    """
    if not (isinstance(polygon, Polygon) and polygon.is_valid):
        msg = f"Invalid polygon supplied: {explain_validity(polygon)}"
        raise ValueError(msg)

    if spacing <= 0:
        raise ValueError("spacing must be a positive number")

    raw: LineString | MultiLineString = line_merge(
        Centerline(polygon, spacing / 10).geometry
    )

    segments: Sequence[LineString] = (
        list(raw.geoms) if isinstance(raw, MultiLineString) else [raw]
    )

    graph: nx.Graph = nx.Graph()
    for idx, seg in enumerate(segments):
        start, end = _round_pt(*seg.coords[0]), _round_pt(*seg.coords[-1])
        graph.add_edge(start, end, weight=seg.length, idx=idx)

    _prune_adjacent_leaves(graph)

    kept_indices = {d["idx"] for *_, d in graph.edges(data=True)}
    cleaned: Iterable[LineString] = (
        seg for i, seg in enumerate(segments) if i in kept_indices
    )

    merged = linemerge(list(cleaned))
    return merged


def _spaced_points(lines: LineString | MultiLineString, spacing: float) -> List[Point]:
    """Return approximately equally‑spaced points *on every branch* of *lines*.

    * ``spacing`` refers to the *desired* interval. Because the total branch length is
      finite, the exact distance between adjacent points is *adjusted* (up to ±½ spacing)
      so that the points are perfectly evenly distributed along each branch.
    * No point is placed **exactly** at a junction or branch‑start – the first point is
      offset by ½ adjusted spacing to avoid duplicates between branches.
    """
    if spacing <= 0:
        raise ValueError("spacing must be a positive number")

    merged: LineString | MultiLineString = line_merge(lines)
    branches: Iterable[LineString] = (
        merged.geoms if isinstance(merged, MultiLineString) else [merged]
    )

    points: List[Point] = []
    for branch in branches:
        length = branch.length
        if length < spacing:
            continue

        # number of internal points, excluding start & end
        n_points = max(int((length - spacing) // spacing) + 1, 1)
        adjusted = length / n_points

        for k in range(n_points):
            # offset so first point not exactly at vertex
            dist = (k + 0.5) * adjusted
            points.append(branch.interpolate(dist))

    return points


def _voronoi_split(poly: Polygon, seeds: Sequence[Point]) -> MultiPolygon:
    """Generate a Voronoi diagram from *seeds* and clip it to *poly*."""
    if not seeds:
        raise ValueError("At least one seed point is required for a Voronoi diagram")

    diagram = voronoi_diagram(MultiPoint(list(seeds)), envelope=poly)

    cells: List[Polygon] = []
    for cell in diagram.geoms:
        clipped = cell.intersection(poly)
        if clipped.is_empty:
            continue
        if isinstance(clipped, Polygon):
            cells.append(clipped)
        elif isinstance(clipped, MultiPolygon):
            cells.extend(clipped.geoms)

    result = MultiPolygon(cells)
    return result


def split_polygon(polygon: Polygon, spacing: float = 5.0) -> MultiPolygon:
    """Split *polygon* into Voronoi cells derived from an approximated centreline.

    Parameters
    ----------
    polygon
        The input geometry to split. Must be a valid, *single* ``Polygon`` (holes are OK).
    spacing
        Approximate distance between consecutive seed points along the centreline.

    Returns
    -------
    MultiPolygon
        A collection of sub‑polygons whose union equals (or is very close to)
        the original *polygon*.

    Notes
    -----
    * A smaller ``spacing`` yields more seed points and therefore more resulting
      polygons at the expense of increased run time.
    * Because of floating‑point arithmetic and the nature of the Voronoi algorithm,
      the union of the returned cells may be *slightly* smaller than the original
      geometry. Buffering the input by a *tiny* amount (~1e‑8) beforehand can help
      if strict area preservation is important.
    """

    try:
        centreline = _extract_centerline(polygon, spacing)
        seeds = _spaced_points(centreline, spacing)
        result = _voronoi_split(polygon, seeds)
        logging.info(
            f"Split polygon into {len(result.geoms)} cells with spacing {spacing}"
        )
        return result
    except Exception:
        # Provide context and re‑raise – callers can catch the same *exc*
        raise
