from pyscript import document, fetch as pyfetch
from pyscript.ffi import create_proxy
import asyncio
import json

document.getElementById("shape-options").style.display = "flex"

image = document.getElementById("annotation-image")
zone = document.getElementById("annotation-zone")
confirm_button = document.getElementById("annotation-confirm")
cancel_button = document.getElementById("annotation-cancel")
backspace_button = document.getElementById("annotation-backspace")
delete_button = document.getElementById("annotation-delete")
previous_button = document.getElementById("annotation-previous")
next_button = document.getElementById("annotation-next")
save_button = document.getElementById("annotation-save")

object_list = document.getElementById("object-types")

confirm_button.style.display = "none"
cancel_button.style.display = "none"
backspace_button.style.display = "none"
delete_button.style.display = "none"
previous_button.style.display = "none"
next_button.style.display = "none"
shape_type = ""
bbox_pos = None
new_shape = None
selected_shape = None


def make_shape_container():
    shape = document.createElementNS("http://www.w3.org/2000/svg", "svg")
    shape.setAttribute("width", "100%")
    shape.setAttribute("height", "100%")
    shape.setAttribute("viewBox", f"0 0 {image.naturalWidth} {image.naturalHeight}")
    shape.classList.add("shape-container")

    return shape


async def get_all_objects():
    response = await pyfetch("/api/object-types")
    if response.ok:
        return await response.json()


def follow_cursor(event):
    rect = zone.getBoundingClientRect()
    x = event.clientX - rect.left
    y = event.clientY - rect.top
    vertical_ruler.style.left = str(x) + "px"
    horizontal_ruler.style.top = str(y) + "px"


def change_object_type(event):
    global selected_shape
    if selected_shape is None:
        return
    selected_shape.setAttribute("data-object-type", event.currentTarget.value)


change_object_type_proxy = create_proxy(change_object_type)


def list_shapes():
    shapes = list(zone.getElementsByClassName("shape"))
    json_shapes = []
    for shape in shapes:
        shape_dict = {}
        if shape.tagName == "rect":
            shape_dict["type"] = "bbox"
            shape_dict["shape"] = {
                "x": float(shape.getAttribute("x")) / image.naturalWidth,
                "y": float(shape.getAttribute("y")) / image.naturalHeight,
                "w": float(shape.getAttribute("width")) / image.naturalWidth,
                "h": float(shape.getAttribute("height")) / image.naturalHeight
            }
        elif shape.tagName == "polygon" or shape.tagName == "polyline":
            if shape.tagName == "polygon":
                shape_dict["type"] = "polygon"
            elif shape.tagName == "polyline":
                shape_dict["type"] = "polyline"

            points = shape.getAttribute("points").split(" ")
            json_points = []
            for point in points:
                x, y = point.split(",")
                x, y = float(x), float(y)
                json_points.append({
                    "x": x / image.naturalWidth,
                    "y": y / image.naturalHeight
                })

            shape_dict["shape"] = json_points
        elif shape.tagName == "circle" and shape.classList.contains("shape-point"):
            shape_dict["type"] = "point"
            shape_dict["shape"] = {
                "x": float(shape.getAttribute("cx")) / image.naturalWidth,
                "y": float(shape.getAttribute("cy")) / image.naturalHeight
            }
        else:
            continue

        shape_dict["object"] = shape.getAttribute("data-object-type")
        json_shapes.append(shape_dict)

    return json_shapes


