import ast
from llvmlite import ir
from pythonbpf.expr_pass import eval_expr
from enum import Enum
from .helper_utils import HelperHandlerRegistry, get_or_create_ptr_from_arg, get_flags_val, _handle_fstring_print, _simple_string_print


class BPFHelperID(Enum):
    BPF_MAP_LOOKUP_ELEM = 1
    BPF_MAP_UPDATE_ELEM = 2
    BPF_MAP_DELETE_ELEM = 3
    BPF_KTIME_GET_NS = 5
    BPF_PRINTK = 6
    BPF_GET_CURRENT_PID_TGID = 14
    BPF_PERF_EVENT_OUTPUT = 25


@HelperHandlerRegistry.register("ktime")
def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func,
                             local_sym_tab=None, struct_sym_tab=None,
                             local_var_metadata=None):
    """
    Emit LLVM IR for bpf_ktime_get_ns helper function call.
    """
    # func is an arg to just have a uniform signature with other emitters
    helper_id = ir.Constant(ir.IntType(64), BPFHelperID.BPF_KTIME_GET_NS.value)
    fn_type = ir.FunctionType(ir.IntType(64), [], var_arg=False)
    fn_ptr_type = ir.PointerType(fn_type)
    fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
    result = builder.call(fn_ptr, [], tail=False)
    return result, ir.IntType(64)


@HelperHandlerRegistry.register("lookup")
def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func,
                                local_sym_tab=None, struct_sym_tab=None,
                                local_var_metadata=None):
    """
    Emit LLVM IR for bpf_map_lookup_elem helper function call.
    """
    if not call.args or len(call.args) != 1:
        raise ValueError("Map lookup expects exactly one argument (key), got "
                         f"{len(call.args)}")
    key_ptr = get_or_create_ptr_from_arg(call.args[0], builder, local_sym_tab)
    map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())

    fn_type = ir.FunctionType(
        ir.PointerType(),  # Return type: void*
        [ir.PointerType(), ir.PointerType()],  # Args: (void*, void*)
        var_arg=False
    )
    fn_ptr_type = ir.PointerType(fn_type)

    fn_addr = ir.Constant(ir.IntType(
        64), BPFHelperID.BPF_MAP_LOOKUP_ELEM.value)
    fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)

    result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)

    return result, ir.PointerType()


@HelperHandlerRegistry.register("print")
def bpf_printk_emitter(call, map_ptr, module, builder, func,
                       local_sym_tab=None, struct_sym_tab=None,
                       local_var_metadata=None):
    """Emit LLVM IR for bpf_printk helper function call."""
    if not hasattr(func, "_fmt_counter"):
        func._fmt_counter = 0

    if not call.args:
        raise ValueError(
            "bpf_printk expects at least one argument (format string)")

    args = []
    if isinstance(call.args[0], ast.JoinedStr):
        args = _handle_fstring_print(call.args[0], module, builder, func,
                                     local_sym_tab, struct_sym_tab,
                                     local_var_metadata)
    elif isinstance(call.args[0], ast.Constant) and isinstance(call.args[0].value, str):
        # TODO: We are onbly supporting single arguments for now.
        # In case of multiple args, the first one will be taken.
        args = _simple_string_print(call.args[0], module, builder, func)
    else:
        raise NotImplementedError(
            "Only simple string literals or f-strings are supported in bpf_printk.")

    fn_type = ir.FunctionType(
        ir.IntType(64), [ir.PointerType(), ir.IntType(32)], var_arg=True)
    fn_ptr_type = ir.PointerType(fn_type)
    fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_PRINTK.value)
    fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)

    builder.call(fn_ptr, args, tail=True)
    return None


@HelperHandlerRegistry.register("update")
def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func,
                                local_sym_tab=None, struct_sym_tab=None,
                                local_var_metadata=None):
    """
    Emit LLVM IR for bpf_map_update_elem helper function call.
    Expected call signature: map.update(key, value, flags=0)
    """
    if (not call.args or
        len(call.args) < 2 or
            len(call.args) > 3):
        raise ValueError("Map update expects 2 or 3 arguments (key, value, flags), got "
                         f"{len(call.args)}")

    key_arg = call.args[0]
    value_arg = call.args[1]
    flags_arg = call.args[2] if len(call.args) > 2 else None

    key_ptr = get_or_create_ptr_from_arg(key_arg, builder, local_sym_tab)
    value_ptr = get_or_create_ptr_from_arg(value_arg, builder, local_sym_tab)
    flags_val = get_flags_val(flags_arg, builder, local_sym_tab)

    map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
    fn_type = ir.FunctionType(
        ir.IntType(64),
        [ir.PointerType(), ir.PointerType(), ir.PointerType(), ir.IntType(64)],
        var_arg=False
    )
    fn_ptr_type = ir.PointerType(fn_type)

    fn_addr = ir.Constant(ir.IntType(
        64), BPFHelperID.BPF_MAP_UPDATE_ELEM.value)
    fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)

    if isinstance(flags_val, int):
        flags_const = ir.Constant(ir.IntType(64), flags_val)
    else:
        flags_const = flags_val

    result = builder.call(
        fn_ptr, [map_void_ptr, key_ptr, value_ptr, flags_const], tail=False)

    return result, None


