# Re:Form Re:Link receiver for Unreal Engine (4.26 and later, 5.x).
#
# Install: 1. Edit > Plugins: switch on "Python Editor Script Plugin" and "Editor Scripting Utilities", and restart the editor.
#          2. Copy this file into your project at Content/Python/reform_relink_unreal.py.
#          3. Edit > Project Settings > Plugins > Python > Startup Scripts: add  reform_relink_unreal.py  and restart the editor.
# While the editor is open it listens on this computer only (the loopback address, port 47800). In Re:Form, File > Re:Link sends
# the active layer: it is imported as a static mesh in /Game/ReformRelink and placed in the open level as an actor named after the
# layer. Sending the same layer again updates the mesh, and the actor stays where you put it.
#
# Nothing here reaches the internet, and only pages of Re:Form are answered.
#
# The conversation (version 1):
#   GET  /reform-relink/status                        -> {"reformRelink": 1, "scene": "...", "formats": ["fbx", "obj"]}
#   POST /reform-relink/model?name=...&format=...     -> {"reformRelink": 1, "received": true, "name": "...", "note": "..."}

import json
import os
import queue
import re
import tempfile
import threading
import urllib.parse
from http.server import BaseHTTPRequestHandler, HTTPServer
from socketserver import ThreadingMixIn

import unreal

RELINK_PROTOCOL_VERSION = 1
RELINK_PORT = 47800
RELINK_FORMATS = ["fbx", "obj"]
RELINK_CONTENT_FOLDER = "/Game/ReformRelink"
RELINK_LARGEST_MODEL_BYTES = 2 * 1024 * 1024 * 1024
RELINK_SECONDS_TO_WAIT_FOR_IMPORT = 600.0
ALLOWED_ORIGIN_PATTERN = re.compile(r"^(https://(www\.)?reformsculpt\.com|http://(localhost|127\.0\.0\.1)(:\d+)?)$")

waiting_models = queue.Queue()
receiver_state = {"server": None, "thread": None, "tick_handle": None, "problem": ""}


class ThreadedServer(ThreadingMixIn, HTTPServer):
    daemon_threads = True


def origin_is_allowed(origin):
    return (not origin) or bool(ALLOWED_ORIGIN_PATTERN.match(origin))


def safe_model_name(name):
    cleaned = re.sub(r"[^A-Za-z0-9 _-]", "_", name or "").strip()[:48].strip()
    return cleaned if cleaned and not re.match(r"^[_ -]+$", cleaned) else "Sculpt"


def asset_name_of(model_name):
    """The name the mesh has in the Content Browser: letters, digits and underscores only."""
    return re.sub(r"[^A-Za-z0-9_]", "_", model_name)


def project_name():
    try:
        return unreal.Paths.get_base_filename(unreal.Paths.get_project_file_path()) or "Unreal project"
    except Exception:
        return "Unreal project"


