# -*- coding: utf-8 -*-
"""机器人手臂深度圆角优化版 V2：彻底消除所有棱角
针对用户反馈的三大问题区域进行深度优化：
1. 蓝色零件（P04肩叉）顶部方形突起 → 全圆角处理
2. 绿色/青色零件（P05/P08）底部和侧面棱角 → Loft外轮廓圆角
3. 所有连接处的尖锐过渡 → 大圆角平滑过渡
4. 橙色零件（P06/P07）所有方形边缘 → 完整圆角系统
"""
import math
import os
import sys
import time
import cadquery as cq
from arm.cad import arm_build as AB
from arm.paths import OUT, WEB, DOCS
from arm.constants import M, C
from arm.cad.helpers import cyl, box, ring, place, trs, bb, bbox_str, solids_of
from arm.cad.arm_build import (Joint, add, bx, Ycyl, Ybox, Yplate, prism, dogbone, sector,
                       FLANGE_Z, COVER_Z, HUB_FACE_Z, CLAMP_SPAN, HUB_M2_CLOCK,
                       BORE_R, CB_R, PCB_W, D2R, T, R, PARTS)
from arm import kinematics as HL
from OCP.BRepAlgoAPI import BRepAlgoAPI_Common
from arm.kinematics import (HUB, PLATE_T, R_MOT, Z_TIP, Z_WR, Z_EI, Z_SH, Z_RL, Z_R1,
                          Z_R3, Z_HUB3, Z_EL, Z_R6, Z_HUB6,
                          Z_R4, Z_CAP, Z_M6, Z_CHEST0, Z_CHEST1, X_J1_HUB, X_ARM, L1,
                          X_FORK_GAP, X_BAY_EXT,
                          Y_AXIS, Y_HUB, Y_REAR, Y_COVER, Y_H, Y_R_IN, Y_R_OUT,
                          Y_M6, Y_CHEST, X_CHEST, FOOT_H)

VERSION = os.environ.get("ARMVER", "v2")
ITER = int(os.environ.get("ARMITER", "3"))
AB.VERSION = VERSION
AB.ITER = ITER

# ------------------------------------------------ 逐件模块（几何已拆出）
from arm.cad.parts import PART_MODULES
from arm.cad.part_tools import (
    BOSS, CHAMFER_EDGE, CLAMP_BACKOFF, CLAMP_BCD, CLAMP_CLOCK, CLAMP_CLOCK_LOCAL, CLAMP_CLR,
    CLAMP_ROOT, CLAMP_T, CLEARANCE_GAP, DISC_R, FENCE_OPEN_CUT, FILLET_LARGE, FILLET_MEDIUM,
    FILLET_SMALL, FILLET_TINY, FILLET_XLARGE, HOLLOW_FILLET, HOLLOW_WALL, SLEEVE_R, XB, Xcyl,
    Xplate, Y_MID, cut_motor_fence, hollow_channel, hollow_tent_window, hub_m2_holes,
    make_direct_hub, make_pitch_to_roll_bracket, motor_front_clamp,
)
from arm.cad.parts.p01_column import X_COL, Y_COL, Z_LOCK, p_chest, split_p01_poka_yoke
from arm.cad.parts.p02_base import make_p02_poka_yoke, p_base_foot
from arm.cad.parts.p03_girdle import p_girdle
from arm.cad.parts.p04_fork import P04_CH, p_fork
from arm.cad.parts.p05_roll_base import p_shoulder_roll_base
from arm.cad.parts.p06_humerus import P06_HOLLOW_WIN, P06_WIN_XZ, p_humerus_bone
from arm.cad.parts.p07_forearm import P07_CH, p_forearm_bone
from arm.cad.parts.p08_wrist import p_wrist_bracket
from arm.cad.parts.p09_flange import p_tool_flange
# 兼容旧工具：scratch/ 的诊断脚本按 arm_model.<符号> 取用件构造函数与开关

