pythonocc-step-editor/step_editor/parametric_component.py

from __future__ import annotations

import argparse
from datetime import datetime
import json
import re
from pathlib import Path
import sys
import traceback

from .isolated_edit_worker import _execute
from .model import StepModel


COMPONENT_SCHEMA = "step-editor-parametric-component-v1"

_COMPONENT_SEQUENCE_RE = re.compile(r"^(?P<index>\d{3,})_(?P<name>.+)$")
_ACTION_OPERATION_MAP = {
    "push_pull_face": "push_pull_face",
    "push_pull_face_keep_relations": "push_pull_face_keep_relations",
    "move_selected_face_plane_position_local": "move_face_plane_offset_local",
    "move_selected_face_plane_position_by_translation": "translate_face_plane_offset_owning",
    "resize_face_width_local": "resize_face_size_local",
    "resize_face_height_local": "resize_face_size_local",
    "resize_face_width_keep_relations": "resize_face_size_local_keep_relations",
    "resize_face_height_keep_relations": "resize_face_size_local_keep_relations",
    "resize_face_width_owning_scale": "resize_face_size_owning_scale",
    "resize_face_height_owning_scale": "resize_face_size_owning_scale",
    "resize_shell_thickness": "resize_shell_thickness",
    "resize_shell_thickness_owning_scale": "resize_shell_thickness_owning_scale",
    "resize_hole": "resize_cylindrical_hole",
    "resize_multi_selected_holes": "resize_cylindrical_holes_by_refs",
    "move_multi_selected_holes_by_offset": "move_cylindrical_holes_by_offset",
    "suppress_multi_selected_holes": "suppress_cylindrical_holes_by_refs",
    "resize_cylindrical_owning_scale": "resize_cylindrical_owning_scale",
    "resize_hole_depth": "resize_cylindrical_depth",
    "resize_hole_depth_owning_scale": "resize_cylindrical_depth_owning_scale",
    "resize_slot_width": "resize_cylindrical_slot_width",
    "resize_slot_depth": "resize_cylindrical_slot_depth",
    "resize_slot_arc_length": "resize_cylindrical_slot_arc_length",
    "resize_slot_angular_span": "resize_cylindrical_slot_angular_span",
    "resize_slot_total_length": "resize_cylindrical_slot_total_length",
    "resize_slot_center_distance": "resize_cylindrical_slot_center_distance",
    "move_cylindrical_hole_axis": "move_cylindrical_hole_axis",
    "move_cylindrical_slot_axis": "move_cylindrical_slot_axis",
    "suppress_hole": "suppress_cylindrical_hole",
    "resize_boss": "resize_cylindrical_boss",
    "resize_boss_height": "resize_cylindrical_boss_height",
    "resize_cylinder_height": "resize_cylindrical_height",
    "resize_cylindrical_height_owning_scale": "resize_cylindrical_height_owning_scale",
    "move_cylindrical_boss_axis": "move_cylindrical_boss_axis",
    "resize_cone_reference_radius": "resize_cone_reference_radius",
    "resize_cone_semi_angle": "resize_cone_semi_angle",
    "resize_sphere_radius": "resize_sphere_radius",
    "resize_torus_major_radius": "resize_torus_radius",
    "resize_torus_minor_radius": "resize_torus_radius",
    "resize_any_edge_length": "resize_general_edge_length",
    "move_edge_start_point": "move_edge_endpoint",
    "move_edge_end_point": "move_edge_endpoint",
    "move_edge_center_point": "move_edge_center",
    "move_circular_edge_axis_center": "move_circular_edge_axis_center",
    "resize_ellipse_edge_major_radius": "resize_ellipse_edge_axis_radius",
    "resize_ellipse_edge_minor_radius": "resize_ellipse_edge_axis_radius",
    "resize_existing_fillet": "resize_existing_fillet",
    "resize_existing_chamfer": "resize_existing_chamfer",
    "fillet_edge": "fillet_edge",
    "chamfer_edge": "chamfer_edge",
    "chamfer_edge_asymmetric": "chamfer_edge_asymmetric",
    "chamfer_edge_distance_angle": "chamfer_edge_distance_angle",
}


def sanitize_component_name(value: object, fallback: str = "STEP_Parametric") -> str:
    text = str(value or "").strip() or fallback
    text = re.sub(r'[<>:"/\\|?*\x00-\x1f]+', "_", text)
    text = re.sub(r"\s+", "_", text).strip(" ._")
    return text or fallback


