// Re:Form Re:Link receiver for Unity (2020.3 and later).
//
// Install: put this file in your project at Assets/Editor/ReformRelinkReceiver.cs. There is nothing to switch on: 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:
// the file is saved in Assets/ReformRelink/ and placed in the open scene as an object named after the layer. Sending the same
// layer again updates that object where it stands. Window > Re:Form Re:Link shows whether it is listening.
//
// 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": "..."}
#if UNITY_EDITOR
using System;
using System.Collections.Concurrent;
using System.IO;
using System.Net;
using System.Text;
using System.Text.RegularExpressions;
using System.Threading;
using UnityEditor;
using UnityEditor.SceneManagement;
using UnityEngine;
using UnityEngine.SceneManagement;

[InitializeOnLoad]
public static class ReformRelinkReceiver
{
    const int ProtocolVersion = 1;
    const int Port = 47800;
    const string ModelFolder = "Assets/ReformRelink";
    const long LargestModelBytes = 2147483647L;
    const int SecondsToWaitForImport = 300;
    static readonly string[] Formats = { "fbx", "obj" };
    static readonly Regex AllowedOrigin = new Regex(@"^(https://(www\.)?reformsculpt\.com|http://(localhost|127\.0\.0\.1)(:\d+)?)$");

    class WaitingModel
    {
        public string Name;
        public string Format;
        public byte[] Bytes;
        public readonly ManualResetEventSlim Done = new ManualResetEventSlim(false);
        public bool Received;
        public string Note = "";
    }

    static readonly ConcurrentQueue<WaitingModel> waitingModels = new ConcurrentQueue<WaitingModel>();
    static HttpListener listener;
    static Thread listenerThread;
    static string projectName = "Unity project";
    public static string Problem = "";
    public static bool IsListening { get { return listener != null && listener.IsListening; } }

    static ReformRelinkReceiver()
    {
        projectName = Application.productName;
        StartListening();
        EditorApplication.update += TakeWaitingModels;
        AssemblyReloadEvents.beforeAssemblyReload += StopListening;
        EditorApplication.quitting += StopListening;
    }

    [MenuItem("Window/Re:Form Re:Link")]
    static void ShowStatus()
    {
        EditorUtility.DisplayDialog("Re:Form Re:Link", IsListening ? "Listening on this computer only, port " + Port + ". In Re:Form: File > Re:Link." : (Problem.Length > 0 ? Problem : "Not listening."), "OK");
    }

    static void StartListening()
    {
        if (listener != null) return;
        try
        {
            listener = new HttpListener();
            listener.Prefixes.Add("http://127.0.0.1:" + Port + "/");
            listener.Start();
        }
        catch (Exception listenError)
        {
            Problem = "Port " + Port + " is in use (another Re:Link receiver may be running): " + listenError.Message;
            Debug.LogWarning("Re:Link receiver: " + Problem);
            listener = null;
            return;
        }
        Problem = "";
        listenerThread = new Thread(ListenLoop) { IsBackground = true, Name = "Re:Link receiver" };
        listenerThread.Start();
    }

    static void StopListening()
    {
        HttpListener stopping = listener;
        listener = null;
        if (stopping != null) { try { stopping.Stop(); stopping.Close(); } catch (Exception) { } }
    }

    static void ListenLoop()
    {
        HttpListener listening = listener;
        while (listening != null && listening.IsListening)
        {
            HttpListenerContext context;
            try { context = listening.GetContext(); }
            catch (Exception) { return; }   // the listener was stopped
            ThreadPool.QueueUserWorkItem(delegate { try { Answer(context); } catch (Exception) { try { context.Response.Abort(); } catch (Exception) { } } });
        }
    }

    static bool OriginIsAllowed(string origin) { return string.IsNullOrEmpty(origin) || AllowedOrigin.IsMatch(origin); }

    static string SafeModelName(string name)
    {
        string cleaned = Regex.Replace(name ?? "", "[^A-Za-z0-9 _-]", "_").Trim();
        if (cleaned.Length > 48) cleaned = cleaned.Substring(0, 48).Trim();
        return cleaned.Length == 0 || Regex.IsMatch(cleaned, "^[_ -]+$") ? "Sculpt" : cleaned;
    }

    static string JsonText(string text)
    {
        StringBuilder builder = new StringBuilder("\"");
        foreach (char character in text ?? "")
        {
            if (character == '"' || character == '\\') builder.Append('\\').Append(character);
            else if (character < ' ') builder.Append(' ');
            else builder.Append(character);
        }
        return builder.Append('"').ToString();
    }

    static void Send(HttpListenerContext context, int statusCode, string json)
    {
        string origin = context.Request.Headers["Origin"];
        HttpListenerResponse response = context.Response;
        response.StatusCode = statusCode;
        if (!string.IsNullOrEmpty(origin) && statusCode != 403)
        {
            response.AddHeader("Access-Control-Allow-Origin", origin);
            response.AddHeader("Vary", "Origin");
        }
        response.AddHeader("Access-Control-Allow-Methods", "GET, POST, OPTIONS");
        response.AddHeader("Access-Control-Allow-Headers", "Content-Type");
        response.AddHeader("Access-Control-Allow-Private-Network", "true");
        response.AddHeader("Access-Control-Max-Age", "600");
        byte[] payload = Encoding.UTF8.GetBytes(json);
        response.ContentType = "application/json";
        response.ContentLength64 = payload.Length;
        response.OutputStream.Write(payload, 0, payload.Length);
        response.OutputStream.Close();
    }

