import ast
from llvmlite import ir
from .expr_pass import eval_expr


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), 5)
    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)


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 call.args and len(call.args) != 1:
        raise ValueError(
            "Map lookup expects exactly one argument, got " f"{len(call.args)}"
        )
    key_arg = call.args[0]
    if isinstance(key_arg, ast.Name):
        key_name = key_arg.id
        if local_sym_tab and key_name in local_sym_tab:
            key_ptr = local_sym_tab[key_name][0]
        else:
            raise ValueError(
                f"Key variable {key_name} not found in local symbol table."
            )
    elif isinstance(key_arg, ast.Constant) and isinstance(key_arg.value, int):
        # handle constant integer keys
        key_val = key_arg.value
        key_type = ir.IntType(64)
        key_ptr = builder.alloca(key_type)
        key_ptr.align = key_type // 8
        builder.store(ir.Constant(key_type, key_val), key_ptr)
    else:
        raise NotImplementedError(
            "Only simple variable names are supported as keys in map lookup."
        )

    if key_ptr is None:
        raise ValueError("Key pointer is None.")

    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)

    # Helper ID 1 is bpf_map_lookup_elem
    fn_addr = ir.Constant(ir.IntType(64), 1)
    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()


def bpf_printk_emitter(
    call,
    map_ptr,
    module,
    builder,
    func,
    local_sym_tab=None,
    struct_sym_tab=None,
    local_var_metadata=None,
):
    if not hasattr(func, "_fmt_counter"):
        func._fmt_counter = 0

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

    if isinstance(call.args[0], ast.JoinedStr):
        fmt_parts = []
        exprs = []

        for value in call.args[0].values:
            print("Value in f-string:", ast.dump(value))
            if isinstance(value, ast.Constant):
                if isinstance(value.value, str):
                    fmt_parts.append(value.value)
                elif isinstance(value.value, int):
                    fmt_parts.append("%lld")
                    exprs.append(ir.Constant(ir.IntType(64), value.value))
                else:
                    raise NotImplementedError(
                        "Only string and integer constants are supported in f-string."
                    )
            elif isinstance(value, ast.FormattedValue):
                print("Formatted value:", ast.dump(value))
                # TODO: Dirty handling here, only checks for int or str
                if isinstance(value.value, ast.Name):
                    if local_sym_tab and value.value.id in local_sym_tab:
                        var_ptr, var_type = local_sym_tab[value.value.id]
                        if isinstance(var_type, ir.IntType):
                            fmt_parts.append("%lld")
                            exprs.append(value.value)
                        elif var_type == ir.PointerType(ir.IntType(8)):
                            # Case with string
                            fmt_parts.append("%s")
                            exprs.append(value.value)
                        else:
                            raise NotImplementedError(
                                "Only integer and pointer types are supported in formatted values."
                            )
                    else:
                        raise ValueError(
                            f"Variable {value.value.id} not found in local symbol table."
                        )
                elif isinstance(value.value, ast.Attribute):
                    # object field access from struct
                    if (
                        isinstance(value.value.value, ast.Name)
                        and local_sym_tab
                        and value.value.value.id in local_sym_tab
                    ):
                        var_name = value.value.value.id
                        field_name = value.value.attr
                        if local_var_metadata and var_name in local_var_metadata:
                            var_type = local_var_metadata[var_name]
                            if var_type in struct_sym_tab:
                                struct_info = struct_sym_tab[var_type]
                                if field_name in struct_info.fields:
                                    field_type = struct_info.field_type(field_name)
                                    if isinstance(field_type, ir.IntType):
                                        fmt_parts.append("%lld")
                                        exprs.append(value.value)
                                    elif field_type == ir.PointerType(ir.IntType(8)):
                                        fmt_parts.append("%s")
                                        exprs.append(value.value)
                                    else:
                                        raise NotImplementedError(
                                            "Only integer and pointer types are supported in formatted values."
                                        )
                                else:
                                    raise ValueError(
                                        f"Field {field_name} not found in struct {var_type}."
                                    )
                            else:
                                raise ValueError(
                                    f"Struct type {var_type} for variable {var_name} not found in struct symbol table."
                                )
                        else:
                            raise ValueError(
                                f"Metadata for variable {var_name} not found in local variable metadata."
                            )
                    else:
                        raise ValueError(
                            f"Variable {value.value.value.id} not found in local symbol table."
                        )
                else:
                    raise NotImplementedError(
                        "Only simple variable names are supported in formatted values."
                    )
            else:
                raise NotImplementedError("Unsupported value type in f-string.")

        fmt_str = "".join(fmt_parts) + "\n" + "\0"
        fmt_name = f"{func.name}____fmt{func._fmt_counter}"
        func._fmt_counter += 1

        fmt_gvar = ir.GlobalVariable(
            module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=fmt_name
        )
        fmt_gvar.global_constant = True
        fmt_gvar.initializer = ir.Constant(  # type: ignore
            ir.ArrayType(ir.IntType(8), len(fmt_str)), bytearray(fmt_str.encode("utf8"))
        )
        fmt_gvar.linkage = "internal"
        fmt_gvar.align = 1  # type: ignore

        fmt_ptr = builder.bitcast(fmt_gvar, ir.PointerType())

        args = [fmt_ptr, ir.Constant(ir.IntType(32), len(fmt_str))]

        # Only 3 args supported in bpf_printk
        if len(exprs) > 3:
            print(
                "Warning: bpf_printk supports up to 3 arguments, extra arguments will be ignored."
            )

        for expr in exprs[:3]:
            print(f"{ast.dump(expr)}")
            val, _ = eval_expr(
                func,
                module,
                builder,
                expr,
                local_sym_tab,
                None,
                struct_sym_tab,
                local_var_metadata,
            )
            if val:
                if isinstance(val.type, ir.PointerType):
                    val = builder.ptrtoint(val, ir.IntType(64))
                elif isinstance(val.type, ir.IntType):
                    if val.type.width < 64:
                        val = builder.sext(val, ir.IntType(64))
                else:
                    print(
                        "Warning: Only integer and pointer types are supported in bpf_printk arguments. Others will be converted to 0."
                    )
                    val = ir.Constant(ir.IntType(64), 0)
                args.append(val)
            else:
                print(
                    "Warning: Failed to evaluate expression for bpf_printk argument. It will be converted to 0."
                )
                args.append(ir.Constant(ir.IntType(64), 0))
        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), 6)
        fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
        return builder.call(fn_ptr, args, tail=True)

    for arg in call.args:
        if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
            fmt_str = arg.value + "\n" + "\0"
            fmt_name = f"{func.name}____fmt{func._fmt_counter}"
            func._fmt_counter += 1

            fmt_gvar = ir.GlobalVariable(
                module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=fmt_name
            )
            fmt_gvar.global_constant = True
            fmt_gvar.initializer = ir.Constant(  # type: ignore
                ir.ArrayType(ir.IntType(8), len(fmt_str)),
                bytearray(fmt_str.encode("utf8")),
            )
            fmt_gvar.linkage = "internal"
            fmt_gvar.align = 1  # type: ignore

            fmt_ptr = builder.bitcast(fmt_gvar, ir.PointerType())

            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), 6)
            fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)

            builder.call(
                fn_ptr, [fmt_ptr, ir.Constant(ir.IntType(32), len(fmt_str))], tail=True
            )
    return None


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

    # Handle key
    if isinstance(key_arg, ast.Name):
        key_name = key_arg.id
        if local_sym_tab and key_name in local_sym_tab:
            key_ptr = local_sym_tab[key_name][0]
        else:
            raise ValueError(
                f"Key variable {key_name} not found in local symbol table."
            )
    elif isinstance(key_arg, ast.Constant) and isinstance(key_arg.value, int):
        # Handle constant integer keys
        key_val = key_arg.value
        key_type = ir.IntType(64)
        key_ptr = builder.alloca(key_type)
        key_ptr.align = key_type.width // 8
        builder.store(ir.Constant(key_type, key_val), key_ptr)
    else:
        raise NotImplementedError(
            "Only simple variable names and integer constants are supported as keys in map update."
        )

    # Handle value
    if isinstance(value_arg, ast.Name):
        value_name = value_arg.id
        if local_sym_tab and value_name in local_sym_tab:
            value_ptr = local_sym_tab[value_name][0]
        else:
            raise ValueError(
                f"Value variable {value_name} not found in local symbol table."
            )
    elif isinstance(value_arg, ast.Constant) and isinstance(value_arg.value, int):
        # Handle constant integers
        value_val = value_arg.value
        value_type = ir.IntType(64)
        value_ptr = builder.alloca(value_type)
        value_ptr.align = value_type.width // 8
        builder.store(ir.Constant(value_type, value_val), value_ptr)
    else:
        raise NotImplementedError(
            "Only simple variable names and integer constants are supported as values in map update."
        )

    # Handle flags argument (defaults to 0)
    if flags_arg is not None:
        if isinstance(flags_arg, ast.Constant) and isinstance(flags_arg.value, int):
            flags_val = flags_arg.value
        elif isinstance(flags_arg, ast.Name):
            flags_name = flags_arg.id
            if local_sym_tab and flags_name in local_sym_tab:
                # Assume it's a stored integer value, load it
                flags_ptr = local_sym_tab[flags_name][0]
                flags_val = builder.load(flags_ptr)
            else:
                raise ValueError(
                    f"Flags variable {flags_name} not found in local symbol table."
                )
        else:
            raise NotImplementedError(
                "Only integer constants and simple variable names are supported as flags in map update."
            )
    else:
        flags_val = 0

    if key_ptr is None or value_ptr is None:
        raise ValueError("Key pointer or value pointer is None.")

    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)

    # helper id
    fn_addr = ir.Constant(ir.IntType(64), 2)
    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


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)
    """
    # Check for correct number of arguments
    if not call.args or len(call.args) != 1:
        raise ValueError(
            "Map delete expects exactly 1 argument (key), got " f"{len(call.args)}"
        )

    key_arg = call.args[0]

    # Handle key argument
    if isinstance(key_arg, ast.Name):
        key_name = key_arg.id
        if local_sym_tab and key_name in local_sym_tab:
            key_ptr = local_sym_tab[key_name][0]
        else:
            raise ValueError(
                f"Key variable {key_name} not found in local symbol table."
            )
    elif isinstance(key_arg, ast.Constant) and isinstance(key_arg.value, int):
        # Handle constant integer keys
        key_val = key_arg.value
        key_type = ir.IntType(64)
        key_ptr = builder.alloca(key_type)
        key_ptr.align = key_type.width // 8
        builder.store(ir.Constant(key_type, key_val), key_ptr)
    else:
        raise NotImplementedError(
            "Only simple variable names and integer constants are supported as keys in map delete."
        )

    if key_ptr is None:
        raise ValueError("Key pointer is None.")

    # Cast map pointer to void*
    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)

    # Helper ID 3 is bpf_map_delete_elem
    fn_addr = ir.Constant(ir.IntType(64), 3)
    fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)

    # Call the helper function
    result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)

    return result, None


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), 14)
    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), 25)
        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."
        )


helper_func_list = {
    "lookup": bpf_map_lookup_elem_emitter,
    "print": bpf_printk_emitter,
    "ktime": bpf_ktime_get_ns_emitter,
    "update": bpf_map_update_elem_emitter,
    "delete": bpf_map_delete_elem_emitter,
    "pid": bpf_get_current_pid_tgid_emitter,
    "output": bpf_perf_event_output_handler,
}


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
        if func_name in helper_func_list:
            # it is not a map method call
            return helper_func_list[func_name](
                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]
                if method_name in helper_func_list:
                    print(local_var_metadata)
                    return helper_func_list[method_name](
                        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]
                if method_name in helper_func_list:
                    return helper_func_list[method_name](
                        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