def put_shapes(json_shapes):
    for shape in json_shapes:
        new_shape = make_shape_container()
        zone_rect = zone.getBoundingClientRect()

        if shape["type"] == "bbox":
            rectangle = document.createElementNS("http://www.w3.org/2000/svg", "rect")
            rectangle.setAttribute("x", str(shape["shape"]["x"] * image.naturalWidth))
            rectangle.setAttribute("y", str(shape["shape"]["y"] * image.naturalHeight))
            rectangle.setAttribute("width", str(shape["shape"]["w"] * image.naturalWidth))
            rectangle.setAttribute("height", str(shape["shape"]["h"] * image.naturalHeight))
            rectangle.setAttribute("fill", "none")
            rectangle.setAttribute("data-object-type", shape["object"] or "")
            rectangle.classList.add("shape-bbox")
            rectangle.classList.add("shape")
            new_shape.appendChild(rectangle)
        elif shape["type"] == "polygon" or shape["type"] == "polyline":
            polygon = document.createElementNS("http://www.w3.org/2000/svg", shape["type"])
            points = " ".join(
                [f"{point['x'] * image.naturalWidth},{point['y'] * image.naturalHeight}" for point in shape["shape"]])
            polygon.setAttribute("points", points)
            polygon.setAttribute("fill", "none")
            polygon.setAttribute("data-object-type", shape["object"] or "")
            polygon.classList.add(f"shape-{shape['type']}")
            polygon.classList.add("shape")
            new_shape.appendChild(polygon)
        elif shape["type"] == "point":
            point = document.createElementNS("http://www.w3.org/2000/svg", "circle")
            point.setAttribute("cx", str(shape["shape"]["x"] * image.naturalWidth))
            point.setAttribute("cy", str(shape["shape"]["y"] * image.naturalHeight))
            point.setAttribute("r", "0")
            point.classList.add("shape-point")
            point.classList.add("shape")
            point.setAttribute("data-object-type", shape["object"] or "")
            new_shape.appendChild(point)

        zone.appendChild(new_shape)


async def load_shapes():
    resource_id = document.getElementById("resource-id").value
    response = await pyfetch(f"/picture/{resource_id}/get-annotations")
    if response.ok:
        shapes = await response.json()
        return shapes


async def save_shapes(event):
    shapes = list_shapes()
    resource_id = document.getElementById("resource-id").value
    print("Saving shapes:", shapes)
    response = await pyfetch(f"/picture/{resource_id}/save-annotations",
        method="POST",
        headers={
            "Content-Type": "application/json"
        },
        body=json.dumps(shapes)
    )
    if response.ok:
        return await response


save_shapes_proxy = create_proxy(save_shapes)
save_button.addEventListener("click", save_shapes_proxy)
delete_shape_key_proxy = create_proxy(lambda event: event.key == "Delete" and delete_shape_proxy(None))
next_shape_key_proxy = create_proxy(lambda event: event.key == "ArrowRight" and next_shape_proxy(None))
previous_shape_key_proxy = create_proxy(lambda event: event.key == "ArrowLeft" and previous_shape_proxy(None))


async def focus_shape(shape):
    global selected_shape

    if shape_type != "select":
        return
    if selected_shape is not None:
        selected_shape.classList.remove("selected")

    selected_shape = shape

    selected_shape.classList.add("selected")

    objects = await get_all_objects()

    delete_button.style.display = "block"
    next_button.style.display = "block"
    previous_button.style.display = "block"
    document.addEventListener("keydown", delete_shape_key_proxy)
    document.addEventListener("keydown", next_shape_key_proxy)
    document.addEventListener("keydown", previous_shape_key_proxy)

    object_list.innerHTML = ""

    new_radio = document.createElement("input")
    new_radio.setAttribute("type", "radio")
    new_radio.setAttribute("name", "object-type")
    new_radio.setAttribute("value", "")
    new_label = document.createElement("label")
    new_label.appendChild(new_radio)
    new_label.append("Undefined")
    object_list.appendChild(new_label)
    new_radio.addEventListener("change", change_object_type_proxy)

    selected_object = selected_shape.getAttribute("data-object-type")
    if not selected_object:
        new_radio.setAttribute("checked", "")

    for object, description in objects.items():
        new_radio = document.createElement("input")
        new_radio.setAttribute("type", "radio")
        new_radio.setAttribute("name", "object-type")
        new_radio.setAttribute("value", object)
        if selected_object == object:
            new_radio.setAttribute("checked", "")
        new_label = document.createElement("label")
        new_label.appendChild(new_radio)
        new_label.append(object)
        object_list.appendChild(new_label)
        new_radio.addEventListener("change", change_object_type_proxy)


async def select_shape(event):
    await focus_shape(event.target)


async def next_shape(event):
    global selected_shape
    if selected_shape is None:
        return

    selected_svg = selected_shape.parentNode

    while selected_svg is not None:
        next_sibling = selected_svg.nextElementSibling
        if next_sibling and next_sibling.classList.contains("shape-container"):
            selected_svg = next_sibling
            break
        elif next_sibling is None:
            # If no more siblings, loop back to the first child
            selected_svg = selected_svg.parentNode.firstElementChild
            while selected_svg is not None and not selected_svg.classList.contains(
                    "shape-container"):
                selected_svg = selected_svg.nextElementSibling
            break
        else:
            selected_svg = next_sibling

    if selected_svg:
        shape = selected_svg.firstElementChild
        await focus_shape(shape)


