import json
import math

import bpy

from ...properties import (
    image_edit_refine_selection_property as refine_selection_property,
)
from ...runtime import JobOperatorBase
from ..balloon.operator import (
    deserialize_matrix,
    drawing_in_region,
    hidden_string_property,
    screen_path_to_image_pixels,
    serialize_matrix,
    simplify_screen_path,
)
from ..common import HIDDEN, cleanup_temporary_input, push_undo
from ..image import (
    is_image_empty,
    plane_bounds,
    prepare_image_input,
    remove_temporary_input,
    video_total_frames,
)
from .common import (
    active_view3d_tool_id,
    apply_image_mask_settings,
    crop_pixels_to_alpha,
    draw_dashed_line,
    draw_image_mask_operator_settings,
    draw_polygon_fill,
    ImageEditWarning,
    image_edit_poll,
    image_mask_refine_request,
    polygon_bounds,
    replace_empty_image_pixels,
    report_image_edit_exception,
    resolve_image_edit_click,
)


GESTURE_TOOL_IDS = {
    "BOX": "anyimage.image_box",
    "LASSO": "anyimage.image_lasso",
    "LINE": "anyimage.image_line",
    "POLYLINE": "anyimage.image_polyline",
}
LINE_SNAP_RADIANS = math.radians(5.0)
POLYLINE_CLOSE_RADIUS = 8.0


def points_inside_polygon(points, polygon):
    import numpy as np

    points = np.asarray(points, dtype=np.float64)
    polygon = np.asarray(polygon, dtype=np.float64)
    inside = np.zeros(len(points), dtype=bool)
    previous = polygon[-1]
    for current in polygon:
        crosses = (current[1] > points[:, 1]) != (previous[1] > points[:, 1])
        edge_x = (
            (previous[0] - current[0])
            * (points[:, 1] - current[1])
            / (previous[1] - current[1] + 1e-300)
            + current[0]
        )
        inside ^= crosses & (points[:, 0] < edge_x)
        previous = current
    return inside


def polygon_mask(image_size, polygon):
    import numpy as np

    width, height = (int(value) for value in image_size)
    if width <= 0 or height <= 0:
        raise ValueError("The source image has no readable dimensions")
    polygon_bounds(polygon)
    y, x = np.indices((height, width), dtype=np.float64)
    centers = np.column_stack((x.ravel() + 0.5, y.ravel() + 0.5))
    return points_inside_polygon(centers, polygon).reshape((height, width)).astype(
        np.float32
    )


def refined_mask_from_pixels(mask_pixels, mask_size, bounds, source_size):
    import numpy as np

    mask_width, mask_height = (int(value) for value in mask_size)
    source_width, source_height = (int(value) for value in source_size)
    local = np.flipud(
        np.asarray(mask_pixels, dtype=np.float32).reshape(
            (mask_height, mask_width, 4)
        )
    )[:, :, 3]
    result = np.zeros((source_height, source_width), dtype=np.float32)
    left, top, right, bottom = (int(value) for value in bounds)
    source_left = max(left, 0)
    source_top = max(top, 0)
    source_right = min(right, source_width)
    source_bottom = min(bottom, source_height)
    if source_right <= source_left or source_bottom <= source_top:
        return result
    result[source_top:source_bottom, source_left:source_right] = local[
        source_top - top : source_bottom - top,
        source_left - left : source_right - left,
    ]
    return result


def edit_image_mask_pixels(source_pixels, source_size, selection_mask, invert=False):
    """Apply a mask while preserving source RGB, then tightly crop the canvas."""
    import numpy as np

    width, height = (int(value) for value in source_size)
    rgba = np.flipud(
        np.asarray(source_pixels, dtype=np.float32).reshape((height, width, 4))
    ).copy()
    mask = np.asarray(selection_mask, dtype=np.float32)
    if mask.shape != (height, width):
        raise ValueError("The refined selection does not match the source image")
    mask = np.clip(mask, 0.0, 1.0)
    if invert:
        mask = 1.0 - mask
    rgba[:, :, 3] *= mask
    return crop_pixels_to_alpha(np.flipud(rgba).ravel(), (0, 0, width, height))


