# Std lib
from typing import List, Union

# 3rd party lib
import numpy as np
from numpy.typing import NDArray
from shapely import buffer, Polygon, MultiPolygon, prepare

# Local lib
from .helper_functions import interpolate_curve
from ..beams import BeamType
from ..mill import Mill


def contour_inward(shape: Union[Polygon, MultiPolygon], mill: Mill, beam: BeamType,
                   rasterize: bool = True) -> List[NDArray]:
    beam_radius = beam.get_radius(units="nm")
    overlap = mill.get_overlap()
    step = 2 * beam_radius * (1 - overlap)
    dwelltime = int(mill.get_dwell_time("ns")) / 100

    contours = []
    prepare(shape)  # Precompute spatial indexes for faster operations

    def add_linestrings(geom):
        if isinstance(geom, Polygon):
            if not geom.exterior.is_empty:
                contours.append(geom.exterior)
            contours.extend(geom.interiors)
        elif isinstance(geom, MultiPolygon):
            for poly in geom.geoms:
                if not poly.exterior.is_empty:
                    contours.append(poly.exterior)
                contours.extend(poly.interiors)

    buffered = buffer(shape, -beam_radius, join_style="mitre")
    quad_segs = 32  # Reduced from 100 for initial buffer
    k = 0

    while True:
        if k != 0:
            quad_segs = max(8, int(quad_segs * 0.8))  # Decrease resolution inward
            buffered = buffer(buffered, -step, join_style="round", quad_segs=quad_segs)
            if buffered.is_empty:
                break
        add_linestrings(buffered)
        k += 1

    streamlines = []
    if not rasterize:
        for line in contours:
            coords = np.array(line.coords)
            streamline = np.column_stack((
                np.zeros(len(coords)),
                coords[:, 0],
                coords[:, 1],
                np.zeros(len(coords))
            ))
            streamlines.append(streamline)
        return streamlines

    for line in contours:
        line_length = line.length
        nr_steps = int(np.rint(line_length / step))

        if line.is_closed and nr_steps == 3:
            if line_length / 2 > beam_radius * 1.5:
                points = interpolate_curve(line, step)
                centroid = line.centroid
                if line.distance(centroid) > beam_radius:
                    points = np.vstack([points, [centroid.x, centroid.y]])
            else:
                points = np.array([[line.centroid.x, line.centroid.y]])
        elif not line.is_closed and nr_steps == 2:
            points = np.array([[line.centroid.x, line.centroid.y]])
        else:
            points = interpolate_curve(line, step)

        if points.size == 0:
            continue

        codes = np.full(len(points), 2, dtype=int)
        codes[0] = 1
        dwelltimes = np.full(len(points), dwelltime)
        streamline = np.column_stack((dwelltimes, points[:, 0], points[:, 1], codes))

        streamlines.append(streamline)

    return streamlines


def contour_outward(shape: Polygon | MultiPolygon, mill: Mill, beam: BeamType, rasterize: bool = True) -> List[NDArray]:
    return contour_inward(shape, mill, beam, rasterize)[::-1]

