from pyscript import document
from pyodide.ffi import create_proxy
from pyodide.http import pyfetch
import asyncio

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")

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

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


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)


async def select_shape(event):
    global selected_shape
    if shape_type != "select":
        return
    if selected_shape is not None:
        selected_shape.classList.remove("selected")

    print("Selected shape:", event.target)

    selected_shape = event.target
    selected_shape.classList.add("selected")

    objects = await get_all_objects()

    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)

    print(objects)


def unselect_shape(event):
    global selected_shape

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

    object_list.innerHTML = ""


select_shape_proxy = create_proxy(select_shape)
unselect_shape_proxy = create_proxy(unselect_shape)


# 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"
    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()
        new_shape = None
    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"
    document.body.style.cursor = "auto"
    cancel_button.style.display = "none"
    helper_message.innerText = "Select a shape type then click on the image to begin defining it"

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.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

    print(new_shape.outerHTML)
    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 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
    zone.removeEventListener("click", make_polygon_proxy)
    document.removeEventListener("keydown", close_polygon_proxy)
    document.removeEventListener("keydown", cancel_polygon_proxy)
    confirm_button.style.display = "none"
    cancel_button.style.display = "none"
    confirm_button.removeEventListener("click", close_polygon_proxy)
    cancel_button.removeEventListener("click", cancel_polygon_proxy)
    polygon_points.clear()

    document.body.style.cursor = "auto"


def cancel_polygon(event):
    if event is not None and hasattr(event, "key") and event.key != "Escape":
        return
    # Delete the polygon
    new_shape.remove()
    zone.removeEventListener("click", make_polygon_proxy)
    document.removeEventListener("keydown", close_polygon_proxy)
    document.removeEventListener("keydown", cancel_polygon_proxy)
    confirm_button.style.display = "none"
    cancel_button.style.display = "none"
    confirm_button.removeEventListener("click", close_polygon_proxy)
    cancel_button.removeEventListener("click", cancel_polygon_proxy)

    polygon_points.clear()

    document.body.style.cursor = "auto"


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 not polygon_points:
        new_shape = document.createElementNS("http://www.w3.org/2000/svg", "svg")
        new_shape.setAttribute("width", "100%")
        new_shape.setAttribute("height", "100%")
        zone_rect = zone.getBoundingClientRect()
        new_shape.setAttribute("viewBox", f"0 0 {image.naturalWidth} {image.naturalHeight}")
        new_shape.classList.add("shape-container")

    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"
        document.body.style.cursor = "crosshair"
    elif 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")

        if not polygon_points:
            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"
            polygon = document.createElementNS("http://www.w3.org/2000/svg", "polygon")
            polygon.setAttribute("fill", "none")
            polygon.setAttribute("data-object-type", "")
            polygon.classList.add("shape-polygon")
            polygon.classList.add("shape")
            new_shape.appendChild(polygon)
            zone.appendChild(new_shape)
            document.body.style.cursor = "crosshair"


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))