    static string Receipt(bool received, string name, string note)
    {
        return "{\"reformRelink\": " + ProtocolVersion + ", \"received\": " + (received ? "true" : "false") + ", \"name\": " + JsonText(name) + ", \"note\": " + JsonText(note) + "}";
    }

    static void Answer(HttpListenerContext context)
    {
        HttpListenerRequest request = context.Request;
        if (!OriginIsAllowed(request.Headers["Origin"])) { Send(context, 403, "{}"); return; }
        string path = request.Url.AbsolutePath;
        if (request.HttpMethod == "OPTIONS") { Send(context, 200, "{\"reformRelink\": " + ProtocolVersion + "}"); return; }
        if (request.HttpMethod == "GET" && path == "/reform-relink/status")
        {
            Send(context, 200, "{\"reformRelink\": " + ProtocolVersion + ", \"scene\": " + JsonText(projectName) + ", \"formats\": [\"fbx\", \"obj\"]}");
            return;
        }
        if (request.HttpMethod == "POST" && path == "/reform-relink/model")
        {
            string modelName = SafeModelName(request.QueryString["name"]);
            string format = (request.QueryString["format"] ?? "").ToLowerInvariant();
            long byteCount = request.ContentLength64;
            if (Array.IndexOf(Formats, format) < 0 || byteCount <= 0 || byteCount > LargestModelBytes)
            {
                Send(context, 200, Receipt(false, modelName, "this receiver takes FBX or OBJ files"));
                return;
            }
            byte[] bytes = new byte[byteCount];
            int read = 0;
            while (read < bytes.Length)
            {
                int count = request.InputStream.Read(bytes, read, bytes.Length - read);
                if (count <= 0) break;
                read += count;
            }
            if (read < bytes.Length) { Send(context, 200, Receipt(false, modelName, "the file did not arrive whole")); return; }
            WaitingModel waiting = new WaitingModel { Name = modelName, Format = format, Bytes = bytes };
            waitingModels.Enqueue(waiting);
            // the import runs on the editor's own thread (TakeWaitingModels); this thread waits for it so the answer is the truth
            if (!waiting.Done.Wait(TimeSpan.FromSeconds(SecondsToWaitForImport))) { Send(context, 200, Receipt(false, modelName, "the editor was busy and did not take the model in time")); return; }
            Send(context, 200, Receipt(waiting.Received, modelName, waiting.Note));
            return;
        }
        Send(context, 404, Receipt(false, "", "no such path"));
    }

    // The editor's own thread: imports whatever has arrived since the last look.
    static void TakeWaitingModels()
    {
        WaitingModel waiting;
        while (waitingModels.TryDequeue(out waiting))
        {
            try { waiting.Note = PlaceModel(waiting); waiting.Received = true; }
            catch (Exception importError) { waiting.Note = "Unity could not import the file: " + importError.Message; waiting.Received = false; }
            waiting.Done.Set();
        }
    }

    static string PlaceModel(WaitingModel waiting)
    {
        if (!AssetDatabase.IsValidFolder(ModelFolder)) AssetDatabase.CreateFolder("Assets", "ReformRelink");
        string assetPath = ModelFolder + "/" + waiting.Name + "." + waiting.Format;
        File.WriteAllBytes(assetPath, waiting.Bytes);
        AssetDatabase.ImportAsset(assetPath, ImportAssetOptions.ForceUpdate | ImportAssetOptions.ForceSynchronousImport);
        GameObject modelAsset = AssetDatabase.LoadAssetAtPath<GameObject>(assetPath);
        if (modelAsset == null)
        {
            AssetDatabase.DeleteAsset(assetPath);
            throw new Exception("it holds no model");
        }
        int triangleCount = 0;
        foreach (MeshFilter meshFilter in modelAsset.GetComponentsInChildren<MeshFilter>(true))
        {
            Mesh mesh = meshFilter.sharedMesh;
            if (mesh == null) continue;
            for (int subMeshIndex = 0; subMeshIndex < mesh.subMeshCount; subMeshIndex++) triangleCount += (int)(mesh.GetIndexCount(subMeshIndex) / 3);
        }
        if (triangleCount == 0)
        {
            AssetDatabase.DeleteAsset(assetPath);
            throw new Exception("it holds no mesh");
        }
        string triangleText = " (" + triangleCount.ToString("N0", System.Globalization.CultureInfo.InvariantCulture) + " triangles)";
        Scene scene = SceneManager.GetActiveScene();
        if (!scene.IsValid()) return "Saved in the project as " + assetPath + triangleText + ". No scene is open: open one and send again to place it.";
        // an object placed by an earlier send is an instance of this same file: the editor has just updated it where it stands
        foreach (GameObject rootObject in scene.GetRootGameObjects())
        {
            foreach (Transform child in rootObject.GetComponentsInChildren<Transform>(true))
            {
                if (PrefabUtility.GetCorrespondingObjectFromSource(child.gameObject) == modelAsset) return "Replaced the one sent before" + triangleText + ".";
            }
        }
        GameObject placed = PrefabUtility.InstantiatePrefab(modelAsset, scene) as GameObject;
        if (placed == null) throw new Exception("it could not be placed in the scene");
        placed.name = waiting.Name;
        EditorSceneManager.MarkSceneDirty(scene);
        return "Added to the scene" + triangleText + ".";
    }
}
#endif