import json
import math
import time

from .background import remove_image_background


def polygon_bounds(polygon):
    import numpy as np

    points = np.asarray(polygon, dtype=np.float64)
    if points.ndim != 2 or points.shape[0] < 3 or points.shape[1] != 2:
        raise ValueError("The selection does not define a valid image region")
    if not np.isfinite(points).all():
        raise ValueError("The selection contains invalid image coordinates")
    left = int(math.floor(float(points[:, 0].min())))
    top = int(math.floor(float(points[:, 1].min())))
    right = int(math.ceil(float(points[:, 0].max())))
    bottom = int(math.ceil(float(points[:, 1].max())))
    if right <= left or bottom <= top:
        raise ValueError("The selection does not define a valid image region")
    return left, top, right, bottom


def polygon_selection_mask(image_size, polygon, fill_rule="EVEN_ODD"):
    """Rasterize a polygon with an explicit fill rule."""
    import numpy as np

    width, height = (int(value) for value in image_size)
    points = np.asarray(polygon, dtype=np.float64)
    if width <= 0 or height <= 0:
        raise ValueError("The selection canvas has no readable dimensions")
    if points.ndim != 2 or points.shape[0] < 3 or points.shape[1] != 2:
        raise ValueError("The selection does not define a valid image region")
    if not np.isfinite(points).all():
        raise ValueError("The selection contains invalid image coordinates")
    if fill_rule not in {"EVEN_ODD", "NONZERO"}:
        raise ValueError(f"Unknown polygon fill rule: {fill_rule}")

    scanlines = np.arange(height, dtype=np.float64) + 0.5
    crossings = [[] for _index in range(height)]
    for start, end in zip(points, np.roll(points, -1, axis=0)):
        selected = np.flatnonzero(
            (start[1] > scanlines) != (end[1] > scanlines)
        )
        if not len(selected):
            continue
        crossing_x = start[0] + (
            (scanlines[selected] - start[1])
            * (end[0] - start[0])
            / (end[1] - start[1])
        )
        winding = 1 if end[1] > start[1] else -1
        for row, value in zip(selected, crossing_x):
            crossings[int(row)].append((float(value), winding))

    mask = np.zeros((height, width), dtype=bool)
    for row, values in enumerate(crossings):
        values.sort()
        state = 0
        index = 0
        while index < len(values):
            crossing_x = values[index][0]
            winding = 0
            count = 0
            while index < len(values) and values[index][0] == crossing_x:
                winding += values[index][1]
                count += 1
                index += 1
            state = (
                state + winding
                if fill_rule == "NONZERO"
                else state ^ (count % 2)
            )
            if index >= len(values) or state == 0:
                continue
            following_x = values[index][0]
            first = max(0, int(math.ceil(crossing_x - 0.5)))
            last = min(width, int(math.floor(following_x - 0.5)) + 1)
            mask[row, first:last] = True
    return mask

def create_selection_canvas(input_path, polygon, fill_rule="EVEN_ODD"):
    """Create an unbounded selection canvas with transparent black outside."""
    import numpy as np
    from PIL import Image

    with Image.open(input_path) as opened_image:
        source = opened_image.convert("RGBA")
    left, top, right, bottom = polygon_bounds(polygon)
    if right <= 0 or bottom <= 0 or left >= source.width or top >= source.height:
        raise ValueError("The selection does not overlap the source image")
    canvas = Image.new("RGBA", (right - left, bottom - top), (0, 0, 0, 0))
    canvas.paste(source, (-left, -top))
    relative = [(float(x) - left, float(y) - top) for x, y in polygon]
    selection_mask = polygon_selection_mask(canvas.size, relative, fill_rule)
    pixels = np.asarray(canvas, dtype=np.uint8).copy()
    pixels[~selection_mask] = 0
    return Image.fromarray(pixels, mode="RGBA"), selection_mask, (
        left,
        top,
        right,
        bottom,
    )


def apply_selection_mask(foreground, selection_mask):
    """Clear pixels outside the source selection after background removal."""
    import numpy as np
    from PIL import Image

    pixels = np.asarray(foreground.convert("RGBA"), dtype=np.uint8).copy()
    if pixels.shape[:2] != selection_mask.shape:
        pixels = np.asarray(
            foreground.convert("RGBA").resize(
                (selection_mask.shape[1], selection_mask.shape[0]),
                Image.Resampling.BILINEAR,
            ),
            dtype=np.uint8,
        )
    pixels[~selection_mask] = 0
    pixels[pixels[:, :, 3] == 0, :3] = 0
    return Image.fromarray(pixels, mode="RGBA")


def crop_to_alpha(image, bounds):
    """Crop an RGBA mask image to its nonzero Alpha bounds."""
    import numpy as np
    from PIL import Image

    pixels = np.asarray(image.convert("RGBA"), dtype=np.uint8)
    visible_y, visible_x = np.nonzero(pixels[:, :, 3] > 0)
    if not len(visible_x):
        raise ValueError("The selection contains no visible image pixels")
    content_left = int(visible_x.min())
    content_top = int(visible_y.min())
    content_right = int(visible_x.max()) + 1
    content_bottom = int(visible_y.max()) + 1
    output = Image.new(
        "RGBA",
        (
            content_right - content_left,
            content_bottom - content_top,
        ),
        (0, 0, 0, 0),
    )
    output.paste(image, (-content_left, -content_top))
    left, top, _right, _bottom = bounds
    return output, (
        left + content_left,
        top + content_top,
        left + content_right,
        top + content_bottom,
    )


def run(context, parameters):
    """Refine a gesture mask with BEN2 and return its source-image bounds."""
    started = time.perf_counter()
    context.progress(0.05, "Preparing image selection")
    context.check_cancelled()
    output_directory = context.directory
    output_directory.mkdir(parents=True, exist_ok=True)
    canvas, selection_mask, bounds = create_selection_canvas(
        parameters["input"], parameters["polygon"]
    )
    selection_path = output_directory / "selection-input.png"
    canvas.save(selection_path, format="PNG", compress_level=1)
    model_manager = context.resource("model_manager")
    context.progress(0.2, f"Loading BEN2 for {parameters['device'].upper()}")
    foreground, _load_ms, _inference_ms = remove_image_background(
        selection_path,
        model_manager.directory(parameters["model"]),
        parameters["device"],
        model_manager,
        context.check_cancelled,
    )
    context.check_cancelled()
    context.progress(0.9, "Applying selection mask")
    output, bounds = crop_to_alpha(
        apply_selection_mask(foreground, selection_mask), bounds
    )
    output_path = output_directory / "image.png"
    output.save(output_path, format="PNG", compress_level=1)
    bounds_path = output_directory / "bounds.json"
    bounds_path.write_text(json.dumps(bounds, separators=(",", ":")), encoding="utf-8")
    elapsed_ms = (time.perf_counter() - started) * 1000.0
    message = f"Image selection refined in {elapsed_ms:.0f} ms"
    context.progress(1.0, message)
    return {"image": output_path.name, "bounds": bounds_path.name}