#include <luna/Check.h>
#include <luna/DebugLog.h>
#include <luna/SourceLocation.h>
#include <luna/UBSAN.h>

using namespace UBSAN::UBInfo;

[[noreturn]] void ub_panic(SourceLocation caller = SourceLocation::current())
{
    fail_at(caller, "Undefined behavior");
}

#define DISPLAY(loc) loc.file, loc.line, loc.column

extern "C"
{
    void __ubsan_handle_builtin_unreachable(UnreachableInfo* info)
    {
        const auto& location = info->location;
        dbgln("ubsan: __builtin_unreachable reached at %s:%d:%d", DISPLAY(location));
        ub_panic();
    }

    void __ubsan_handle_pointer_overflow(PointerOverflowInfo* info)
    {
        const auto& location = info->location;
        dbgln("ubsan: pointer overflow occurred at %s:%d:%d", DISPLAY(location));
        ub_panic();
    }

    void __ubsan_handle_shift_out_of_bounds(ShiftOutOfBoundsInfo* info)
    {
        const auto& location = info->location;
        dbgln("ubsan: shift out of bounds for type %s at %s:%d:%d", info->lhs_type->name, DISPLAY(location));
        ub_panic();
    }

    void __ubsan_handle_load_invalid_value(InvalidValueInfo* info)
    {
        const auto& location = info->location;
        dbgln("ubsan: load invalid value for type %s at %s:%d:%d", info->type->name, DISPLAY(location));
        ub_panic();
    }

    void __ubsan_handle_out_of_bounds(OutOfBoundsInfo* info, usize index)
    {
        const auto& location = info->location;
        dbgln("ubsan: out of bounds array (of type %s) access (index %zu of type %s) at %s:%d:%d",
              info->array_type->name, index, info->index_type->name, DISPLAY(location));
        ub_panic();
    }

    void ubsan_handle_generic_overflow(OverflowInfo* info, const char* overflow_type)
    {
        const auto& location = info->location;
        dbgln("ubsan: %s overflow (value cannot fit into type %s) at %s:%d:%d", overflow_type, info->type->name,
              DISPLAY(location));
        ub_panic();
    }

#define UBSAN_OVERFLOW_BINARY(operation)                                                                               \
    void __ubsan_handle_##operation##_overflow(OverflowInfo* info, usize, usize)                                       \
    {                                                                                                                  \
        ubsan_handle_generic_overflow(info, #operation);                                                               \
    }

#define UBSAN_OVERFLOW_UNARY(operation)                                                                                \
    void __ubsan_handle_##operation##_overflow(OverflowInfo* info, usize)                                              \
    {                                                                                                                  \
        ubsan_handle_generic_overflow(info, #operation);                                                               \
    }

    UBSAN_OVERFLOW_BINARY(add);
    UBSAN_OVERFLOW_BINARY(sub);
    UBSAN_OVERFLOW_BINARY(mul);
    UBSAN_OVERFLOW_UNARY(negate);
    UBSAN_OVERFLOW_BINARY(divrem);

#define is_aligned(value, alignment) !(value & (alignment - 1))

    const char* g_type_check_kinds[] = {
        "load of",
        "store to",
        "reference binding to",
        "member access within",
        "member call on",
        "constructor call on",
        "downcast of",
        "downcast of",
        "upcast of",
        "cast to virtual base of",
    };

    void __ubsan_handle_type_mismatch(TypeMismatchInfo* info, usize pointer)
    {
        const auto& location = info->location;

        if (pointer == 0) { dbgln("ubsan: null pointer access at %s:%d:%d", DISPLAY(location)); }
        else if (info->alignment != 0 && is_aligned(pointer, info->alignment))
        {
            dbgln("ubsan: unaligned pointer access (address %p) at %s:%d:%d", (void*)pointer, DISPLAY(location));
        }
        else
        {
            dbgln("ubsan: %s address %p with insufficient space for object of type %s at %s:%d:%d",
                  g_type_check_kinds[info->type_check_kind], (void*)pointer, info->type->name, DISPLAY(location));
        }
        ub_panic();
    }

    void __ubsan_handle_type_mismatch_v1(TypeMismatchInfo_v1* v1_info, usize pointer)
    {
        TypeMismatchInfo info = {
            .location = v1_info->location,
            .type = v1_info->type,
            .alignment = 1UL << v1_info->log_alignment,
            .type_check_kind = v1_info->type_check_kind,
        };
        __ubsan_handle_type_mismatch(&info, pointer);
    }

    void __ubsan_handle_nonnull_arg(NonnullArgInfo* info, intptr_t)
    {
        dbgln("ubsan: null argument at %s:%d:%d", DISPLAY(info->location));
        ub_panic();
    }
}