step2urdf/app/services/step_text.py

from __future__ import annotations

import re
from collections import defaultdict
from pathlib import Path

from app.models import PartNode


_X2_RE = re.compile(r"\\X2\\([0-9A-Fa-f]+)\\X0\\")
_PRODUCT_RE = re.compile(
    r"#(\d+)\s*=\s*PRODUCT\s*\(\s*'([^']*)'\s*,\s*'([^']*)'\s*,\s*'([^']*)'",
    re.IGNORECASE,
)
# PRODUCT_DEFINITION_FORMATION(... , #product, ...)
_PDF_RE = re.compile(
    r"#(\d+)\s*=\s*PRODUCT_DEFINITION_FORMATION\s*\([^;]*?#(\d+)\s*\)\s*;",
    re.IGNORECASE | re.DOTALL,
)
# PRODUCT_DEFINITION(..., #pdf, ...)
_PD_RE = re.compile(
    r"#(\d+)\s*=\s*PRODUCT_DEFINITION\s*\([^;]*?#(\d+)\s*\)\s*;",
    re.IGNORECASE | re.DOTALL,
)
# NEXT_ASSEMBLY_USAGE_OCCURRENCE('name', ..., #parent_pd, #child_pd, ...)
_NAUO_RE = re.compile(
    r"#(\d+)\s*=\s*NEXT_ASSEMBLY_USAGE_OCCURRENCE\s*\(\s*'([^']*)'\s*,"
    r"[^;]*?#(\d+)\s*,\s*#(\d+)\s*",
    re.IGNORECASE | re.DOTALL,
)


def decode_step_string(s: str) -> str:
    def repl(m: re.Match[str]) -> str:
        hexpart = m.group(1)
        chars: list[str] = []
        for i in range(0, len(hexpart), 4):
            chars.append(chr(int(hexpart[i : i + 4], 16)))
        return "".join(chars)

    return _X2_RE.sub(repl, s)


def parse_step_products(step_path: Path) -> dict:
    text = step_path.read_text(encoding="utf-8", errors="replace")

    products: dict[str, str] = {}
    for m in _PRODUCT_RE.finditer(text):
        eid, name, *_ = m.groups()
        products[eid] = decode_step_string(name)

    # pdf_id -> product_id
    pdf_to_product: dict[str, str] = {}
    for m in _PDF_RE.finditer(text):
        pdf_id, product_id = m.groups()
        pdf_to_product[pdf_id] = product_id

    # pd_id -> pdf_id
    pd_to_pdf: dict[str, str] = {}
    for m in _PD_RE.finditer(text):
        pd_id, pdf_id = m.groups()
        pd_to_pdf[pd_id] = pdf_id

    def pd_name(pd_id: str) -> str:
        pdf = pd_to_pdf.get(pd_id)
        if not pdf:
            return f"pd_{pd_id}"
        prod = pdf_to_product.get(pdf)
        if not prod:
            return f"pdf_{pdf}"
        return products.get(prod, f"product_{prod}")

    children: dict[str, list[str]] = defaultdict(list)
    parents: dict[str, str] = {}
    edge_names: dict[tuple[str, str], str] = {}

    for m in _NAUO_RE.finditer(text):
        _eid, edge_name, parent_pd, child_pd = m.groups()
        children[parent_pd].append(child_pd)
        parents[child_pd] = parent_pd
        edge_names[(parent_pd, child_pd)] = decode_step_string(edge_name)

    # Roots = PDs that appear as parents or products but have no parent
    all_pd = set(pd_to_pdf) | set(children) | set(parents)
    roots = [pd for pd in all_pd if pd not in parents]
    if not roots and products:
        # fallback: flat product list
        nodes = [
            PartNode(id=pid, name=name, parent_id=None)
            for pid, name in products.items()
        ]
        return {
            "root_name": nodes[0].name if nodes else "assembly",
            "parts": nodes,
            "product_names": sorted({p.name for p in nodes}),
            "stats": {
                "products": len(products),
                "nauo": 0,
                "axis2_placement": text.upper().count("AXIS2_PLACEMENT_3D"),
                "circles": len(re.findall(r"\bCIRCLE\b", text, flags=re.I)),
                "cylinders": text.upper().count("CYLINDRICAL_SURFACE"),
                "mode": "flat_products",
            },
        }

    # Prefer the largest tree root as assembly root
    def subtree_size(pd: str, seen: set[str] | None = None) -> int:
        seen = seen or set()
        if pd in seen:
            return 0
        seen.add(pd)
        return 1 + sum(subtree_size(c, seen) for c in children.get(pd, []))

    roots_sorted = sorted(roots, key=subtree_size, reverse=True)
    root_pd = roots_sorted[0] if roots_sorted else None

    nodes: list[PartNode] = []
    for pd in sorted(all_pd, key=lambda x: int(x) if x.isdigit() else 0):
        nodes.append(
            PartNode(
                id=pd,
                name=pd_name(pd),
                parent_id=parents.get(pd),
                children=list(children.get(pd, [])),
            )
        )

    product_names = sorted({decode_step_string(n) for n in products.values()})
    root_name = pd_name(root_pd) if root_pd else ""
    if not root_name or root_name.startswith(("pd_", "pdf_", "product_")):
        preferred = next((n for n in product_names if "装配体" in n or "assembly" in n.lower()), None)
        root_name = preferred or (product_names[0] if product_names else "assembly")

    return {
        "root_name": root_name,
        "parts": nodes,
        "product_names": product_names,
        "stats": {
            "products": len(products),
            "product_definitions": len(pd_to_pdf),
            "nauo": len(edge_names),
            "roots": len(roots),
            "axis2_placement": text.upper().count("AXIS2_PLACEMENT_3D"),
            "circles": len(re.findall(r"\bCIRCLE\b", text, flags=re.I)),
            "cylinders": text.upper().count("CYLINDRICAL_SURFACE"),
            "mode": "assembly_tree",
            "geometry_backend": "text_only",
        },
    }