def default_component_root(project_root: Path | None = None) -> Path:
    root = project_root or Path(__file__).resolve().parent.parent
    return root / "nodes"


def next_component_dir(root: Path, component_name: object) -> Path:
    base_name = sanitize_component_name(component_name)
    max_index = -1
    if root.is_dir():
        for child in root.iterdir():
            if not child.is_dir():
                continue
            match = _COMPONENT_SEQUENCE_RE.match(child.name)
            if match:
                max_index = max(max_index, int(match.group("index")))
    return root / f"{max_index + 1:03d}_{base_name}"


def component_name_from_step(step_path: object) -> str:
    try:
        stem = Path(str(step_path)).stem
    except Exception:
        stem = ""
    return sanitize_component_name(f"{stem}_STEP参数化组件" if stem else "STEP参数化组件")


def numeric_text(value: object) -> str:
    text = str(value if value is not None else "").strip()
    return text


def json_script_literal(value: object) -> str:
    return json.dumps(value, ensure_ascii=False, indent=4)


def _float_or_text(value: object) -> object:
    if isinstance(value, (int, float)):
        return value
    text = str(value if value is not None else "").strip()
    try:
        return float(text)
    except ValueError:
        return text


def _as_list3(value: object) -> list[float] | None:
    if isinstance(value, str):
        chunks = [chunk.strip() for chunk in value.strip().strip("()[]").replace(";", ",").split(",") if chunk.strip()]
    elif isinstance(value, (tuple, list)):
        chunks = list(value)
    else:
        return None
    if len(chunks) != 3:
        return None
    try:
        return [float(chunks[0]), float(chunks[1]), float(chunks[2])]
    except (TypeError, ValueError):
        return None


def _target_object_id(spec: dict[str, object], selected_kind: str | None, selected_face_id: int | None, selected_edge_id: int | None) -> int | None:
    if spec.get("source_face_id") not in {"", None}:
        return int(spec["source_face_id"])
    action = str(spec.get("action") or "")
    kind = str(selected_kind or "")
    if "edge" in action and selected_edge_id is not None and kind == "edge":
        return int(selected_edge_id)
    if selected_face_id is not None:
        return int(selected_face_id)
    if selected_edge_id is not None:
        return int(selected_edge_id)
    return None


def _target_kind_for_action(action: str, selected_kind: str | None) -> str:
    if "edge" in action:
        return "edge"
    if selected_kind in {"face", "feature", "edge"}:
        return str(selected_kind)
    return "face"


def operation_for_action(action: object) -> str | None:
    return _ACTION_OPERATION_MAP.get(str(action or ""))


def _axis_arg_for_face_size(action: str) -> str | None:
    if "face_width" in action:
        return "width"
    if "face_height" in action:
        return "height"
    return None


def _mode_arg_for_torus(action: str) -> str | None:
    if action == "resize_torus_major_radius":
        return "major"
    if action == "resize_torus_minor_radius":
        return "minor"
    return None


def _endpoint_arg_for_edge(action: str) -> str | None:
    if action == "move_edge_start_point":
        return "start"
    if action == "move_edge_end_point":
        return "end"
    return None


def _ellipse_axis_arg(action: str) -> str | None:
    if action == "resize_ellipse_edge_major_radius":
        return "major"
    if action == "resize_ellipse_edge_minor_radius":
        return "minor"
    return None