def edit_image_object(source_object, selection_mask, invert=False):
    import numpy as np

    image = source_object.data
    source_size = tuple(image.size)
    source_pixels = np.empty(len(image.pixels), dtype=np.float32)
    image.pixels.foreach_get(source_pixels)
    pixels, crop_bounds = edit_image_mask_pixels(
        source_pixels, source_size, selection_mask, invert
    )
    left, top, right, bottom = crop_bounds
    return replace_empty_image_pixels(
        source_object,
        pixels,
        (right - left, bottom - top),
        crop_bounds,
    )


def gesture_property():
    return bpy.props.EnumProperty(
        items=(
            ("BOX", "Box", "Edit with a rectangular gesture"),
            ("LASSO", "Lasso", "Edit with a freehand gesture"),
            ("LINE", "Line", "Edit one side of a line"),
            ("POLYLINE", "Polyline", "Edit with a clicked polygon"),
        ),
        default="BOX",
        options={"HIDDEN", "SKIP_SAVE"},
    )


def box_screen_path(start, end):
    x0, y0 = start
    x1, y1 = end
    if abs(x1 - x0) < 2.0 or abs(y1 - y0) < 2.0:
        return ()
    return ((x0, y0), (x1, y0), (x1, y1), (x0, y1))


def snapped_line_endpoint(start, end, snap):
    if not snap:
        return tuple(float(value) for value in end)
    delta_x = float(end[0]) - float(start[0])
    delta_y = float(end[1]) - float(start[1])
    length = math.hypot(delta_x, delta_y)
    if length <= 0.0:
        return tuple(float(value) for value in end)
    angle = round(math.atan2(delta_y, delta_x) / LINE_SNAP_RADIANS)
    angle *= LINE_SNAP_RADIANS
    return (
        float(start[0]) + math.cos(angle) * length,
        float(start[1]) + math.sin(angle) * length,
    )


def line_half_plane_polygon(start, end, rectangle_size):
    """Clip a rectangle to the left side of one directed line."""
    width, height = (float(value) for value in rectangle_size)
    if math.dist(start, end) < 2.0:
        return ()
    vertices = [(0.0, 0.0), (width, 0.0), (width, height), (0.0, height)]

    def side(point):
        return (end[0] - start[0]) * (point[1] - start[1]) - (
            end[1] - start[1]
        ) * (point[0] - start[0])

    result = []
    previous = vertices[-1]
    previous_side = side(previous)
    for current in vertices:
        current_side = side(current)
        if (current_side >= 0.0) != (previous_side >= 0.0):
            factor = previous_side / (previous_side - current_side)
            result.append(
                (
                    previous[0] + (current[0] - previous[0]) * factor,
                    previous[1] + (current[1] - previous[1]) * factor,
                )
            )
        if current_side >= 0.0:
            result.append(current)
        previous = current
        previous_side = current_side
    return tuple(result)


def image_line_half_plane_polygon(line, image_size):
    """Keep the image-space half-plane shown by the screen-space preview."""
    if len(line) != 2:
        raise ValueError("Drag a longer line across the image plane")
    return line_half_plane_polygon(line[1], line[0], image_size)