async def previous_shape(event):
    global selected_shape
    if selected_shape is None:
        return

    selected_svg = selected_shape.parentNode

    while selected_svg is not None:
        next_sibling = selected_svg.previousElementSibling
        if next_sibling and next_sibling.classList.contains("shape-container"):
            selected_svg = next_sibling
            break
        elif next_sibling is None:
            # If no more siblings, loop back to the last child
            selected_svg = selected_svg.parentNode.lastElementChild
            while selected_svg is not None and not selected_svg.classList.contains(
                    "shape-container"):
                selected_svg = selected_svg.previousElementSibling
            break
        else:
            selected_svg = next_sibling

    if selected_svg:
        shape = selected_svg.firstElementChild
        await focus_shape(shape)


def unselect_shape(event):
    global selected_shape

    if selected_shape is not None:
        selected_shape.classList.remove("selected")
        selected_shape = None

    object_list.innerHTML = ""
    delete_button.style.display = "none"
    next_button.style.display = "none"
    previous_button.style.display = "none"
    document.removeEventListener("keydown", delete_shape_key_proxy)
    document.removeEventListener("keydown", next_shape_key_proxy)
    document.removeEventListener("keydown", previous_shape_key_proxy)


def delete_shape(event):
    global selected_shape
    if selected_shape is None:
        return
    # Shape is SVG shape inside SVG so we need to remove the parent SVG
    selected_shape.parentNode.remove()
    selected_shape = None
    object_list.innerHTML = ""
    delete_button.style.display = "none"
    next_button.style.display = "none"
    previous_button.style.display = "none"
    document.removeEventListener("keydown", delete_shape_key_proxy)
    document.removeEventListener("keydown", next_shape_key_proxy)
    document.removeEventListener("keydown", previous_shape_key_proxy)


select_shape_proxy = create_proxy(select_shape)
unselect_shape_proxy = create_proxy(unselect_shape)
delete_shape_proxy = create_proxy(delete_shape)
next_shape_proxy = create_proxy(next_shape)
previous_shape_proxy = create_proxy(previous_shape)

delete_button.addEventListener("click", delete_shape_proxy)
next_button.addEventListener("click", next_shape_proxy)
previous_button.addEventListener("click", previous_shape_proxy)

# These are functions usable in JS
cancel_bbox_proxy = create_proxy(lambda event: cancel_bbox(event))
make_bbox_proxy = create_proxy(lambda event: make_bbox(event))
make_polygon_proxy = create_proxy(lambda event: make_polygon(event))
follow_cursor_proxy = create_proxy(follow_cursor)


def switch_shape(event):
    global shape_type
    object_list.innerHTML = ""
    unselect_shape(None)
    shape = event.currentTarget.id
    shape_type = shape
    if shape_type == "select":
        # Add event listeners to existing shapes
        print(len(list(document.getElementsByClassName("shape"))), "shapes found")
        for shape in document.getElementsByClassName("shape"):
            print("Adding event listener to shape:", shape)
            shape.addEventListener("click", select_shape_proxy)
        image.addEventListener("click", unselect_shape_proxy)
        helper_message.innerText = "Click on a shape to select"
        # Cancel the current shape creation
        if shape_type == "shape-bbox":
            cancel_bbox(None)
        elif shape_type == "shape-polygon":
            cancel_polygon(None)
        elif shape_type == "shape-polyline":
            cancel_polygon(None)
    else:
        # Remove event listeners for selection
        for shape in document.getElementsByClassName("shape"):
            print("Removing event listener from shape:", shape)
            shape.removeEventListener("click", select_shape_proxy)
        image.removeEventListener("click", unselect_shape_proxy)
        helper_message.innerText = "Select a shape type then click on the image to begin defining it"
    print("Shape is now of type:", shape)


vertical_ruler = document.getElementById("annotation-ruler-vertical")
horizontal_ruler = document.getElementById("annotation-ruler-horizontal")
vertical_ruler_2 = document.getElementById("annotation-ruler-vertical-secondary")
horizontal_ruler_2 = document.getElementById("annotation-ruler-horizontal-secondary")
helper_message = document.getElementById("annotation-helper-message")

