import tensorrt_rtx as trt
import argparse
import os
import onnx
from pathlib import Path
from tqdm import tqdm

# Path for your persistent hardware profile
MASTER_CACHE_FILE = "..\\TRT_RTX_ENGINE.cache"

def get_timing_cache(config):
    if os.path.exists(MASTER_CACHE_FILE):
        with open(MASTER_CACHE_FILE, "rb") as f:
            cache_data = f.read()
            return config.create_timing_cache(cache_data)
    return config.create_timing_cache(b"")

def save_timing_cache(timing_cache):
    with open(MASTER_CACHE_FILE, "wb") as f:
        f.write(timing_cache.serialize())


def get_onnx_precision(onnx_path):
    """Detects if the ONNX model is FP32 or FP16."""
    model = onnx.load(onnx_path)
    # Check the data type of the first input
    dtype = model.graph.input[0].type.tensor_type.elem_type
    if dtype == onnx.TensorProto.FLOAT16:
        return "fp16"
    return "fp32"

def build_engines(input_dir, output_dir):
    logger  = trt.Logger(trt.Logger.ERROR)
    builder = trt.Builder(logger)
    
    # Limit workspace (e.g., to 2GB or 4GB depending on your GPU)
    # 1 << 30 is 1GB
    vram_limit  = 4 * (1 << 30)    
    path        = Path(input_dir)
    trt_version = trt.__version__ 
    
    # Grab all .onnx files except 'rife'
    onnx_files = list(f for f in path.glob("*.onnx") if "rife" not in f.name.lower())
    
    # Progress Bar configuration
    pbar = tqdm(onnx_files, desc="Converting Models", unit="model")
    
    if not onnx_files:
        print(f"No ONNX files found in {input_dir}")
        return
    
    for onnx_path in pbar:
        network   = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
        parser    = trt.OnnxParser(network, logger)
        config    = builder.create_builder_config()        
        precision = get_onnx_precision(onnx_path)
        
        
        # Enforce tactic limits and FP16
        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, vram_limit)
        config.builder_optimization_level = 5
        
        if builder.platform_has_fast_fp16:
            precision = "fp16"
            config.set_flag(trt.BuilderFlag.FP16)
        
        
        # Use Master Cache
        timing_cache = get_timing_cache(config)
        config.set_timing_cache(timing_cache, False)    
        
        # Parse ONNX
        with open(onnx_path, 'rb') as model:
            if not parser.parse(model.read()):
                print(f"Error parsing {onnx_path}")
                continue            
        
        # Dynamic Shape Profile
        profile    = builder.create_optimization_profile()
        input_name = network.get_input(0).name
        min_shape  = (1, 3, 256, 256)
        opt_shape  = (1, 3, 512, 512)
        max_shape  = (1, 3, 1280, 1280)        
        
        profile.set_shape(input_name, min_shape, opt_shape, max_shape)
        config.add_optimization_profile(profile)
        
        # Filename: ONNXFILE_MIN_OPT_MAX_VERSION.trt
        def fmt(s): return "x".join(map(str, s))
        base_name   = onnx_path.stem
        engine_name = f"{base_name}_{fmt(min_shape)}_{fmt(opt_shape)}_{fmt(max_shape)}_{trt_version}.trt"
        engine_path = f"{output_dir}\\{engine_name}"  
        
        # Build serialized engine
        print(f"Building {engine_path}...")
        serialized_engine = builder.build_serialized_network(network, config)
        
        # Update Master Cache with new findings
        save_timing_cache(timing_cache)
        print(f"Updated Master Cache.")
            
        if serialized_engine:
            with open(engine_path, "wb") as f:
                f.write(serialized_engine)
            print(f"Saved {engine_path} and updated Master Cache.")
        else:
            print(f"FAILED TO SAVE {engine_path}...")    

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("onnx_dir", help="Directory containing .onnx files")
    parser.add_argument("trt_dir", help="Directory to put .trt files")
    args = parser.parse_args()
    build_engines(args.onnx_dir, args.trt_dir)