# c_features.py
import ctypes
import json
import os

lib = ctypes.CDLL("./cpplib/libsuri_features.so")  # path to so lib

# arguments and return types
lib.extract_flow_features_json.argtypes = [ctypes.c_char_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_char_p)]
lib.extract_flow_features_json.restype = ctypes.c_int

lib.tag_anomalies_json.argtypes = [ctypes.c_char_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_char_p)]
lib.tag_anomalies_json.restype = ctypes.c_int

lib.free_allocated.argtypes = [ctypes.c_void_p]
lib.free_allocated.restype = None


def call_c_function(func, events, required_cols=None):
    in_str = json.dumps(events).encode()
    col_str = json.dumps(required_cols).encode() if required_cols else None
    out_ptr = ctypes.c_char_p()

    code = func(in_str, col_str, ctypes.byref(out_ptr))
    if code != 0:
        raise RuntimeError(f"C++ feature extraction failed with code {code}")

    result = json.loads(out_ptr.value.decode())
    lib.free_allocated(out_ptr)
    return result

def extract_features(event_type: str, events, required_cols=None):
    funcs = {
        "flow": lib.extract_flow_features_json,
        "http": lib.extract_http_features_json,
        "dns": lib.extract_dns_features_json,
        "tls": lib.extract_tls_features_json
    }
    return call_c_function(funcs[event_type], events, required_cols)

def tag_anomalies(events, predictions):
    events_str = json.dumps(events).encode()
    preds_str = json.dumps(predictions).encode()
    out_ptr = ctypes.c_char_p()
    code = lib.tag_anomalies_json(events_str, preds_str, ctypes.byref(out_ptr))
    if code != 0:
        print(f"[C++ wrapper] Error - C function return code- {code}")
        print("Input data:")
        print(json.dumps(events[:2], indent=2))  # we display the first 2 events
    

        raise RuntimeError(f"C++ tagging failed with code {code}")
    result = json.loads(out_ptr.value.decode())
    lib.free_allocated(out_ptr)
    return result
    