helper_message.innerText = "Select a shape type then click on the image to begin defining it"


def cancel_bbox(event):
    global bbox_pos, new_shape

    # Key must be ESCAPE
    if event is not None and hasattr(event, "key") and event.key != "Escape":
        return

    if new_shape is not None and event is not None:
        # Require event so the shape is kept when it ends normally
        new_shape.remove()
    zone.removeEventListener("click", make_bbox_proxy)
    document.removeEventListener("keydown", cancel_bbox_proxy)
    cancel_button.removeEventListener("click", cancel_bbox_proxy)

    bbox_pos = None
    vertical_ruler.style.display = "none"
    horizontal_ruler.style.display = "none"
    vertical_ruler_2.style.display = "none"
    horizontal_ruler_2.style.display = "none"
    zone.style.cursor = "auto"
    cancel_button.style.display = "none"
    helper_message.innerText = "Select a shape type then click on the image to begin defining it"
    new_shape = None


def make_bbox(event):
    global new_shape, bbox_pos
    zone_rect = zone.getBoundingClientRect()

    if bbox_pos is None:
        helper_message.innerText = "Now define the second point"

        bbox_pos = [(event.clientX - zone_rect.left) / zone_rect.width,
                    (event.clientY - zone_rect.top) / zone_rect.height]
        vertical_ruler_2.style.left = str(bbox_pos[0] * 100) + "%"
        horizontal_ruler_2.style.top = str(bbox_pos[1] * 100) + "%"
        vertical_ruler_2.style.display = "block"
        horizontal_ruler_2.style.display = "block"

    else:
        x0, y0 = bbox_pos.copy()
        x1 = (event.clientX - zone_rect.left) / zone_rect.width
        y1 = (event.clientY - zone_rect.top) / zone_rect.height

        rectangle = document.createElementNS("http://www.w3.org/2000/svg", "rect")

        new_shape = make_shape_container()
        zone_rect = zone.getBoundingClientRect()

        new_shape.appendChild(rectangle)
        zone.appendChild(new_shape)

        minx = min(x0, x1)
        miny = min(y0, y1)
        maxx = max(x0, x1)
        maxy = max(y0, y1)

        rectangle.setAttribute("x", str(minx * image.naturalWidth))
        rectangle.setAttribute("y", str(miny * image.naturalHeight))
        rectangle.setAttribute("width", str((maxx - minx) * image.naturalWidth))
        rectangle.setAttribute("height", str((maxy - miny) * image.naturalHeight))
        rectangle.setAttribute("fill", "none")
        rectangle.setAttribute("data-object-type", "")
        rectangle.classList.add("shape-bbox")
        rectangle.classList.add("shape")

        # Add event listeners to the new shape
        rectangle.addEventListener("click", select_shape_proxy)

        cancel_bbox(None)


polygon_points = []


def make_polygon(event):
    global new_shape, polygon_points

    polygon = new_shape.children[0]

    zone_rect = zone.getBoundingClientRect()

    polygon_points.append(((event.clientX - zone_rect.left) / zone_rect.width,
                           (event.clientY - zone_rect.top) / zone_rect.height))

    # Update the polygon
    polygon.setAttribute("points", " ".join(
            [f"{point[0] * image.naturalWidth},{point[1] * image.naturalHeight}" for point in
             polygon_points]))


def reset_polygon():
    global new_shape, polygon_points

    zone.removeEventListener("click", make_polygon_proxy)
    document.removeEventListener("keydown", close_polygon_proxy)
    document.removeEventListener("keydown", cancel_polygon_proxy)
    document.removeEventListener("keydown", backspace_polygon_proxy)
    confirm_button.style.display = "none"
    cancel_button.style.display = "none"
    backspace_button.style.display = "none"
    confirm_button.removeEventListener("click", close_polygon_proxy)
    cancel_button.removeEventListener("click", cancel_polygon_proxy)
    backspace_button.removeEventListener("click", backspace_polygon_proxy)
    polygon_points.clear()

    zone.style.cursor = "auto"
    new_shape = None


def close_polygon(event):
    if event is not None and hasattr(event, "key") and event.key != "Enter":
        return
    # Polygon is already there, but we need to remove the events
    reset_polygon()