def build():
    print("=" * 80, flush=True)
    print("【第一阶段 1/5】开始生成手臂各级构件 3D 几何实体模型...", flush=True)
    print("=" * 80, flush=True)
    t_stage1 = time.time()

    print("  [0/9] 载入电机标准件模组 (STEP)...", end="", flush=True)
    t0 = time.time()
    motor = cq.importers.importStep(AB.SRC_STEP).val()
    motor = place(motor, T(0, 0, -2.57))
    AB.MOTOR, AB.HALF, AB.CLAMP = motor, None, None
    print(f" 完成 ({time.time()-t0:.2f}s) ✓", flush=True)

    JS = {}
    for tag, cn, rear, d, xd, holder, driven in HL.JOINTS:
        JS[tag] = Joint(tag, cn, rear, d, xd, lug_phase=90.0 if tag == "J1" else 0.0)
    j1, j2, j3, j4, j5, j6 = (JS["J%d" % k] for k in range(1, 7))

    PARTS[:] = []
    AB.CENTS.clear()
    AB.GRP = {}
    par = {"J1": "root", "J2": "J1", "J3": "J2", "J4": "J3", "J5": "J4", "J6": "J5"}
    motor_of = {"J1": "M1_J1", "J2": "M2_J2", "J3": "M3_J3",
                "J4": "M4_J4", "J5": "M5_J5", "J6": "M6_J6"}
    for t in ("J1", "J2", "J3", "J4", "J5", "J6"):
        AB.GRP[motor_of[t]] = par[t]
    AB.JAW = {}
    # 逐件模块：几何、件号/配色/说明、装配帧归属都在 arm/cad/parts/ 里，
    # 这里只按 BOM 顺序收集 specs()，并把件号登记到装配帧图上。
    parts_to_build = []
    for _m in PART_MODULES:
        for _spec in _m.specs(JS):
            AB.GRP[_spec[0]] = _m.JOINT
            parts_to_build.append(_spec)


    total_parts = len(parts_to_build)
    for idx, (name, cn, fn, color, note) in enumerate(parts_to_build, 1):
        print(f"  [{idx}/{total_parts}] 正在构建 {name} ({cn})...", end="", flush=True)
        t0 = time.time()
        sh = fn()
        dt = time.time() - t0
        vol = AB.vol_of(sh) / 1000.0
        n_solids = len(solids_of(sh))
        print(f" 完成 ✓  ({dt:.2f}s | 体积: {vol:.1f} cm³ | 实体数: {n_solids})", flush=True)
        add(name, cn, sh, color, note=note)

    print("  [装配] 正在将 M1~M6 电机定位至各关节坐标位姿...", end="", flush=True)
    t0 = time.time()
    for k, (tag, cn, rear, d, xd, holder, driven_name) in enumerate(HL.JOINTS):
        j = JS[tag]
        add("M%d_%s" % (k + 1, tag), "%s 电机模组" % tag, place(motor, j.t), 0x9ca3af,
            printed=False, note="复用 电机.STEP")
    print(f" 完成 ✓ ({time.time()-t0:.2f}s)", flush=True)
    print(f"-> 构件模型全部生成完成，总耗时: {time.time()-t_stage1:.2f}s\n", flush=True)
    return JS


CHAIN = ["root", "J1", "J2", "J3", "J4", "J5", "J6"]
SEP = 74.0