def component_edit_config_from_spec(
    *,
    parameter_row: dict[str, str],
    spec: dict[str, object],
    selected_kind: str | None,
    selected_face_id: int | None,
    selected_edge_id: int | None,
    step_path: Path | None,
) -> dict[str, object] | None:
    action = str(spec.get("action") or "")
    operation = operation_for_action(action)
    if action in {"resize_multi_selected_holes", "move_multi_selected_holes_by_offset"}:
        refs = [dict(item) for item in (spec.get("multi_hole_refs") or []) if isinstance(item, dict)]
        if not operation or len(refs) < 2:
            return None
        value_type = str(spec.get("value_type", "number"))
        default_value = parameter_row.get("default", "")
        target_value: object
        if value_type == "vector3":
            target_value = _as_list3(default_value) or _as_list3(spec.get("current_raw")) or default_value
        elif value_type in {"number", "positive", "integer", "integer_or_empty"}:
            target_value = _float_or_text(default_value)
        else:
            target_value = default_value
        target_arg: object = {"param": parameter_row["name"]}
        transform = str(spec.get("target_transform") or "")
        if transform:
            target_arg = {
                "param": parameter_row["name"],
                "transform": transform,
                "context": spec.get("transform_context", {}),
            }
        return {
            "parameter": parameter_row["name"],
            "displayName": parameter_row.get("displayName", parameter_row["name"]),
            "targetKind": "multi_feature",
            "targetId": -1,
            "uiAction": action,
            "operation": operation,
            "args": [refs, target_arg],
            "default": target_value,
            "valueType": value_type,
            "scope": spec.get("scope_key", spec.get("scope_default", "")),
            "scopeLabel": spec.get("scope_label", spec.get("scope_text", "")),
            "sourceStep": str(step_path or ""),
            "parameterKey": spec.get("key", ""),
        }
    target_id = _target_object_id(spec, selected_kind, selected_face_id, selected_edge_id)
    if not operation or target_id is None:
        return None
    value_type = str(spec.get("value_type", "number"))
    default_value = parameter_row.get("default", "")
    target_value: object
    if value_type == "vector3":
        target_value = _as_list3(default_value) or _as_list3(spec.get("current_raw")) or default_value
    elif value_type in {"number", "positive", "integer", "integer_or_empty"}:
        target_value = _float_or_text(default_value)
    else:
        target_value = default_value
    args: list[object] = [int(target_id)]
    target_arg: object = {"param": parameter_row["name"]}
    transform = str(spec.get("target_transform") or "")
    if transform:
        target_arg = {
            "param": parameter_row["name"],
            "transform": transform,
            "context": spec.get("transform_context", {}),
        }
    if action == "suppress_hole":
        target_value = ""
    else:
        args.append(target_arg)
    axis_arg = _axis_arg_for_face_size(action)
    if axis_arg is not None:
        args.append(axis_arg)
    elif action in {"resize_torus_major_radius", "resize_torus_minor_radius"}:
        args.append(_mode_arg_for_torus(action))
    elif action in {"move_edge_start_point", "move_edge_end_point"}:
        args.insert(1, _endpoint_arg_for_edge(action))
    elif action in {"resize_ellipse_edge_major_radius", "resize_ellipse_edge_minor_radius"}:
        args.append(_ellipse_axis_arg(action))
    elif action in {
        "resize_hole_depth",
        "resize_hole_depth_owning_scale",
        "resize_slot_width",
        "resize_slot_depth",
        "resize_slot_arc_length",
        "resize_slot_total_length",
        "resize_slot_center_distance",
    }:
        manual_id = spec.get("manual_bottom_face_id")
        if manual_id in {"", None}:
            manual_id = spec.get("slot_pair_manual_face_id")
        args.append("" if manual_id in {"", None} else manual_id)
    return {
        "parameter": parameter_row["name"],
        "displayName": parameter_row.get("displayName", parameter_row["name"]),
        "targetKind": _target_kind_for_action(action, selected_kind),
        "targetId": int(target_id),
        "uiAction": action,
        "operation": operation,
        "args": args,
        "default": target_value,
        "valueType": value_type,
        "scope": spec.get("scope_key", spec.get("scope_default", "")),
        "scopeLabel": spec.get("scope_label", spec.get("scope_text", "")),
        "sourceStep": str(step_path or ""),
        "parameterKey": spec.get("key", ""),
    }