class RelinkRequestHandler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, format_text, *arguments):
        pass

    def send_answer(self, status_code, body):
        payload = json.dumps(body).encode("utf-8")
        self.send_response(status_code)
        origin = self.headers.get("Origin", "")
        if origin:
            self.send_header("Access-Control-Allow-Origin", origin)
            self.send_header("Vary", "Origin")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.send_header("Access-Control-Allow-Private-Network", "true")
        self.send_header("Access-Control-Max-Age", "600")
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def refuse_other_pages(self):
        if origin_is_allowed(self.headers.get("Origin", "")):
            return False
        self.send_response(403)
        self.send_header("Content-Length", "0")
        self.end_headers()
        return True

    def do_OPTIONS(self):
        if self.refuse_other_pages():
            return
        self.send_answer(200, {"reformRelink": RELINK_PROTOCOL_VERSION})

    def do_GET(self):
        if self.refuse_other_pages():
            return
        address = urllib.parse.urlparse(self.path)
        if address.path != "/reform-relink/status":
            self.send_answer(404, {"reformRelink": RELINK_PROTOCOL_VERSION, "received": False, "note": "no such path"})
            return
        self.send_answer(200, {"reformRelink": RELINK_PROTOCOL_VERSION, "scene": receiver_state.get("project_name", "Unreal project"), "formats": RELINK_FORMATS})

    def do_POST(self):
        if self.refuse_other_pages():
            return
        address = urllib.parse.urlparse(self.path)
        if address.path != "/reform-relink/model":
            self.send_answer(404, {"reformRelink": RELINK_PROTOCOL_VERSION, "received": False, "note": "no such path"})
            return
        query = urllib.parse.parse_qs(address.query)
        model_name = safe_model_name(query.get("name", [""])[0])
        format_name = query.get("format", [""])[0].lower()
        byte_count = int(self.headers.get("Content-Length", "0") or "0")
        if format_name not in RELINK_FORMATS or byte_count <= 0 or byte_count > RELINK_LARGEST_MODEL_BYTES:
            self.send_answer(200, {"reformRelink": RELINK_PROTOCOL_VERSION, "received": False, "name": model_name, "note": "this receiver takes FBX or OBJ files"})
            return
        file_bytes = self.rfile.read(byte_count)
        waiting = {"name": model_name, "format": format_name, "bytes": file_bytes, "done": threading.Event(), "note": "", "received": False}
        waiting_models.put(waiting)
        # the import runs on the editor's own thread (the tick below); this thread waits for it so the answer is the truth
        if not waiting["done"].wait(RELINK_SECONDS_TO_WAIT_FOR_IMPORT):
            self.send_answer(200, {"reformRelink": RELINK_PROTOCOL_VERSION, "received": False, "name": model_name, "note": "the editor was busy and did not take the model in time"})
            return
        self.send_answer(200, {"reformRelink": RELINK_PROTOCOL_VERSION, "received": waiting["received"], "name": model_name, "note": waiting["note"]})


def triangle_count_of(static_mesh):
    try:
        return int(static_mesh.get_num_triangles(0))
    except Exception:
        pass
    try:
        total = 0
        for section_index in range(unreal.EditorStaticMeshLibrary.get_number_materials(static_mesh)):
            total += len(unreal.ProceduralMeshLibrary.get_section_from_static_mesh(static_mesh, 0, section_index)[1]) // 3
        return total
    except Exception:
        return 0