@HelperHandlerRegistry.register("delete")
def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func,
                                local_sym_tab=None, struct_sym_tab=None,
                                local_var_metadata=None):
    """
    Emit LLVM IR for bpf_map_delete_elem helper function call.
    Expected call signature: map.delete(key)
    """
    if not call.args or len(call.args) != 1:
        raise ValueError("Map delete expects exactly one argument (key), got "
                         f"{len(call.args)}")
    key_ptr = get_or_create_ptr_from_arg(call.args[0], builder, local_sym_tab)
    map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())

    # Define function type for bpf_map_delete_elem
    fn_type = ir.FunctionType(
        ir.IntType(64),  # Return type: int64 (status code)
        [ir.PointerType(), ir.PointerType()],  # Args: (void*, void*)
        var_arg=False
    )
    fn_ptr_type = ir.PointerType(fn_type)

    fn_addr = ir.Constant(ir.IntType(
        64), BPFHelperID.BPF_MAP_DELETE_ELEM.value)
    fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)

    result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)

    return result, None


@HelperHandlerRegistry.register("pid")
def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func,
                                     local_sym_tab=None, struct_sym_tab=None,
                                     local_var_metadata=None):
    """
    Emit LLVM IR for bpf_get_current_pid_tgid helper function call.
    """
    # func is an arg to just have a uniform signature with other emitters
    helper_id = ir.Constant(ir.IntType(
        64), BPFHelperID.BPF_GET_CURRENT_PID_TGID.value)
    fn_type = ir.FunctionType(ir.IntType(64), [], var_arg=False)
    fn_ptr_type = ir.PointerType(fn_type)
    fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
    result = builder.call(fn_ptr, [], tail=False)

    # Extract the lower 32 bits (PID) using bitwise AND with 0xFFFFFFFF
    mask = ir.Constant(ir.IntType(64), 0xFFFFFFFF)
    pid = builder.and_(result, mask)
    return pid, ir.IntType(64)


def bpf_perf_event_output_handler(call, map_ptr, module, builder, func,
                                  local_sym_tab=None, struct_sym_tab=None,
                                  local_var_metadata=None):
    if len(call.args) != 1:
        raise ValueError("Perf event output expects exactly one argument (data), got "
                         f"{len(call.args)}")
    data_arg = call.args[0]
    ctx_ptr = func.args[0]  # First argument to the function is ctx

    if isinstance(data_arg, ast.Name):
        data_name = data_arg.id
        if local_sym_tab and data_name in local_sym_tab:
            data_ptr = local_sym_tab[data_name][0]
        else:
            raise ValueError(
                f"Data variable {data_name} not found in local symbol table.")
        # Check is data_name is a struct
        if local_var_metadata and data_name in local_var_metadata:
            data_type = local_var_metadata[data_name]
            if data_type in struct_sym_tab:
                struct_info = struct_sym_tab[data_type]
                size_val = ir.Constant(ir.IntType(64), struct_info.size)
            else:
                raise ValueError(
                    f"Struct type {data_type} for variable {data_name} not found in struct symbol table.")
        else:
            raise ValueError(
                f"Metadata for variable {data_name} not found in local variable metadata.")

        # BPF_F_CURRENT_CPU is -1 in 32 bit
        flags_val = ir.Constant(ir.IntType(64), 0xFFFFFFFF)

        map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
        data_void_ptr = builder.bitcast(data_ptr, ir.PointerType())
        fn_type = ir.FunctionType(
            ir.IntType(64),
            [ir.PointerType(ir.IntType(8)), ir.PointerType(), ir.IntType(64),
             ir.PointerType(), ir.IntType(64)],
            var_arg=False
        )
        fn_ptr_type = ir.PointerType(fn_type)

        # helper id
        fn_addr = ir.Constant(ir.IntType(
            64), BPFHelperID.BPF_PERF_EVENT_OUTPUT.value)
        fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)

        result = builder.call(
            fn_ptr, [ctx_ptr, map_void_ptr, flags_val, data_void_ptr, size_val], tail=False)
        return result, None
    else:
        raise NotImplementedError(
            "Only simple object names are supported as data in perf event output.")


def handle_helper_call(call, module, builder, func,
                       local_sym_tab=None, map_sym_tab=None,
                       struct_sym_tab=None, local_var_metadata=None):
    print(local_var_metadata)
    if isinstance(call.func, ast.Name):
        func_name = call.func.id
        hdl_func = HelperHandlerRegistry.get_handler(func_name)
        if hdl_func:
            # it is not a map method call
            return hdl_func(call, None, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
        else:
            raise NotImplementedError(
                f"Function {func_name} is not implemented as a helper function.")
    elif isinstance(call.func, ast.Attribute):
        # likely a map method call
        if isinstance(call.func.value, ast.Call) and isinstance(call.func.value.func, ast.Name):
            map_name = call.func.value.func.id
            method_name = call.func.attr
            if map_sym_tab and map_name in map_sym_tab:
                map_ptr = map_sym_tab[map_name]
                hdl_func = HelperHandlerRegistry.get_handler(method_name)
                if hdl_func:
                    print(local_var_metadata)
                    return hdl_func(
                        call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
                else:
                    raise NotImplementedError(
                        f"Map method {method_name} is not implemented as a helper function.")
            else:
                raise ValueError(
                    f"Map variable {map_name} not found in symbol tables.")
        elif isinstance(call.func.value, ast.Name):
            obj_name = call.func.value.id
            method_name = call.func.attr
            if map_sym_tab and obj_name in map_sym_tab:
                map_ptr = map_sym_tab[obj_name]
                hdl_func = HelperHandlerRegistry.get_handler(method_name)
                if hdl_func:
                    return hdl_func(
                        call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
                else:
                    raise NotImplementedError(
                        f"Map method {method_name} is not implemented as a helper function.")
            else:
                raise ValueError(
                    f"Map variable {obj_name} not found in symbol tables.")
        else:
            raise NotImplementedError(
                "Attribute not supported for map method calls.")
    return None