def render_component_main_py(component: dict[str, object]) -> str:
    project_root = str(Path(__file__).resolve().parent.parent)
    component_name = sanitize_component_name(component.get("componentName") or component_name_from_step(component.get("sourceStep")))
    output_parameter = {
        "name": "output_step",
        "displayName": "输出STEP",
        "type": "file",
        "ioRole": "output",
        "default": "",
    }
    return f'''# -*- coding: utf-8 -*-
"""
STEP 参数化组件。

这个文件按 FlowEditor 节点脚本方式生成：
  1. INPUT_PARAMETERS 是从软件“导出参数”勾选行直接嵌入的输入参数。
  2. PARAMETERS 会额外加上输出 STEP 文件参数，供节点设计器生成输出端口。
  3. execute(inputs, params, context) 是 FlowEditor 调用入口。
  4. main() 只用于本地命令行调试。
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
import sys

PROJECT_ROOT = {json.dumps(project_root, ensure_ascii=False)}
if PROJECT_ROOT and PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

from step_editor.parametric_component import run_embedded_component

INPUT_PARAMETERS = {json_script_literal(component.get("parameters", []))}

OUTPUT_PARAMETERS = [
    {json_script_literal(output_parameter)}
]

PARAMETERS = INPUT_PARAMETERS + OUTPUT_PARAMETERS

COMPONENT = {json_script_literal(component)}

NODE_INFO = {{
    "typeName": {json.dumps(component_name, ensure_ascii=False)},
    "displayName": {json.dumps(component_name, ensure_ascii=False)},
    "category": "几何参数化",
    "icon": "icon.svg",
    "parameters": PARAMETERS,
}}


def _value_from_inputs(name, inputs, params, default=""):
    value = inputs.get(name) if isinstance(inputs, dict) else None
    if value in (None, "") and isinstance(params, dict):
        value = params.get(name)
    if value in (None, ""):
        value = default
    return value


def _component_input_values(inputs, params):
    values = {{}}
    for item in INPUT_PARAMETERS:
        name = item.get("name")
        if not name:
            continue
        values[name] = _value_from_inputs(name, inputs, params, item.get("default", ""))
    return values


def _default_output_step(work_dir, output_dir):
    output_root = output_dir or os.path.join(work_dir, "output")
    os.makedirs(output_root, exist_ok=True)
    return os.path.join(output_root, COMPONENT.get("outputName") or "modified.step")


def run(inputs=None, output_step=None, work_dir=None):
    return run_embedded_component(COMPONENT, inputs=inputs, output_step=output_step, work_dir=work_dir)


def execute(inputs, params, context):
    """
    FlowEditor 调用入口。

    inputs：上游节点传入值，优先级高于 params。
    params：节点属性面板参数。
    context：FlowEditor 上下文，常见字段包括 work_dir / input_dir / output_dir。
    """
    inputs = inputs or {{}}
    params = params or {{}}
    context = context or {{}}
    work_dir = context.get("work_dir") or os.getcwd()
    output_dir = context.get("output_dir") or os.path.join(work_dir, "output")
    output_step = _default_output_step(work_dir, output_dir)
    result = run(
        inputs=_component_input_values(inputs, params),
        output_step=output_step,
        work_dir=work_dir,
    )
    if not result.get("ok"):
        raise RuntimeError(result.get("error") or json.dumps(result, ensure_ascii=False))
    return {{
        "output_step": result.get("outputStep", output_step),
        "outputStep": result.get("outputStep", output_step),
        "sourceStep": result.get("sourceStep", ""),
        "messages": result.get("messages", []),
    }}


def main(argv=None):
    parser = argparse.ArgumentParser(description="Run generated STEP parametric component.")
    parser.add_argument("--inputs", default="", help="JSON file or JSON object with input parameter values.")
    parser.add_argument("--output-step", default="", help="Output STEP file path.")
    parser.add_argument("--work-dir", default="", help="Runtime output directory.")
    args = parser.parse_args(argv)
    inputs = args.inputs
    if inputs:
        candidate = Path(inputs)
        if candidate.is_file():
            inputs = json.loads(candidate.read_text(encoding="utf-8-sig"))
        else:
            inputs = json.loads(inputs)
    else:
        inputs = {{}}
    result = run(inputs=inputs, output_step=args.output_step or None, work_dir=args.work_dir or None)
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0 if result.get("ok") else 2


if __name__ == "__main__":
    raise SystemExit(main())
'''


def export_parametric_component(
    *,
    parameters: list[dict[str, str]],
    edits: list[dict[str, object]],
    source_step: Path | None,
    component_root: Path | None = None,
    component_name: str | None = None,
) -> Path:
    if not parameters:
        raise ValueError("No input parameters selected.")
    root = component_root or default_component_root()
    target_dir = next_component_dir(root, component_name or component_name_from_step(source_step))
    target_dir.mkdir(parents=True, exist_ok=True)
    component = {
        "schema": COMPONENT_SCHEMA,
        "createdAt": datetime.now().isoformat(timespec="seconds"),
        "componentName": component_name or component_name_from_step(source_step),
        "sourceStep": str(source_step or ""),
        "parameters": parameters,
        "edits": edits,
        "outputName": "modified.step",
    }
    target = target_dir / "main.py"
    target.write_text(render_component_main_py(component), encoding="utf-8")
    return target


def _input_values(inputs: object) -> dict[str, object]:
    if not isinstance(inputs, dict):
        return {}
    if isinstance(inputs.get("inputs"), dict):
        return dict(inputs["inputs"])
    rows = inputs.get("parameters")
    if isinstance(rows, list):
        result: dict[str, object] = {}
        for row in rows:
            if isinstance(row, dict) and row.get("name"):
                result[str(row["name"])] = row.get("value", row.get("default", ""))
        return result
    return dict(inputs)


