create/src/Mod/Create/kc_format.py

"""
kc_format.py — .kc file format support.

Handles two responsibilities:
1. Round-trip preservation: caches silo/ ZIP entries before FreeCAD's C++
   save rewrites the ZIP from scratch, then re-injects them after save.
2. Manifest auto-creation: ensures every .kc file has silo/manifest.json.
"""

import json
import os
import uuid
import zipfile
from datetime import datetime, timezone

import FreeCAD

# Cache: filepath -> {entry_name: bytes}
_silo_cache = {}

# Pre-reinject hooks: called with (doc, filename, entries) before ZIP write.
_pre_reinject_hooks = []


def register_pre_reinject(callback):
    """Register a callback invoked before silo/ entries are written to ZIP.

    Signature: callback(doc, filename, entries) -> None
    ``entries`` is a dict {entry_name: bytes} or None. Mutate in place.
    """
    _pre_reinject_hooks.append(callback)


def _metadata_save_hook(doc, filename, entries):
    """Write dirty metadata back to the silo/ cache before ZIP write."""
    obj = doc.getObject("SiloMetadata")
    if obj is None or not hasattr(obj, "Proxy"):
        return
    proxy = obj.Proxy
    if proxy is None or not proxy.is_dirty():
        return
    entries["silo/metadata.json"] = obj.RawContent.encode("utf-8")


register_pre_reinject(_metadata_save_hook)


def _manifest_enrich_hook(doc, filename, entries):
    """Populate silo_instance and part_uuid from the tracked Silo object."""
    raw = entries.get("silo/manifest.json")
    if raw is None:
        return
    try:
        manifest = json.loads(raw)
    except (json.JSONDecodeError, ValueError):
        return

    changed = False

    # Populate part_uuid from SiloItemId if available.
    for obj in doc.Objects:
        if hasattr(obj, "SiloItemId") and obj.SiloItemId:
            if manifest.get("part_uuid") != obj.SiloItemId:
                manifest["part_uuid"] = obj.SiloItemId
                changed = True
            break

    # Populate silo_instance from Silo settings.
    if not manifest.get("silo_instance"):
        try:
            import silo_commands

            api_url = silo_commands._get_api_url()
            if api_url:
                # Strip /api suffix to get base instance URL.
                instance = api_url.rstrip("/")
                if instance.endswith("/api"):
                    instance = instance[:-4]
                manifest["silo_instance"] = instance
                changed = True
        except Exception:
            pass

    if changed:
        entries["silo/manifest.json"] = (json.dumps(manifest, indent=2) + "\n").encode(
            "utf-8"
        )


register_pre_reinject(_manifest_enrich_hook)


def _solver_context_hook(doc, filename, entries):
    """Pack solver context into silo/solver/context.json for assemblies."""
    try:
        for obj in doc.Objects:
            if obj.TypeId == "Assembly::AssemblyObject":
                ctx = obj.getSolveContext()
                if ctx:  # non-empty means we have grounded parts
                    entries["silo/solver/context.json"] = (
                        json.dumps(ctx, indent=2) + "\n"
                    ).encode("utf-8")
                break  # one assembly per document
    except Exception as exc:
        FreeCAD.Console.PrintWarning(f"kc_format: solver context hook failed: {exc}\n")


register_pre_reinject(_solver_context_hook)


KC_VERSION = "1.0"


def _default_manifest():
    """Generate a default silo/manifest.json for new .kc files."""
    now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    username = os.environ.get("USER", os.environ.get("USERNAME", "unknown"))
    return {
        "kc_version": KC_VERSION,
        "silo_instance": None,
        "part_uuid": str(uuid.uuid4()),
        "revision_hash": None,
        "created_at": now,
        "modified_at": now,
        "created_by": username,
    }


class _KcFormatObserver:
    """Document observer that preserves silo/ entries across saves."""

    def slotStartSaveDocument(self, doc, filename):
        """Before save: cache silo/ entries from the existing file."""
        if not filename.lower().endswith(".kc"):
            return
        if not os.path.isfile(filename):
            return
        try:
            with zipfile.ZipFile(filename, "r") as zf:
                entries = {}
                for name in zf.namelist():
                    if name.startswith("silo/"):
                        entries[name] = zf.read(name)
                if entries:
                    _silo_cache[filename] = entries
        except Exception:
            pass

    def slotFinishSaveDocument(self, doc, filename):
        """After save: re-inject cached silo/ entries and ensure manifest."""
        if not filename.lower().endswith(".kc"):
            _silo_cache.pop(filename, None)
            return
        entries = _silo_cache.pop(filename, None)
        if entries is None:
            entries = {}
        for _hook in _pre_reinject_hooks:
            try:
                _hook(doc, filename, entries)
            except Exception as exc:
                FreeCAD.Console.PrintWarning(
                    f"kc_format: pre_reinject hook failed: {exc}\n"
                )
        try:
            with zipfile.ZipFile(filename, "a") as zf:
                existing = set(zf.namelist())
                # Re-inject cached silo/ entries
                if entries:
                    for name, data in entries.items():
                        if name not in existing:
                            zf.writestr(name, data)
                            existing.add(name)
                # Ensure silo/manifest.json exists
                if "silo/manifest.json" not in existing:
                    manifest = _default_manifest()
                    zf.writestr(
                        "silo/manifest.json",
                        json.dumps(manifest, indent=2) + "\n",
                    )
                else:
                    # Update modified_at timestamp
                    raw = zf.read("silo/manifest.json")
                    manifest = json.loads(raw)
                    now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
                    if manifest.get("modified_at") != now:
                        manifest["modified_at"] = now
                        # ZipFile append mode can't overwrite; write new entry
                        # (last duplicate wins in most ZIP readers)
                        zf.writestr(
                            "silo/manifest.json",
                            json.dumps(manifest, indent=2) + "\n",
                        )
        except Exception as e:
            FreeCAD.Console.PrintWarning(
                f"kc_format: failed to update .kc silo/ entries: {e}\n"
            )


def update_manifest_fields(filename, updates):
    """Update fields in an existing .kc manifest after save.

    *filename*: path to the .kc file.
    *updates*: dict of field_name -> value to merge into the manifest.

    Used by silo_commands to write ``revision_hash`` after a successful
    upload (which happens after the ZIP has already been written by save).
    """
    if not filename or not filename.lower().endswith(".kc"):
        return
    if not os.path.isfile(filename):
        return
    try:
        with zipfile.ZipFile(filename, "a") as zf:
            if "silo/manifest.json" not in zf.namelist():
                return
            raw = zf.read("silo/manifest.json")
            manifest = json.loads(raw)
            manifest.update(updates)
            zf.writestr(
                "silo/manifest.json",
                json.dumps(manifest, indent=2) + "\n",
            )
    except Exception as e:
        FreeCAD.Console.PrintWarning(f"kc_format: failed to update manifest: {e}\n")


def register():
    """Connect to application-level save signals."""
    FreeCAD.addDocumentObserver(_KcFormatObserver())