from __future__ import annotations
import argparse, json, time
from pathlib import Path
import tensorrt as trt
import torch

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

COMPONENTS = {
    "ss-flow-solo": ("onnx/anigen/ss-flow-solo/model.onnx", "engines/l4-sm89/anigen-ss-flow-solo.plan"),
    "ss-decoder": ("onnx/anigen/ss-decoder/model.onnx", "engines/l4-sm89/anigen-ss-decoder.plan"),
    "slat-flow-core": ("onnx/anigen/slat-flow-core/model.onnx", "engines/l4-sm89/anigen-slat-flow-core.plan"),
    "slat-decoder-core": ("onnx/anigen/slat-decoder-core/model.onnx", "engines/l4-sm89/anigen-slat-decoder-core.plan"),
    "skin-decoder": ("onnx/anigen/skin-decoder/model.onnx", "engines/l4-sm89/anigen-skin-decoder.plan"),
}

PROFILES = {
    "slat-flow-core": {
        "geo": ((1,128,1024),(1,4096,1024),(1,16384,1024)),
        "skin": ((1,128,512),(1,4096,512),(1,16384,512)),
        "skl": ((1,16,512),(1,1024,512),(1,8192,512)),
        "mod_geo": ((1,1024),(1,1024),(1,1024)),
        "mod_skin": ((1,512),(1,512),(1,512)),
        "mod_skl": ((1,512),(1,512),(1,512)),
        "cond": ((1,1374,1024),(1,1374,1024),(1,1374,1024)),
    },
    "skin-decoder": {
        "vertex_features": ((1,128,4),(1,25000,4),(1,300000,4)),
        "joint_features": ((1,2,4),(1,24,4),(1,64,4)),
        "parents": ((1,2),(1,24),(1,64)),
    },
}


def build(name: str, src: Path, dst: Path, workspace_gib: int, opt_level: int):
    builder = trt.Builder(LOGGER)
    flags = 0
    if hasattr(trt.NetworkDefinitionCreationFlag, "STRONGLY_TYPED"):
        flags |= 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
    network = builder.create_network(flags)
    parser = trt.OnnxParser(network, LOGGER)
    if not parser.parse_from_file(str(src)):
        raise RuntimeError("\n".join(str(parser.get_error(i)) for i in range(parser.num_errors)))
    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_gib * (1 << 30))
    if hasattr(config, "builder_optimization_level"):
        config.builder_optimization_level = opt_level
    if name in PROFILES:
        profile = builder.create_optimization_profile()
        by_name = PROFILES[name]
        for i in range(network.num_inputs):
            tensor = network.get_input(i)
            if tensor.name in by_name:
                mn, op, mx = by_name[tensor.name]
                profile.set_shape(tensor.name, mn, op, mx)
        config.add_optimization_profile(profile)
    started = time.time()
    blob = builder.build_serialized_network(network, config)
    if blob is None:
        raise RuntimeError("TensorRT build returned None")
    dst.parent.mkdir(parents=True, exist_ok=True)
    dst.write_bytes(blob)
    return {
        "seconds": round(time.time() - started, 3),
        "bytes": dst.stat().st_size,
        "inputs": {network.get_input(i).name: list(network.get_input(i).shape) for i in range(network.num_inputs)},
        "outputs": {network.get_output(i).name: list(network.get_output(i).shape) for i in range(network.num_outputs)},
        "profile": PROFILES.get(name),
    }


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--root", default="/tmp/cf-anigen-onnx"); ap.add_argument("--component", choices=list(COMPONENTS)+["all"], default="all"); ap.add_argument("--workspace-gib", type=int, default=6); ap.add_argument("--opt-level", type=int, default=4); a=ap.parse_args()
    root=Path(a.root); names=list(COMPONENTS) if a.component=='all' else [a.component]; result={}
    for name in names:
        src_rel,dst_rel=COMPONENTS[name]; src,dst=root/src_rel,root/dst_rel
        if not src.exists(): print(f"SKIP {name}: {src} missing",flush=True); continue
        print(f"BUILD {name}: {src} -> {dst}",flush=True); result[name]=build(name,src,dst,a.workspace_gib,a.opt_level)
    result['environment']={'gpu':torch.cuda.get_device_name(0),'cc':list(torch.cuda.get_device_capability(0)),'cuda':torch.version.cuda,'tensorrt':trt.__version__}
    meta=root/f'trt_build_anigen_{a.component}_meta.json'; meta.write_text(json.dumps(result,indent=2)); print(json.dumps(result,indent=2),flush=True)

if __name__=='__main__': main()