def _parse_vector3(value: object) -> list[float]:
    vector = _as_list3(value)
    if vector is None:
        raise ValueError(f"Expected 3D vector value, got {value!r}.")
    return vector


def _apply_arg_transform(value: object, transform: str, context: object) -> object:
    context = context if isinstance(context, dict) else {}
    if transform == "plane_target_position_to_offset":
        current = float(context.get("current_plane_position"))
        return float(value) - current
    if transform == "radius_to_diameter":
        return float(value) * 2.0
    if transform == "diameter_to_radius":
        return float(value) * 0.5
    if transform == "degrees_to_radians":
        import math

        return math.radians(float(value))
    if transform == "slot_open_angle_degrees_to_angular_span":
        import math

        return math.tau - math.radians(float(value))
    if transform == "target_center_to_translation":
        current = _parse_vector3(context.get("current_center"))
        target = _parse_vector3(value)
        return [target[index] - current[index] for index in range(3)]
    return value


def _resolve_arg(value: object, values: dict[str, object], defaults: dict[str, object]) -> object:
    if isinstance(value, dict) and "param" in value:
        name = str(value.get("param") or "")
        resolved = values.get(name, defaults.get(name, ""))
        transform = str(value.get("transform") or "")
        if transform:
            return _apply_arg_transform(resolved, transform, value.get("context"))
        return resolved
    return value


def run_embedded_component(
    component: dict[str, object],
    *,
    inputs: object | None = None,
    output_step: str | Path | None = None,
    work_dir: str | Path | None = None,
) -> dict[str, object]:
    started = datetime.now().isoformat(timespec="seconds")
    try:
        source_step = Path(str(component.get("sourceStep") or "")).expanduser()
        if not source_step.is_file():
            return {"ok": False, "error": f"Source STEP does not exist: {source_step}", "startedAt": started}
        output_path = Path(output_step) if output_step else None
        if output_path is None:
            output_root = Path(work_dir) if work_dir else Path.cwd()
            output_path = output_root / str(component.get("outputName") or "modified.step")
        output_path.parent.mkdir(parents=True, exist_ok=True)
        parameters = [row for row in component.get("parameters", []) if isinstance(row, dict)]
        defaults = {str(row.get("name") or ""): row.get("default", "") for row in parameters if row.get("name")}
        values = _input_values(inputs)
        model = StepModel.load(source_step)
        messages: list[str] = []
        for edit in component.get("edits", []):
            if not isinstance(edit, dict):
                continue
            operation = str(edit.get("operation") or "")
            args = [_resolve_arg(arg, values, defaults) for arg in list(edit.get("args") or [])]
            messages.append(_execute(model, operation, args))
        model.export_all(output_path)
        return {
            "ok": True,
            "startedAt": started,
            "finishedAt": datetime.now().isoformat(timespec="seconds"),
            "sourceStep": str(source_step),
            "outputStep": str(output_path),
            "messages": messages,
            "parameters": parameters,
        }
    except Exception as exc:
        return {
            "ok": False,
            "startedAt": started,
            "finishedAt": datetime.now().isoformat(timespec="seconds"),
            "error": str(exc),
            "traceback": traceback.format_exc(),
        }


def _load_inputs(path_or_json: str) -> object:
    if not path_or_json:
        return {}
    candidate = Path(path_or_json)
    if candidate.is_file():
        return json.loads(candidate.read_text(encoding="utf-8-sig"))
    return json.loads(path_or_json)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Run a STEP parametric component JSON.")
    parser.add_argument("component", help="Component JSON file.")
    parser.add_argument("--inputs", default="", help="JSON file or inline JSON object.")
    parser.add_argument("--output-step", default="", help="Output STEP path.")
    parser.add_argument("--work-dir", default="", help="Runtime work directory.")
    parsed = parser.parse_args(argv)
    try:
        component = json.loads(Path(parsed.component).read_text(encoding="utf-8-sig"))
        result = run_embedded_component(
            component,
            inputs=_load_inputs(parsed.inputs),
            output_step=parsed.output_step or None,
            work_dir=parsed.work_dir or None,
        )
    except Exception as exc:
        result = {"ok": False, "error": str(exc), "traceback": traceback.format_exc()}
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0 if result.get("ok") else 2


if __name__ == "__main__":
    raise SystemExit(main())