def place_model(waiting):
    """Imports one waiting model as a static mesh and places it in the open level. Runs on the editor's own thread."""
    model_name = waiting["name"]
    asset_name = asset_name_of(model_name)
    temporary_directory = tempfile.mkdtemp(prefix="reform_relink_")
    file_path = os.path.join(temporary_directory, asset_name + "." + waiting["format"])
    with open(file_path, "wb") as model_file:
        model_file.write(waiting["bytes"])
    try:
        options = unreal.FbxImportUI()
        options.set_editor_property("import_mesh", True)
        options.set_editor_property("import_as_skeletal", False)
        options.set_editor_property("import_materials", False)
        options.set_editor_property("import_textures", False)
        options.set_editor_property("import_animations", False)
        options.set_editor_property("mesh_type_to_import", unreal.FBXImportType.FBXIT_STATIC_MESH)
        options.static_mesh_import_data.set_editor_property("combine_meshes", True)
        options.static_mesh_import_data.set_editor_property("remove_degenerates", False)  # every triangle of the sculpt arrives
        task = unreal.AssetImportTask()
        task.set_editor_property("filename", file_path)
        task.set_editor_property("destination_path", RELINK_CONTENT_FOLDER)
        task.set_editor_property("destination_name", asset_name)
        task.set_editor_property("automated", True)
        task.set_editor_property("replace_existing", True)
        task.set_editor_property("save", False)
        task.set_editor_property("options", options)
        unreal.AssetToolsHelpers.get_asset_tools().import_asset_tasks([task])
    finally:
        try:
            os.remove(file_path)
            os.rmdir(temporary_directory)
        except OSError:
            pass
    static_mesh = unreal.EditorAssetLibrary.load_asset(RELINK_CONTENT_FOLDER + "/" + asset_name) if unreal.EditorAssetLibrary.does_asset_exist(RELINK_CONTENT_FOLDER + "/" + asset_name) else None
    if static_mesh is None or not isinstance(static_mesh, unreal.StaticMesh):
        raise RuntimeError("the file held no mesh")
    triangle_count = triangle_count_of(static_mesh)
    triangle_text = (" (" + format(triangle_count, ",") + " triangles)") if triangle_count > 0 else ""
    try:
        level_actors = unreal.EditorLevelLibrary.get_all_level_actors()
    except Exception:
        level_actors = None
    if level_actors is None:
        return "Saved in the project as " + RELINK_CONTENT_FOLDER + "/" + asset_name + triangle_text + ". No level is open: open one and send again to place it."
    for actor in level_actors:
        if isinstance(actor, unreal.StaticMeshActor) and actor.static_mesh_component is not None and actor.static_mesh_component.static_mesh == static_mesh:
            return "Replaced the one sent before" + triangle_text + "."   # the actor shows the mesh as it is now, where it stands
    actor = unreal.EditorLevelLibrary.spawn_actor_from_object(static_mesh, unreal.Vector(0.0, 0.0, 0.0), unreal.Rotator(0.0, 0.0, 0.0))
    if actor is None:
        return "Saved in the project as " + RELINK_CONTENT_FOLDER + "/" + asset_name + triangle_text + ". No level is open: open one and send again to place it."
    actor.set_actor_label(model_name)
    return "Added to the level" + triangle_text + "."


def take_waiting_models(seconds_since_last_tick=0.0):
    """The editor's tick: imports whatever has arrived since the last look."""
    while True:
        try:
            waiting = waiting_models.get_nowait()
        except queue.Empty:
            break
        try:
            waiting["note"] = place_model(waiting)
            waiting["received"] = True
        except Exception as import_error:
            waiting["note"] = "Unreal could not import the file: " + " ".join(str(import_error).split())[:160]
            waiting["received"] = False
        waiting["done"].set()


def start_receiver(register_tick=True):
    if receiver_state["server"] is not None:
        return
    receiver_state["project_name"] = project_name()
    try:
        server = ThreadedServer(("127.0.0.1", RELINK_PORT), RelinkRequestHandler)
    except OSError as listen_error:
        receiver_state["problem"] = "Port " + str(RELINK_PORT) + " is in use (another Re:Link receiver may be running): " + str(listen_error)
        unreal.log_warning("Re:Link receiver: " + receiver_state["problem"])
        return
    thread = threading.Thread(target=server.serve_forever, name="Re:Link receiver", daemon=True)
    thread.start()
    receiver_state["server"] = server
    receiver_state["thread"] = thread
    receiver_state["problem"] = ""
    if register_tick:
        receiver_state["tick_handle"] = unreal.register_slate_post_tick_callback(take_waiting_models)
    unreal.log("Re:Link receiver: listening on this computer only, port " + str(RELINK_PORT) + ". In Re:Form: File > Re:Link.")


def stop_receiver():
    if receiver_state["tick_handle"] is not None:
        unreal.unregister_slate_post_tick_callback(receiver_state["tick_handle"])
        receiver_state["tick_handle"] = None
    server = receiver_state["server"]
    if server is not None:
        server.shutdown()
        server.server_close()
    receiver_state["server"] = None
    receiver_state["thread"] = None


# One receiver for the editor, however many times the script is run: the one before is stopped first.
earlier_stop = getattr(unreal, "reform_relink_stop_receiver", None)
if earlier_stop is not None:
    try:
        earlier_stop()
    except Exception:
        pass
if os.environ.get("REFORM_RELINK_DO_NOT_START") != "1":
    start_receiver()
    unreal.reform_relink_stop_receiver = stop_receiver