class EditImageMask(bpy.types.Operator):
    bl_idname = "anyimage.edit_image_mask"
    bl_label = "Edit Image Mask"
    bl_description = "Edit and tightly crop an Image Empty with a gesture"
    bl_options = {"REGISTER", "UNDO"}

    source_object_name: hidden_string_property()
    source_matrix_data: hidden_string_property()
    selection_data: hidden_string_property()
    refine_selection: refine_selection_property()
    invert: bpy.props.BoolProperty(
        name="Invert",
        description="Keep the opposite side of the refined gesture mask",
        default=False,
        options={"SKIP_SAVE"},
    )
    gesture: gesture_property()

    _path = None
    _cursor = None
    _handle = None
    _area_pointer = 0
    _region_pointer = 0
    _region_size = (0, 0)

    def draw(self, _context):
        draw_image_mask_operator_settings(self.layout, self)

    @classmethod
    def poll(cls, context):
        return image_edit_poll(context)

    def _screen_polygon(self):
        if self.gesture == "BOX":
            return box_screen_path(self._path[0], self._cursor)
        if self.gesture == "LASSO":
            return tuple(simplify_screen_path(self._path))
        if self.gesture == "LINE":
            return line_half_plane_polygon(
                self._path[0], self._cursor, self._region_size
            )
        vertices = list(self._path)
        if self._cursor is not None:
            vertices.append(self._cursor)
        return tuple(vertices)

    def _draw_selection(self):
        if not self._path or not drawing_in_region(
            self._area_pointer, self._region_pointer
        ):
            return
        polygon = self._screen_polygon()
        if len(polygon) >= 3:
            draw_polygon_fill(
                polygon,
                (1.0, 0.25, 0.2, 0.16)
                if self.invert
                else (1.0, 1.0, 1.0, 0.16),
            )
        if self.gesture == "LINE":
            draw_dashed_line((self._path[0], self._cursor))
            return
        if len(polygon) > 1:
            draw_dashed_line((*polygon, polygon[0]))
        if (
            self.gesture == "POLYLINE"
            and len(self._path) >= 3
            and self._cursor is not None
            and math.dist(self._cursor, self._path[0]) <= POLYLINE_CLOSE_RADIUS
        ):
            self._draw_polyline_close_indicator()

    def _draw_polyline_close_indicator(self):
        import gpu
        from gpu_extras.batch import batch_for_shader

        center_x, center_y = self._path[0]
        vertices = tuple(
            (
                center_x + math.cos(index * math.tau / 24.0) * POLYLINE_CLOSE_RADIUS,
                center_y + math.sin(index * math.tau / 24.0) * POLYLINE_CLOSE_RADIUS,
            )
            for index in range(25)
        )
        shader = gpu.shader.from_builtin("UNIFORM_COLOR")
        batch = batch_for_shader(shader, "LINE_STRIP", {"pos": vertices})
        gpu.state.blend_set("ALPHA")
        gpu.state.line_width_set(2.0)
        shader.bind()
        shader.uniform_float("color", (1.0, 1.0, 1.0, 0.95))
        batch.draw(shader)
        gpu.state.line_width_set(1.0)
        gpu.state.blend_set("NONE")

    def _set_status(self, context):
        invert = "On" if self.invert else "Off"
        if self.gesture == "POLYLINE":
            action = "click points; click start or Enter to finish; Backspace removes"
        elif self.gesture == "LINE":
            action = "drag line; Ctrl snaps 5°; release to edit"
        else:
            action = "drag and release to edit"
        context.workspace.status_text_set(
            f"Image {self.gesture.title()}: {action}; F toggles Invert ({invert}); "
            "RMB or Esc cancels"
        )

    def _finish(self, context):
        if self._handle is not None:
            bpy.types.SpaceView3D.draw_handler_remove(self._handle, "WINDOW")
            self._handle = None
        context.workspace.status_text_set(None)
        if context.area is not None:
            context.area.tag_redraw()

    def cancel(self, context):
        self._finish(context)

    def invoke(self, context, event):
        source_object, should_edit = resolve_image_edit_click(
            context, event, self.report
        )
        if not should_edit:
            return {"FINISHED"} if source_object is not None else {"CANCELLED"}
        if source_object is None:
            return {"CANCELLED"}
        if video_total_frames(source_object.data) > 1:
            self.report({"WARNING"}, "Image tools only support a still image")
            return {"CANCELLED"}
        apply_image_mask_settings(self, context)
        self.invert = False
        self.source_object_name = source_object.name
        self.source_matrix_data = serialize_matrix(source_object.matrix_world)
        self.selection_data = ""
        point = (float(event.mouse_region_x), float(event.mouse_region_y))
        self._path = [point]
        self._cursor = point
        self._area_pointer = context.area.as_pointer()
        self._region_pointer = context.region.as_pointer()
        self._region_size = (context.region.width, context.region.height)
        self._handle = bpy.types.SpaceView3D.draw_handler_add(
            self._draw_selection, (), "WINDOW", "POST_PIXEL"
        )
        self._set_status(context)
        context.window_manager.modal_handler_add(self)
        context.area.tag_redraw()
        return {"RUNNING_MODAL"}

    def _complete(self, context):
        source_object = bpy.data.objects.get(self.source_object_name)
        if not is_image_empty(source_object):
            raise RuntimeError("The source Image Empty is no longer available")
        matrix = deserialize_matrix(self.source_matrix_data)
        if self.gesture == "LINE":
            line = screen_path_to_image_pixels(
                (self._path[0], self._cursor),
                context.region,
                context.region_data,
                matrix,
                plane_bounds(source_object),
                tuple(source_object.data.size),
            )
            polygon = image_line_half_plane_polygon(
                line, tuple(source_object.data.size)
            )
        else:
            path = self._screen_polygon()
            if len(path) < 3:
                self._finish(context)
                return {"CANCELLED"}
            polygon = screen_path_to_image_pixels(
                path,
                context.region,
                context.region_data,
                matrix,
                plane_bounds(source_object),
                tuple(source_object.data.size),
            )
        polygon_bounds(polygon)
        self.selection_data = json.dumps(polygon, separators=(",", ":"))
        self._finish(context)
        return self.execute(context)

    def modal(self, context, event):
        active_tool = active_view3d_tool_id(context)
        if active_tool is not None and active_tool != GESTURE_TOOL_IDS[self.gesture]:
            self._finish(context)
            return {"CANCELLED"}
        if event.type in {"ESC", "RIGHTMOUSE"} and event.value == "PRESS":
            self._finish(context)
            return {"CANCELLED"}
        if event.type == "F" and event.value == "PRESS":
            self.invert = not self.invert
            self._set_status(context)
            context.area.tag_redraw()
            return {"RUNNING_MODAL"}
        if self.gesture == "POLYLINE":
            return self._modal_polyline(context, event)
        if event.type == "MOUSEMOVE":
            point = (float(event.mouse_region_x), float(event.mouse_region_y))
            if self.gesture == "LASSO":
                self._path.append(point)
                self._cursor = point
            elif self.gesture == "LINE":
                self._cursor = snapped_line_endpoint(self._path[0], point, event.ctrl)
            else:
                self._cursor = point
            context.area.tag_redraw()
            return {"RUNNING_MODAL"}
        if event.type == "LEFTMOUSE" and event.value == "RELEASE":
            point = (float(event.mouse_region_x), float(event.mouse_region_y))
            if self.gesture == "LASSO":
                self._path.append(point)
                self._cursor = point
            elif self.gesture == "LINE":
                self._cursor = snapped_line_endpoint(self._path[0], point, event.ctrl)
                if math.dist(self._path[0], self._cursor) < 2.0:
                    self._finish(context)
                    return {"CANCELLED"}
            else:
                self._cursor = point
            try:
                return self._complete(context)
            except (RuntimeError, TypeError, ValueError) as error:
                report_image_edit_exception(self, error)
                self._finish(context)
                return {"CANCELLED"}
        return {"PASS_THROUGH"}

    def _modal_polyline(self, context, event):
        if event.type == "BACK_SPACE" and event.value == "PRESS":
            if len(self._path) > 1:
                self._path.pop()
            context.area.tag_redraw()
            return {"RUNNING_MODAL"}
        if event.type in {"RET", "NUMPAD_ENTER"} and event.value == "PRESS":
            return self._confirm_polyline(context)
        if event.type == "MOUSEMOVE":
            self._cursor = (
                float(event.mouse_region_x),
                float(event.mouse_region_y),
            )
            context.area.tag_redraw()
            return {"RUNNING_MODAL"}
        if event.type == "LEFTMOUSE" and event.value == "PRESS":
            point = (float(event.mouse_region_x), float(event.mouse_region_y))
            if (
                len(self._path) >= 3
                and math.dist(point, self._path[0]) <= POLYLINE_CLOSE_RADIUS
            ):
                return self._confirm_polyline(context)
            if math.dist(point, self._path[-1]) >= 2.0:
                self._path.append(point)
            context.area.tag_redraw()
            return {"RUNNING_MODAL"}
        return {"PASS_THROUGH"}

    def _confirm_polyline(self, context):
        self._cursor = self._path[0]
        try:
            return self._complete(context)
        except (RuntimeError, TypeError, ValueError) as error:
            report_image_edit_exception(self, error)
            self._finish(context)
            return {"CANCELLED"}

    def execute(self, context):
        source_object = bpy.data.objects.get(self.source_object_name)
        if not is_image_empty(source_object):
            self.report({"ERROR"}, "The source Image Empty is no longer available")
            return {"CANCELLED"}
        try:
            polygon = json.loads(self.selection_data)
            polygon_bounds(polygon)
            if video_total_frames(source_object.data) > 1:
                raise RuntimeError("Image tools only support a still image")
            if self.refine_selection:
                source_path, delete_input = prepare_image_input(
                    source_object.data, context.scene
                )
                result = bpy.ops.anyimage.refine_image_mask(
                    input_path=str(source_path),
                    delete_input=delete_input,
                    source_object_name=self.source_object_name,
                    selection_data=self.selection_data,
                    invert=self.invert,
                )
                if "CANCELLED" in result:
                    if delete_input:
                        remove_temporary_input(source_path)
                    return {"CANCELLED"}
                return {"FINISHED"}
            edit_image_object(
                source_object,
                polygon_mask(tuple(source_object.data.size), polygon),
                self.invert,
            )
        except (json.JSONDecodeError, RuntimeError, TypeError, ValueError) as error:
            report_image_edit_exception(self, error)
            return {"CANCELLED"}
        return {"FINISHED"}


