HartOMat/backend/app/services/part_key_service.py

"""Part key generation and scene manifest building for the USD pipeline.

The `resolved_material_assignments` JSONB schema written by `generate_usd_master_task`:
    {part_key: {"source_name": str, "prim_path": str}}

The `manual_material_overrides` JSONB schema written by `PUT /cad/{id}/part-materials` (Priority 4):
    {part_key: material_name_str}

The `source_material_assignments` JSONB schema written by the Excel importer (future):
    {source_part_name: material_name_str}

No pxr imports — all data is read from JSONB columns, never from USD files directly.
"""
from __future__ import annotations

import hashlib
import re


# ── Part key generation ───────────────────────────────────────────────────────

_AF_RE = re.compile(r'_AF\d+$', re.IGNORECASE)


def generate_part_key(
    xcaf_label_path: str,
    source_name: str,
    existing_keys: set[str] | None = None,
) -> str:
    """Deterministic slug from source_name, max 64 chars, unique within assembly.

    - Strips `_AF\\d+` OCC suffix from source_name before slugifying.
    - Falls back to sha256 digest of xcaf_label_path if slug is empty.
    - Deduplicates by appending _2, _3, ... if existing_keys is provided.
    """
    base = _AF_RE.sub('', source_name) if source_name else ''
    # Split camelCase before slugifying: "RingOuter" → "Ring_Outer"
    base = re.sub(r'([a-z])([A-Z])', r'\1_\2', base)
    slug = re.sub(r'[^a-z0-9]+', '_', base.lower()).strip('_')
    if not slug:
        slug = f"part_{hashlib.sha256(xcaf_label_path.encode()).hexdigest()[:8]}"
    slug = slug[:50]

    if existing_keys is None:
        return slug

    key = slug
    n = 2
    while key in existing_keys:
        key = f"{slug}_{n}"
        n += 1
    existing_keys.add(key)
    return key


# ── Scene manifest building ───────────────────────────────────────────────────

def build_scene_manifest(cad_file, usd_asset=None) -> dict:
    """Build a scene manifest dict from CadFile ORM object.

    Source of part list (priority order):
    1. `resolved_material_assignments` — keyed by partKey (set by generate_usd_master_task)
    2. `parsed_objects["objects"]` — list of source name strings from STEP extraction
    3. Empty manifest if neither is available.

    Material assignment priority per part:
    1. `manual_material_overrides[part_key]` — provenance "manual"
    2. `resolved_material_assignments[part_key]["material"]` — provenance "auto"
    3. substring match in `source_material_assignments` against source_name — provenance "source"
    4. None, is_unassigned=True — provenance "default"
    """
    cad_id = str(cad_file.id)
    resolved = cad_file.resolved_material_assignments or {}
    manual = cad_file.manual_material_overrides or {}
    source = cad_file.source_material_assignments or {}

    parts: list[dict] = []
    unmatched_source_rows: list[str] = []
    unassigned_parts: list[str] = []

    if resolved:
        # Build from resolved assignments (USD pipeline has run)
        for part_key, meta in resolved.items():
            source_name = meta.get("source_name", "") if isinstance(meta, dict) else ""
            prim_path = meta.get("prim_path") if isinstance(meta, dict) else None

            effective_material, provenance = _resolve_material(
                part_key, source_name, manual, resolved, source
            )
            is_unassigned = effective_material is None

            parts.append({
                "part_key": part_key,
                "source_name": source_name,
                "prim_path": prim_path,
                "effective_material": effective_material,
                "assignment_provenance": provenance,
                "is_unassigned": is_unassigned,
            })
            if is_unassigned:
                unassigned_parts.append(part_key)

    elif cad_file.parsed_objects:
        # Fall back to parsed_objects from STEP extraction
        object_names: list[str] = cad_file.parsed_objects.get("objects") or []
        seen_keys: set[str] = set()
        for source_name in object_names:
            part_key = generate_part_key(source_name, source_name, seen_keys)
            effective_material, provenance = _resolve_material(
                part_key, source_name, manual, resolved, source
            )
            is_unassigned = effective_material is None

            parts.append({
                "part_key": part_key,
                "source_name": source_name,
                "prim_path": None,
                "effective_material": effective_material,
                "assignment_provenance": provenance,
                "is_unassigned": is_unassigned,
            })
            if is_unassigned:
                unassigned_parts.append(part_key)

    # Find source rows not matched to any part
    matched_source_names = {p["source_name"].lower() for p in parts}
    for src_key in source:
        if not any(
            src_key.lower() in sn or sn in src_key.lower()
            for sn in matched_source_names
        ):
            unmatched_source_rows.append(src_key)

    return {
        "cad_file_id": cad_id,
        "parts": parts,
        "unmatched_source_rows": unmatched_source_rows,
        "unassigned_parts": unassigned_parts,
    }


def _resolve_material(
    part_key: str,
    source_name: str,
    manual: dict,
    resolved: dict,
    source: dict,
) -> tuple[str | None, str]:
    """Return (material_name, provenance) for one part using priority order."""
    # 1. Manual override
    if part_key in manual and manual[part_key]:
        return str(manual[part_key]), "manual"

    # 2. Auto-resolved from USD pipeline
    meta = resolved.get(part_key)
    if isinstance(meta, dict) and meta.get("material"):
        return str(meta["material"]), "auto"

    # 3. Substring match in source_material_assignments against source_name
    sn_lower = source_name.lower()
    for src_key, src_mat in source.items():
        if src_key.lower() in sn_lower or sn_lower in src_key.lower():
            if src_mat:
                return str(src_mat), "source"

    # 4. Unassigned
    return None, "default"


# ── Effective assignments for render pipeline ─────────────────────────────────

def get_effective_assignments(cad_file) -> dict[str, str]:
    """Return {part_key: material_name} merged from all three layers.

    Used by the render pipeline when building the material map (Priority 5).
    """
    manifest = build_scene_manifest(cad_file)
    return {
        p["part_key"]: p["effective_material"]
        for p in manifest["parts"]
        if p["effective_material"] is not None
    }