LightWork3D/src/Mod/FreeCADExchange/ExchangeWriteBack.py

# FreeCADExchange write-back helpers.

import json
import os
from datetime import datetime
from pathlib import Path
import traceback
import uuid

import FreeCAD as App

import DeviceImport
import TerminalObjects as TerminalObjects

try:
    import TerminalImport
except ImportError:
    TerminalImport = None

try:
    import FreeCADGui as Gui
except ImportError:
    Gui = None


STATE_WRITEBACK_OBSERVER = "_qet_exchange_writeback_observer"
ENV_JSON_PATH = "QET_2D_TO_3D_JSON"


class ExchangeWriteBackError(RuntimeError):
    pass


def _append_debug_log(message):
    try:
        DeviceImport._append_debug_log(message)
    except Exception:
        pass


def _doc_name(doc):
    if doc is None:
        return "<none>"
    return getattr(doc, "Name", "") or "<unnamed>"


def _doc_path(doc):
    if doc is None:
        return ""
    return getattr(doc, "FileName", "") or ""


def _doc_object_count(doc):
    if doc is None:
        return -1
    try:
        return len(list(getattr(doc, "Objects", []) or []))
    except Exception:
        return -1


def _project_uuid_from_payload(payload):
    if isinstance(payload, dict):
        value = (payload.get("project_uuid") or "").strip()
        if value:
            return value
    return ""


def _root_group(doc):
    try:
        return doc.getObject(TerminalObjects.ROOT_GROUP_NAME)
    except Exception:
        return None


def _is_device_group(obj):
    if obj is None:
        return False
    try:
        if not obj.Name.startswith(DeviceImport.DEVICE_GROUP_PREFIX):
            return False
        return "QetElementUuid" in getattr(obj, "PropertiesList", [])
    except Exception:
        return False


def _iter_device_groups(doc):
    root = _root_group(doc)
    if root is not None:
        for child in list(getattr(root, "Group", []) or []):
            if _is_device_group(child):
                yield child
        return

    for obj in doc.Objects:
        if _is_device_group(obj):
            yield obj


def _iter_terminal_objects(device_group):
    terminal_container = TerminalObjects.find_child_group_by_kind(
        device_group,
        TerminalObjects.TERMINAL_GROUP_KIND,
    )
    if terminal_container is None:
        return []
    return TerminalObjects.collect_terminal_objects(terminal_container)


def _scene_path_from_doc(doc, scene_path=""):
    candidate = (scene_path or "").strip()
    if candidate:
        return candidate

    env_scene = os.environ.get("QET_FREECAD_SCENE_FILE", "").strip()
    if env_scene:
        return env_scene

    file_name = getattr(doc, "FileName", "").strip()
    if file_name:
        return file_name

    return ""


def _output_path_for_scene(scene_path):
    scene_path = (scene_path or "").strip()
    if not scene_path:
        return ""

    path = Path(scene_path)
    if path.suffix.lower() == ".fcstd":
        return str(path.with_name("3d_to_2d.json"))
    if path.is_dir():
        return str(path / "3d_to_2d.json")
    if path.name.lower().endswith(".fcstd"):
        return str(path.with_name("3d_to_2d.json"))
    return str(path.parent / "3d_to_2d.json")


def _output_path_for_exchange_json():
    json_path = os.environ.get(ENV_JSON_PATH, "").strip()
    if not json_path:
        return ""
    return str(Path(json_path).with_name("3d_to_2d.json"))


def _input_path_for_scene(scene_path):
    scene_path = (scene_path or "").strip()
    if not scene_path:
        return ""
    path = Path(scene_path)
    if path.suffix.lower() == ".fcstd":
        return str(path.with_name("2d_to_3d.json"))
    if path.is_dir():
        return str(path / "2d_to_3d.json")
    return str(path.parent / "2d_to_3d.json")


def _load_json_payload(path):
    path_text = (path or "").strip()
    if not path_text:
        return None
    try:
        candidate = Path(path_text)
        if not candidate.is_file():
            return None
        return json.loads(candidate.read_text(encoding="utf-8"))
    except Exception as exc:
        _append_debug_log("write-back could not load payload {0}: {1}".format(path_text, exc))
        return None