class RefineImageMask(JobOperatorBase, bpy.types.Operator):
    bl_idname = "anyimage.refine_image_mask"
    bl_label = "Refine Image Mask"
    job_type = "refine-image-mask"
    starting_message = "Refining image selection"

    input_path: bpy.props.StringProperty(options=HIDDEN)
    delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN)
    source_object_name: hidden_string_property()
    selection_data: hidden_string_property()
    invert: bpy.props.BoolProperty(default=False, options=HIDDEN)

    def request(self, _context):
        return image_mask_refine_request(self.input_path, self.selection_data)

    def response(self, _context, result):
        source_object = bpy.data.objects.get(self.source_object_name)
        if not is_image_empty(source_object):
            raise RuntimeError("The source Image Empty is no longer available")
        import numpy as np

        bounds = tuple(json.loads(result.file("bounds").read_text(encoding="utf-8")))
        refined_image = bpy.data.images.load(
            str(result.file("image")), check_existing=False
        )
        try:
            refined_pixels = np.empty(len(refined_image.pixels), dtype=np.float32)
            refined_image.pixels.foreach_get(refined_pixels)
            mask = refined_mask_from_pixels(
                refined_pixels,
                tuple(refined_image.size),
                bounds,
                tuple(source_object.data.size),
            )
        finally:
            bpy.data.images.remove(refined_image, do_unlink=True)
        try:
            edit_image_object(source_object, mask, self.invert)
        except ImageEditWarning as error:
            self.report({"WARNING"}, str(error))
            return str(error)
        push_undo("Edit Image Mask")
        return "Image selection refined"

    def cleanup(self):
        cleanup_temporary_input(self.input_path, self.delete_input)


class ActivateImageEditTool(bpy.types.Operator):
    bl_idname = "anyimage.activate_image_edit_tool"
    bl_label = "Image Edit Tool"
    bl_description = "Activate the Image Box tool in the 3D Viewport"
    bl_options = {"INTERNAL"}

    @classmethod
    def poll(cls, context):
        return image_edit_poll(context)

    def execute(self, _context):
        return bpy.ops.wm.tool_set_by_id(
            name="anyimage.image_box", space_type="VIEW_3D"
        )