def explode_chain(p, JS, cents, gmean):
    """爆炸图偏移：立柱垂直阶梯拉开，底座前后防呆拉开，电机沿轴拉出，机械臂整体随上段分离"""
    nm = p["name"]
    if nm == "P01_A_立柱下段_三角形30%_墙4底5_竖打":
        return [0.0, 0.0, 0.0]
    if nm == "P01_B_立柱中段_三角形30%_墙4底5_竖打":
        return [0.0, 0.0, 60.0]
    if nm == "P01_C_立柱上段_三角形30%_墙4底5_竖打":
        return [0.0, 0.0, 120.0]
    if nm == "P02_Center_底座中核心_螺旋体30%_墙5底6_平打":
        return [0.0, 0.0, 0.0]
    if nm == "P02_Front_前向防翻力臂_螺旋体30%_墙4底5_平打":
        return [65.0, 0.0, 0.0]
    if nm == "P02_Rear_后向防翻力臂_螺旋体30%_墙4底5_平打":
        return [-65.0, 0.0, 0.0]
    if nm == "P03_肩吊臂_螺旋体35%_墙5底5_竖打":
        return [0.0, 0.0, 120.0]

    grp = AB.GRP.get(p["name"], "root")
    if grp == "root":
        return [0.0, 0.0, 0.0]
    k = CHAIN.index(grp)
    off = [0.0, 0.0, 120.0 - SEP * k]
    j = JS.get(grp)
    if j:
        w = 0.0
        if "半联轴器" in nm or "抱箍" in nm:
            w = 34.0 * (1.0 if ("A" in nm[-1:]) else -1.0)
        elif nm.startswith("M"):
            w = 62.0
        for q in range(3):
            off[q] += j.d[q] * w
    return off


if __name__ == "__main__":
    JS = build()

    print("=" * 80, flush=True)
    print("【单实体拓扑完整性校验】验证所有打印构件流形实体状态...", flush=True)
    print("=" * 80, flush=True)
    bad = 0
    printed_list = [p for p in PARTS if p["printed"]]
    total_pr = len(printed_list)
    for idx, p in enumerate(printed_list, 1):
        n = len(solids_of(p["shape"]))
        if n != 1:
            bad += 1
        status = "单实体 ✓" if n == 1 else f"异常 ({n} 个实体) ❌"
        print(f"  [{idx}/{total_pr}] {p['name']:<14} : 体积 {p['vol']/1000.0:6.1f} cm³ | {status}", flush=True)
    if bad == 0:
        print(f"-> 单实体检查全部通过！{total_pr}/{total_pr} 构件均为完美单一流形拓扑 ✓\n", flush=True)
    else:
        print(f"-> 发现 {bad} 个非单实体构件！\n", flush=True)

    AB.explode_vec = explode_chain
    bom = AB.export(JS)

    # 紧固件（螺钉）实例注入：把 docs/screws.json 的 72 颗螺钉按真实孔位还原成三维实例，
    # 追加进 web/mesh_arm.json 的 parts（查看器可点 BOM 螺丝行高亮定位；纯数据，不动 CAD）
    from arm.dataio import screw_mesh
    screw_mesh.main()
    cl = AB.collide()
    from arm.viewer.build_viewer import rebuild_viewer
    rebuild_viewer()
    # 可达包络自动注入：AB.export 刚重写过 mesh_arm.json（包络件会被冲掉），
    # 这里重跑 FK 采样 + 包络网格注入，再重打包一次查看器（约 3~4s，纯数据无 CAD 计算）
    from arm.analysis import ws_sweep, gen_env_mesh
    t_env = time.time()
    ws_sweep.main()
    gen_env_mesh.main()
    rebuild_viewer()
    print(f"-> 可达包络已自动注入查看器（+{time.time() - t_env:.1f}s）✓\n", flush=True)
    print("=" * 80, flush=True)
    print("✅ 机器人手臂编译、导出与查看器发布全流程完成！", flush=True)
    print("=" * 80, flush=True)

    # OCC/cadquery 的静态析构在解释器正常退出时必定段错误（实测：光是
    # `python -c "import arm.cad.helpers"` 就返回 139，与本项目代码无关）。所以入口
    # 脚本跑完后显式刷缓冲再 os._exit —— 否则全流程已经成功、exit code 却是 139,
    # 会把上游脚本的 set -e 判成失败。只在 __main__ 出口生效, 被 import 时不触发。
    sys.stdout.flush()
    os._exit(0)