def _payload_for_writeback(scene_path, payload=None):
    if isinstance(payload, dict):
        return payload

    env_path = os.environ.get(ENV_JSON_PATH, "").strip()
    loaded = _load_json_payload(env_path)
    if isinstance(loaded, dict):
        return loaded

    loaded = _load_json_payload(_input_path_for_scene(scene_path))
    if isinstance(loaded, dict):
        return loaded

    return payload


def _sync_terminals_for_writeback(doc, scene_path, payload):
    if TerminalImport is None or not isinstance(payload, dict):
        return None
    if not isinstance(payload.get("devices"), list) or not payload.get("devices"):
        return None
    try:
        # 保存/写回以当前 2d_to_3d.json 为端子快照，先同步 3D 工程端子，避免旧工程继续回写缺失或重复端子。
        return TerminalImport.import_terminals_from_payload(payload, scene_path)
    except Exception as exc:
        _append_debug_log("write-back terminal sync failed: {0}".format(exc))
        _append_debug_log(traceback.format_exc())
        return None


def sync_terminals_from_current_payload(doc, scene_path="", payload=None):
    scene_path = _scene_path_from_doc(doc, scene_path)
    payload = _payload_for_writeback(scene_path, payload)
    return _sync_terminals_for_writeback(doc, scene_path, payload)


def _format_timestamp():
    return datetime.now().astimezone().isoformat(timespec="seconds")


def _collect_instance_bindings(doc):
    bindings = []
    seen = set()
    for device_group in _iter_device_groups(doc):
        instance_id = getattr(device_group, "QetInstanceId", "").strip()
        if not instance_id:
            continue
        element_uuids = set()
        group_element_uuid = getattr(device_group, "QetElementUuid", "").strip()
        if group_element_uuid:
            element_uuids.add(group_element_uuid)
        for terminal_obj in _iter_terminal_objects(device_group):
            terminal_element_uuid = getattr(terminal_obj, "QetElementUuid", "").strip()
            if terminal_element_uuid:
                element_uuids.add(terminal_element_uuid)
        for element_uuid in sorted(element_uuids):
            key = (element_uuid, instance_id)
            if key in seen:
                continue
            seen.add(key)
            bindings.append(
                {
                    "element_uuid": element_uuid,
                    "device_instance_id": instance_id,
                }
            )
    return bindings


def _stable_terminal_instance_id(project_uuid, device_instance_id, terminal_obj):
    values = [
        project_uuid,
        device_instance_id,
        getattr(terminal_obj, "QetElementUuid", "").strip(),
        getattr(terminal_obj, "QetTerminalUuid", "").strip(),
        getattr(terminal_obj, "QetTemplateSlotName", "").strip(),
        getattr(terminal_obj, "Label", "").strip(),
        getattr(terminal_obj, "Name", "").strip(),
    ]
    seed = "qet-freecad-writeback-terminal|" + "|".join(values)
    return str(uuid.uuid5(uuid.NAMESPACE_URL, seed))


def _writeback_terminal_instance_id(project_uuid, terminal_obj, device_instance_id, used_ids):
    terminal_instance_id = (
        getattr(terminal_obj, "QetTerminalInstanceId", "").strip()
        or getattr(terminal_obj, "QetInstanceId", "").strip()
        or ""
    )
    if (
        not terminal_instance_id
        or terminal_instance_id == device_instance_id
        or terminal_instance_id in used_ids
    ):
        terminal_instance_id = _stable_terminal_instance_id(
            project_uuid,
            device_instance_id,
            terminal_obj,
        )
        suffix = 1
        while terminal_instance_id in used_ids:
            terminal_instance_id = str(uuid.uuid5(
                uuid.NAMESPACE_URL,
                "{0}|{1}".format(terminal_instance_id, suffix),
            ))
            suffix += 1
        TerminalObjects.ensure_string_property(
            terminal_obj,
            "QetTerminalInstanceId",
            "QET Exchange",
            "Stable 3D terminal instance UUID",
            terminal_instance_id,
        )
    used_ids.add(terminal_instance_id)
    return terminal_instance_id