def cancel_polygon(event):
    if event is not None and hasattr(event, "key") and event.key != "Escape":
        return
    # Delete the polygon
    new_shape.remove()
    reset_polygon()


def backspace_polygon(event):
    if event is not None and hasattr(event, "key") and event.key != "Backspace":
        return
    if not polygon_points:
        return
    polygon_points.pop()
    polygon = new_shape.children[0]
    polygon.setAttribute("points", " ".join(
            [f"{point[0] * image.naturalWidth},{point[1] * image.naturalHeight}" for point in
             polygon_points]))


close_polygon_proxy = create_proxy(close_polygon)
cancel_polygon_proxy = create_proxy(cancel_polygon)
backspace_polygon_proxy = create_proxy(backspace_polygon)


def open_shape(event):
    global new_shape, bbox_pos
    if bbox_pos or shape_type == "select":
        return
    print("Creating a new shape of type:", shape_type)

    if shape_type == "shape-bbox":
        helper_message.innerText = ("Define the first point at the intersection of the lines "
                                    "by clicking on the image, or click the cross to cancel")

        cancel_button.addEventListener("click", cancel_bbox_proxy)
        document.addEventListener("keydown", cancel_bbox_proxy)
        cancel_button.style.display = "block"
        bbox_pos = None
        zone.addEventListener("click", make_bbox_proxy)
        vertical_ruler.style.display = "block"
        horizontal_ruler.style.display = "block"
        zone.style.cursor = "crosshair"
    elif shape_type == "shape-polygon" or shape_type == "shape-polyline":
        if shape_type == "shape-polygon":
            helper_message.innerText = ("Click on the image to define the points of the polygon, "
                                        "press escape to cancel, enter to close, or backspace to "
                                        "remove the last point")
        elif shape_type == "shape-polyline":
            helper_message.innerText = ("Click on the image to define the points of the polyline, "
                                        "press escape to cancel, enter to finish, or backspace to "
                                        "remove the last point")

        if not polygon_points and not new_shape:
            new_shape = make_shape_container()
            zone_rect = zone.getBoundingClientRect()

        if not polygon_points and int(new_shape.children.length) == 0:
            zone.addEventListener("click", make_polygon_proxy)
            document.addEventListener("keydown", close_polygon_proxy)
            document.addEventListener("keydown", cancel_polygon_proxy)
            document.addEventListener("keydown", backspace_polygon_proxy)
            cancel_button.addEventListener("click", cancel_polygon_proxy)
            cancel_button.style.display = "block"
            confirm_button.addEventListener("click", close_polygon_proxy)
            confirm_button.style.display = "block"
            backspace_button.addEventListener("click", backspace_polygon_proxy)
            backspace_button.style.display = "block"
            if shape_type == "shape-polygon":
                polygon = document.createElementNS("http://www.w3.org/2000/svg", "polygon")
                polygon.classList.add("shape-polygon")
            elif shape_type == "shape-polyline":
                polygon = document.createElementNS("http://www.w3.org/2000/svg", "polyline")
                polygon.classList.add("shape-polyline")
            polygon.setAttribute("fill", "none")
            polygon.setAttribute("data-object-type", "")
            polygon.classList.add("shape")
            new_shape.appendChild(polygon)
            zone.appendChild(new_shape)
            zone.style.cursor = "crosshair"
    elif shape_type == "shape-point":
        point = document.createElementNS("http://www.w3.org/2000/svg", "circle")
        zone_rect = zone.getBoundingClientRect()
        point.setAttribute("cx", str((event.clientX - zone_rect.left) / zone_rect.width * image.naturalWidth))
        point.setAttribute("cy", str((event.clientY - zone_rect.top) / zone_rect.height * image.naturalHeight))
        point.setAttribute("r", "0")
        point.classList.add("shape-point")
        point.classList.add("shape")
        point.setAttribute("data-object-type", "")

        new_shape = make_shape_container()
        zone_rect = zone.getBoundingClientRect()

        new_shape.appendChild(point)
        zone.appendChild(new_shape)

        new_shape = None



for button in list(document.getElementById("shape-selector").children):
    button.addEventListener("click", create_proxy(switch_shape))
    print("Shape", button.id, "is available")

zone.addEventListener("mousemove", follow_cursor_proxy)
zone.addEventListener("click", create_proxy(open_shape))

# Load existing annotations, if any
put_shapes(await load_shapes())
print("Ready!")