#!/usr/bin/env python3
"""Offline STEP -> triangle cache for Qt OpenGL. No robot/ROS/device operations.

Requires gmsh (OpenCascade). Runtime GUI needs only numpy and the generated NPZ.
"""
import argparse
import hashlib
import json
from pathlib import Path
import gmsh
import numpy as np


def convert(source, output):
    source=Path(source).resolve();output=Path(output).resolve()
    gmsh.initialize(['step-converter','-nopopup'])
    try:
        gmsh.option.setNumber('General.Terminal',0)
        gmsh.option.setString('Geometry.OCCTargetUnit','MM')
        gmsh.model.add('gripper')
        gmsh.model.occ.importShapes(str(source));gmsh.model.occ.synchronize()
        gmsh.option.setNumber('Mesh.MeshSizeMin',.8)
        gmsh.option.setNumber('Mesh.MeshSizeMax',2.)
        gmsh.option.setNumber('Mesh.MeshSizeFromCurvature',18)
        gmsh.option.setNumber('Mesh.ElementOrder',1)
        gmsh.model.mesh.generate(2)
        tags,coords,_=gmsh.model.mesh.getNodes()
        order=np.argsort(tags);sorted_tags=tags[order]
        nodes=np.asarray(coords).reshape(-1,3)[order]
        parts=[];colors=[];labels=[]
        surfaces=gmsh.model.getEntities(2)
        for part,(_,tag) in enumerate(surfaces):
            kinds,_,element_nodes=gmsh.model.mesh.getElements(2,tag)
            r,g,b,a=gmsh.model.getColor(2,tag)
            base=np.array([r,g,b],float)/255.
            if np.max(base)-np.min(base)<.03: base=np.array([.76,.79,.83])
            for kind,indices in zip(kinds,element_nodes):
                if kind!=2: raise ValueError('Only first order triangles are supported')
                indices=np.asarray(indices).reshape(-1,3)
                triangles=nodes[np.searchsorted(sorted_tags,indices)]
                parts.append(triangles/1000.)
                colors.append(np.tile(base,(len(triangles),1)))
                labels.append(np.full(len(triangles),part,np.int32))
        triangles=np.concatenate(parts).astype(np.float32)
        rgb=np.concatenate(colors).astype(np.float32)
        if not np.isfinite(triangles).all() or not len(triangles):raise ValueError('Invalid mesh')
        lo=triangles.min(axis=(0,1));hi=triangles.max(axis=(0,1))
        output.parent.mkdir(parents=True,exist_ok=True)
        np.savez_compressed(output,triangles=triangles,colors=rgb,surface=np.concatenate(labels))
        info=dict(source=source.name,source_sha256=hashlib.sha256(source.read_bytes()).hexdigest(),
                  units='m',source_units='mm',triangles=len(triangles),surfaces=len(surfaces),
                  solids=len(gmsh.model.getEntities(3)),bounds_m=[lo.tolist(),hi.tolist()],
                  dimensions_mm=((hi-lo)*1000).tolist(),articulated=False,
                  note='Original assembled STEP pose; no inferred joints or live finger animation.')
        output.with_suffix('.json').write_text(json.dumps(info,ensure_ascii=False,indent=2)+'\n')
        print(json.dumps(info,ensure_ascii=False))
    finally:gmsh.finalize()

if __name__=='__main__':
    parser=argparse.ArgumentParser();parser.add_argument('source');parser.add_argument('output')
    args=parser.parse_args();convert(args.source,args.output)