def _collect_terminal_bindings(doc):
    bindings = []
    seen = set()
    used_terminal_instance_ids = set()
    project_uuid = _project_uuid_from_doc(doc)
    for device_group in _iter_device_groups(doc):
        instance_id = getattr(device_group, "QetInstanceId", "").strip()
        for terminal_obj in _iter_terminal_objects(device_group):
            terminal_uuid = getattr(terminal_obj, "QetTerminalUuid", "").strip()
            binding_mode = getattr(terminal_obj, "QetTerminalBindingMode", "").strip().lower()
            if (
                TerminalObjects.is_local_terminal_uuid(terminal_uuid)
                or binding_mode == TerminalObjects.TERMINAL_BINDING_MODE_LOCAL
            ):
                continue
            terminal_instance_id = _writeback_terminal_instance_id(
                project_uuid,
                terminal_obj,
                instance_id,
                used_terminal_instance_ids,
            )
            if not terminal_uuid or not terminal_instance_id:
                continue
            key = (terminal_uuid, terminal_instance_id)
            if key in seen:
                continue
            seen.add(key)
            bindings.append(
                {
                    "terminal_uuid": terminal_uuid,
                    "device_instance_id": instance_id,
                    "terminal_instance_id": terminal_instance_id,
                }
            )
    return bindings


def _project_uuid_from_doc(doc, payload=None):
    root = _root_group(doc)
    if root is not None:
        project_uuid = getattr(root, "QetProjectUuid", "").strip()
        if project_uuid:
            return project_uuid
    return _project_uuid_from_payload(payload)


