import json
import os
import sys
from pathlib import Path

import tensorrt as trt
from huggingface_hub import snapshot_download

ROOT = Path(snapshot_download(
    os.environ.get('CF_RUNTIME_REPO','patdev/Companion-Forge-L4-ONNX'),
    repo_type='model', token=os.environ.get('HF_TOKEN'), local_dir='/tmp/cf-custom',
    allow_patterns=['onnx/custom-ops/*.onnx','plugins/tensorrt/companion_sparse_trt.py'],
))
sys.path.insert(0, str(ROOT/'plugins/tensorrt'))
import companion_sparse_trt  # noqa: F401,E402

LOGGER = trt.Logger(trt.Logger.INFO)

PROFILES = {
    'SparseSubdivide': {'feats': (32,8), 'coords': (32,4)},
    'SparseDownsample': {'feats': (64,8), 'coords': (64,4)},
    'SparseUpsample': {'feats': (8,8), 'target_coords': (64,4), 'inverse': (64,)},
    'SparseWindowAttention': {'qkv': (64,3,4,16), 'coords': (64,4)},
    'SparseConv3D': {'feats': (64,8), 'coords': (64,4), 'weight': (16,3,3,3,8), 'bias': (16,)},
    'MeshTopologyExtract': {
        'verts_grid': (274625,3), 'sdf': (274625,), 'cube_idx': (262144,8),
        'beta': (262144,12), 'alpha': (262144,8), 'gamma': (262144,), 'colors_grid': (274625,6),
    },
}


def build(name: str):
    onnx_path = ROOT/'onnx/custom-ops'/f'{name}.onnx'
    builder = trt.Builder(LOGGER)
    network = builder.create_network(0)
    parser = trt.OnnxParser(network, LOGGER)
    if not parser.parse(onnx_path.read_bytes()):
        errs = [str(parser.get_error(i)) for i in range(parser.num_errors)]
        raise RuntimeError(f'{name} parse failed:\n' + '\n'.join(errs))
    cfg = builder.create_builder_config()
    cfg.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 << 30)
    prof = builder.create_optimization_profile()
    input_report = {}
    dynamic = False
    for i in range(network.num_inputs):
        inp = network.get_input(i)
        shape = PROFILES[name].get(inp.name)
        if shape is None:
            raise KeyError(f'{name}: missing profile for {inp.name}')
        input_report[inp.name] = tuple(inp.shape)
        if any(int(d) < 0 for d in inp.shape):
            dynamic = True
            prof.set_shape(inp.name, shape, shape, shape)
    if dynamic:
        cfg.add_optimization_profile(prof)
    blob = builder.build_serialized_network(network, cfg)
    if blob is None:
        raise RuntimeError(f'{name} build_serialized_network returned None')
    out = Path('/tmp/custom-engines')/f'{name}.plan'
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_bytes(bytes(blob))
    return {
        'op': name,
        'onnx_bytes': onnx_path.stat().st_size,
        'plan_bytes': out.stat().st_size,
        'inputs': input_report,
        'outputs': [network.get_output(i).name for i in range(network.num_outputs)],
    }


if __name__ == '__main__':
    report=[]
    for name in PROFILES:
        print(f'=== BUILD {name} ===', flush=True)
        r=build(name); print('OK',json.dumps(r),flush=True); report.append(r)
    Path('/tmp/custom-engines/report.json').write_text(json.dumps(report,indent=2))