pdm/creoparams/ugc.py

"""Lettura del container UGC dei file nativi Creo (.prt, .asm, .drw).

Un file nativo Creo si apre con un header ASCII in chiaro, seguito da un
indice delle sezioni (UGC_TOC) e dal corpo binario. Solo alcune sezioni sono
compresse: la tabella dei parametri non lo e', ed e' quello che sfruttiamo.

Struttura osservata su Creo 9.0.3.0:

    #UGC:2 PART/SHEETMETAL 2038 1760 805 1 1 15 4000 2022414 000002d7 \\
    #- VERS 0 0                                                       \\
    #- CMNM 00e9258400200.prt                                         \\
    #-END_OF_UGC_HEADER
    #Creo  TM  9.0  (c) 2026 by PTC Inc.  All Rights Reserved. 9.0.3.0
    #UGC_TOC 2 32 81 17###...
    ND:0:Model_L05_PX:1 ab4 440 557 2038 a 1 b715 85d3###...
    BasicData 2c3b 15767 1575c 2038 9 -1 e3ea 48f0###...
    ...
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field

HEADER_END = b"#-END_OF_UGC_HEADER"
TOC_MARKER = b"#UGC_TOC"
MAGIC = b"#UGC:"

# Quanto leggiamo per cercare header e TOC. L'indice sta sempre nei primi KB.
_HEAD_WINDOW = 64 * 1024

_BANNER_RE = re.compile(r"^#(Creo|Pro/ENGINEER).*", re.IGNORECASE)
_VERSION_RE = re.compile(r"(\d+(?:\.\d+){1,3})\s*$")


@dataclass
class Section:
    """Una voce dell'indice UGC_TOC.

    `alt_size` e' il terzo campo numerico. Per le sezioni compresse e' maggiore
    di `size` (dimensione decompressa); per quelle non compresse e' leggermente
    minore. Non lo interpretiamo oltre: ci serve solo come indizio di
    compressione.
    """

    name: str
    offset: int
    size: int
    alt_size: int
    extra: tuple = ()

    @property
    def end(self) -> int:
        return self.offset + self.size

    @property
    def looks_compressed(self) -> bool:
        return self.alt_size > self.size

    def contains(self, pos: int) -> bool:
        return self.offset <= pos < self.end


@dataclass
class Container:
    """Un file nativo Creo aperto e con header/TOC interpretati."""

    data: bytes
    kind: str = ""              # PART, PART/SHEETMETAL, DRAWING, ASSEMBLY...
    creo_version: str = ""      # es. "9.0.3.0"
    banner: str = ""            # riga di copyright completa
    common_name: str = ""       # nome originale del modello (campo CMNM)
    header: dict = field(default_factory=dict)
    sections: list = field(default_factory=list)

    @property
    def base_kind(self) -> str:
        return self.kind.split("/", 1)[0]

    @property
    def subtype(self) -> str:
        parts = self.kind.split("/", 1)
        return parts[1] if len(parts) > 1 else ""

    def section_at(self, pos: int) -> str:
        for section in self.sections:
            if section.contains(pos):
                return section.name
        return ""


class NotACreoFile(ValueError):
    """Il file non ha un header UGC riconoscibile."""


def _decode_cmnm(raw: str) -> str:
    """Il campo CMNM e' prefissato da 3 cifre esadecimali con la lunghezza.

    Esempio: "00e9258400200.prt" -> 0x00e = 14 = len("9258400200.prt").
    Se il prefisso non torna, restituiamo il valore grezzo.
    """
    if len(raw) > 3:
        try:
            declared = int(raw[:3], 16)
        except ValueError:
            return raw
        if declared == len(raw) - 3:
            return raw[3:]
    return raw


def _parse_toc_line(line: str) -> Section | None:
    # Le righe del TOC sono paddate a destra con '#'.
    line = line.rstrip("#").strip()
    if not line:
        return None
    parts = line.split()
    if len(parts) < 4:
        return None
    try:
        offset = int(parts[1], 16)
        size = int(parts[2], 16)
        alt_size = int(parts[3], 16)
    except ValueError:
        return None
    return Section(parts[0], offset, size, alt_size, tuple(parts[4:]))


def parse(data: bytes) -> Container:
    """Interpreta header e indice. Non tocca il corpo binario."""
    if not data.startswith(MAGIC):
        raise NotACreoFile("header UGC assente (i primi byte non sono '#UGC:')")

    head_end = data.find(HEADER_END)
    if head_end < 0:
        raise NotACreoFile("marcatore di fine header UGC non trovato")

    container = Container(data=data)
    window = data[: max(head_end + _HEAD_WINDOW, head_end)].decode(
        "latin-1", errors="replace"
    )
    lines = window.split("\n")

    in_toc = False
    for index, line in enumerate(lines):
        stripped = line.rstrip().rstrip("\\").rstrip()

        if index == 0:
            # "#UGC:2 PART/SHEETMETAL 2038 1760 ..."
            tokens = stripped.split()
            container.kind = tokens[1] if len(tokens) > 1 else ""
            continue

        if in_toc:
            if stripped.startswith("#"):
                break  # fine dell'indice, inizia il corpo
            section = _parse_toc_line(stripped)
            if section:
                container.sections.append(section)
                # NEXT_TOC_ENTRY non punta a un secondo indice ASCII: il corpo
                # prosegue con blocchi marcati inline. Ci fermiamo qui.
                if section.name == "NEXT_TOC_ENTRY":
                    break
            continue

        if stripped.startswith("#- "):
            key, _, value = stripped[3:].partition(" ")
            value = value.strip()
            container.header[key] = value
            if key == "CMNM":
                container.common_name = _decode_cmnm(value)
        elif stripped.startswith(TOC_MARKER.decode()):
            in_toc = True
        elif _BANNER_RE.match(stripped):
            container.banner = stripped.lstrip("#").strip()
            match = _VERSION_RE.search(container.banner)
            if match:
                container.creo_version = match.group(1)

    return container


#: Oltre l'indice, il corpo prosegue con blocchi introdotti da una riga
#: "#Nome". I nomi possono contenere ':' e '#' (es. "PicturePrimdata#0").
_BLOCK_RE = re.compile(rb"\n#([A-Za-z][A-Za-z0-9_:#]{2,40})\n")


def iter_blocks(data: bytes) -> "list[tuple[str, bytes, bool]]":
    """I blocchi che seguono l'indice, come (nome, contenuto, compresso).

    Il contenuto e' gia' decompresso quando serve: lo stesso blocco puo'
    essere memorizzato compresso in un file e in chiaro in un altro, quindi
    la scelta va fatta guardando i byte, non il nome.
    """
    from . import lzw

    marks = [(m.start(), m.end(), m.group(1).decode("latin-1"))
             for m in _BLOCK_RE.finditer(data)]

    blocks = []
    for index, (start, body, name) in enumerate(marks):
        end = marks[index + 1][0] if index + 1 < len(marks) else len(data)
        payload = data[body:end]
        if payload.startswith(lzw.MAGIC):
            try:
                payload, _ = lzw.decompress(payload)
            except (ValueError, IndexError):
                pass
            else:
                blocks.append((name, payload, True))
                continue
        blocks.append((name, payload, False))
    return blocks


def read(path) -> Container:
    with open(path, "rb") as handle:
        return parse(handle.read())