def write_back_document(doc=None, scene_path="", payload=None):
    if doc is None:
        doc = App.ActiveDocument
    if doc is None:
        raise ExchangeWriteBackError("No active FreeCAD document is available.")

    _append_debug_log(
        "write_back_document starting: doc={0}, path={1}, objects={2}, requested_scene_path={3}, env_json={4}".format(
            _doc_name(doc),
            _doc_path(doc) or "<unsaved>",
            _doc_object_count(doc),
            scene_path or "<empty>",
            os.environ.get(ENV_JSON_PATH, "").strip() or "<empty>",
        )
    )

    scene_path = _scene_path_from_doc(doc, scene_path)
    output_path = _output_path_for_exchange_json() or _output_path_for_scene(scene_path)
    if not output_path:
        raise ExchangeWriteBackError(
            "Cannot determine the 3d_to_2d.json output path."
        )

    payload = _payload_for_writeback(scene_path, payload)
    _sync_terminals_for_writeback(doc, scene_path, payload)

    project_uuid = _project_uuid_from_doc(doc, payload)
    if not project_uuid:
        raise ExchangeWriteBackError(
            "Cannot determine project_uuid for write-back."
        )

    report = {
        "schema_version": "2.0",
        "project_uuid": project_uuid,
        "generated_at": _format_timestamp(),
        "instances": _collect_instance_bindings(doc),
        "terminals": _collect_terminal_bindings(doc),
        "output_path": output_path,
    }

    output_dir = str(Path(output_path).parent)
    os.makedirs(output_dir, exist_ok=True)
    _append_debug_log(
        "write_back_document collected: project_uuid={0}, scene_path={1}, output_path={2}, instance_count={3}, terminal_count={4}".format(
            project_uuid,
            scene_path or "<empty>",
            output_path,
            len(report["instances"]),
            len(report["terminals"]),
        )
    )
    Path(output_path).write_text(
        json.dumps(
            {
                "schema_version": report["schema_version"],
                "project_uuid": report["project_uuid"],
                "generated_at": report["generated_at"],
                "instances": report["instances"],
                "terminals": report["terminals"],
            },
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )

    _append_debug_log(
        "write_back_document completed: instances={0}, terminals={1}, path={2}".format(
            len(report["instances"]),
            len(report["terminals"]),
            output_path,
        )
    )
    try:
        App.Console.PrintMessage(
            "[FreeCADExchange] Wrote 3d_to_2d.json to {0}\n".format(output_path)
        )
    except Exception:
        pass
    return report


def _is_exchange_document(doc):
    if doc is None:
        return False

    if _root_group(doc) is not None:
        return True

    for obj in doc.Objects:
        if _is_device_group(obj):
            return True
    return False


class _WriteBackObserver:
    def slotCreatedDocument(self, doc):
        _append_debug_log(
            "write-back observer slotCreatedDocument: doc={0}, path={1}, objects={2}".format(
                _doc_name(doc),
                _doc_path(doc) or "<unsaved>",
                _doc_object_count(doc),
            )
        )

    def slotDeletedDocument(self, doc):
        _append_debug_log(
            "write-back observer slotDeletedDocument: doc={0}, path={1}, objects={2}".format(
                _doc_name(doc),
                _doc_path(doc) or "<unsaved>",
                _doc_object_count(doc),
            )
        )

    def slotActivateDocument(self, doc):
        _append_debug_log(
            "write-back observer slotActivateDocument: doc={0}, path={1}, objects={2}".format(
                _doc_name(doc),
                _doc_path(doc) or "<unsaved>",
                _doc_object_count(doc),
            )
        )

    def slotStartSaveDocument(self, doc, name):
        _append_debug_log(
            "write-back observer slotStartSaveDocument: doc={0}, doc_path={1}, target_name={2}, objects={3}, exchange_doc={4}".format(
                _doc_name(doc),
                _doc_path(doc) or "<unsaved>",
                name or "<empty>",
                _doc_object_count(doc),
                _is_exchange_document(doc),
            )
        )
        if not _is_exchange_document(doc):
            return
        try:
            sync_terminals_from_current_payload(doc, scene_path=name)
        except Exception as exc:
            _append_debug_log("write-back terminal sync before save failed: {0}".format(exc))
            _append_debug_log(traceback.format_exc())
            try:
                App.Console.PrintError(
                    "[FreeCADExchange] terminal sync before save failed: {0}\n".format(exc)
                )
            except Exception:
                pass

    def slotFinishSaveDocument(self, doc, name):
        _append_debug_log(
            "write-back observer slotFinishSaveDocument: doc={0}, doc_path={1}, target_name={2}, objects={3}, exchange_doc={4}".format(
                _doc_name(doc),
                _doc_path(doc) or "<unsaved>",
                name or "<empty>",
                _doc_object_count(doc),
                _is_exchange_document(doc),
            )
        )
        if not _is_exchange_document(doc):
            return
        try:
            write_back_document(doc, scene_path=name)
        except Exception as exc:
            _append_debug_log("write-back after save failed: {0}".format(exc))
            _append_debug_log(traceback.format_exc())
            try:
                App.Console.PrintError(
                    "[FreeCADExchange] write-back after save failed: {0}\n".format(exc)
                )
            except Exception:
                pass


def ensure_document_observer_installed():
    if getattr(App, STATE_WRITEBACK_OBSERVER, None) is not None:
        _append_debug_log("write-back observer already installed")
        return getattr(App, STATE_WRITEBACK_OBSERVER)

    observer = _WriteBackObserver()
    try:
        App.addDocumentObserver(observer)
    except Exception as exc:
        _append_debug_log("failed to add write-back observer: {0}".format(exc))
        return None

    setattr(App, STATE_WRITEBACK_OBSERVER, observer)
    _append_debug_log(
        "write-back observer installed: observer_id={0}".format(id(observer))
    )
    return observer


class CommandWriteBack:
    def GetResources(self):
        return {
            "MenuText": "Write Back 3D Binding",
            "ToolTip": "Generate 3d_to_2d.json from the current FreeCAD document",
        }

    def IsActive(self):
        return App.ActiveDocument is not None

    def Activated(self):
        try:
            report = write_back_document(App.ActiveDocument)
            try:
                App.Console.PrintMessage(
                    "[FreeCADExchange] Write-back completed: {0} instances, {1} terminals\n".format(
                        len(report["instances"]),
                        len(report["terminals"]),
                    )
                )
            except Exception:
                pass
        except Exception as exc:
            try:
                App.Console.PrintError(
                    "[FreeCADExchange] Write-back failed: {0}\n".format(exc)
                )
            except Exception:
                pass


_COMMANDS_REGISTERED = False


def register_commands():
    global _COMMANDS_REGISTERED
    if _COMMANDS_REGISTERED:
        return
    if Gui is None:
        return
    try:
        Gui.addCommand("QET_Exchange_WriteBack", CommandWriteBack())
        _COMMANDS_REGISTERED = True
    except Exception as exc:
        _append_debug_log("failed to register write-back command: {0}".format(exc))


register_commands()
ensure_document_observer_installed()