"""
LLDB Formatters for LLVM data types.

Load into LLDB with 'command script import /path/to/lldbDataFormatters.py'
"""

from __future__ import annotations

import collections
import lldb


def __lldb_init_module(debugger, internal_dict):
    debugger.HandleCommand("type category define -e llvm -l c++")
    debugger.HandleCommand(
        "type synthetic add -w llvm "
        f"-l {__name__}.SmallVectorSynthProvider "
        '-x "^llvm::SmallVectorImpl<.+>$"'
    )
    debugger.HandleCommand(
        "type summary add -w llvm "
        '-e -s "size=${svar%#}" '
        '-x "^llvm::SmallVectorImpl<.+>$"'
    )
    debugger.HandleCommand(
        "type synthetic add -w llvm "
        f"-l {__name__}.SmallVectorSynthProvider "
        '-x "^llvm::SmallVector<.+,.+>$"'
    )
    debugger.HandleCommand(
        "type summary add -w llvm "
        '-e -s "size=${svar%#}" '
        '-x "^llvm::SmallVector<.+,.+>$"'
    )
    debugger.HandleCommand(
        "type synthetic add -w llvm "
        f"-l {__name__}.ArrayRefSynthProvider "
        '-x "^llvm::ArrayRef<.+>$"'
    )
    debugger.HandleCommand(
        "type summary add -w llvm "
        '-e -s "size=${svar%#}" '
        '-x "^llvm::ArrayRef<.+>$"'
    )
    debugger.HandleCommand(
        "type summary add -w llvm "
        f"-F {__name__}.SmallStringSummaryProvider "
        '-x "^llvm::SmallString<.+>$"'
    )
    debugger.HandleCommand(
        "type summary add -w llvm "
        f"-F {__name__}.StringRefSummaryProvider "
        "llvm::StringRef"
    )
    debugger.HandleCommand(
        "type summary add -w llvm "
        f"-F {__name__}.ConstStringSummaryProvider "
        "lldb_private::ConstString"
    )
    debugger.HandleCommand(
        "type synthetic add -w llvm "
        f"-l {__name__}.PointerIntPairSynthProvider "
        '-x "^llvm::PointerIntPair<.+>$"'
    )
    debugger.HandleCommand(
        "type synthetic add -w llvm "
        f"-l {__name__}.PointerUnionSynthProvider "
        '-x "^llvm::PointerUnion<.+>$"'
    )
    debugger.HandleCommand(
        "type summary add -w llvm "
        f"-e -F {__name__}.DenseMapSummary "
        '-x "^llvm::DenseMap<.+>$"'
    )
    debugger.HandleCommand(
        "type synthetic add -w llvm "
        f"-l {__name__}.DenseMapSynthetic "
        '-x "^llvm::DenseMap<.+>$"'
    )
    debugger.HandleCommand(
        "type synthetic add -w llvm "
        f"-l {__name__}.DenseSetSynthetic "
        '-x "^llvm::DenseSet<.+>$"'
    )

    debugger.HandleCommand(
        "type synthetic add -w llvm "
        f"-l {__name__}.ExpectedSynthetic "
        '-x "^llvm::Expected<.+>$"'
    )
    debugger.HandleCommand(
        "type summary add -w llvm "
        f"-F {__name__}.SmallBitVectorSummary "
        "llvm::SmallBitVector"
    )


# Pretty printer for llvm::SmallVector/llvm::SmallVectorImpl
class SmallVectorSynthProvider:
    def __init__(self, valobj, internal_dict):
        self.valobj = valobj
        self.update()  # initialize this provider

    def num_children(self):
        return self.size.GetValueAsUnsigned(0)

    def get_child_index(self, name):
        try:
            return int(name.lstrip("[").rstrip("]"))
        except:
            return -1

    def get_child_at_index(self, index):
        # Do bounds checking.
        if index < 0:
            return None
        if index >= self.num_children():
            return None

        offset = index * self.type_size
        return self.begin.CreateChildAtOffset(
            "[" + str(index) + "]", offset, self.data_type
        )

    def update(self):
        self.begin = self.valobj.GetChildMemberWithName("BeginX")
        self.size = self.valobj.GetChildMemberWithName("Size")
        the_type = self.valobj.GetType()
        # If this is a reference type we have to dereference it to get to the
        # template parameter.
        if the_type.IsReferenceType():
            the_type = the_type.GetDereferencedType()

        if the_type.IsPointerType():
            the_type = the_type.GetPointeeType()

        self.data_type = the_type.GetTemplateArgumentType(0)
        self.type_size = self.data_type.GetByteSize()
        assert self.type_size != 0


class ArrayRefSynthProvider:
    """Provider for llvm::ArrayRef"""

    def __init__(self, valobj, internal_dict):
        self.valobj = valobj
        self.update()  # initialize this provider

    def num_children(self):
        return self.length

    def get_child_index(self, name):
        try:
            return int(name.lstrip("[").rstrip("]"))
        except:
            return -1

    def get_child_at_index(self, index):
        if index < 0 or index >= self.num_children():
            return None
        offset = index * self.type_size
        return self.data.CreateChildAtOffset(
            "[" + str(index) + "]", offset, self.data_type
        )

    def update(self):
        self.data = self.valobj.GetChildMemberWithName("Data")
        length_obj = self.valobj.GetChildMemberWithName("Length")
        self.length = length_obj.GetValueAsUnsigned(0)
        self.data_type = self.data.GetType().GetPointeeType()
        self.type_size = self.data_type.GetByteSize()
        assert self.type_size != 0


def SmallStringSummaryProvider(valobj, internal_dict):
    # The underlying SmallVector base class is the first child.
    vector = valobj.GetChildAtIndex(0)
    num_elements = vector.GetNumChildren()
    res = '"'
    for i in range(num_elements):
        c = vector.GetChildAtIndex(i)
        if c:
            res += chr(c.GetValueAsUnsigned())
    res += '"'
    return res


def StringRefSummaryProvider(valobj, internal_dict):
    data_pointer = valobj.GetChildMemberWithName("Data")
    length = valobj.GetChildMemberWithName("Length").unsigned
    if data_pointer.unsigned == 0 or length == 0:
        return '""'

    data = data_pointer.deref
    # StringRef may be uninitialized with length exceeding available memory,
    # potentially causing bad_alloc exceptions. Limit the length to max string summary setting.
    limit_obj = valobj.target.debugger.GetSetting("target.max-string-summary-length")
    if limit_obj:
        length = min(length, limit_obj.GetUnsignedIntegerValue())
    # Get a char[N] type, from the underlying char type.
    array_type = data.type.GetArrayType(length)
    # Cast the char* string data to a char[N] array.
    char_array = data.Cast(array_type)
    # Use the builtin summary for its support of max-string-summary-length and
    # display of non-printable bytes.
    return char_array.summary


def ConstStringSummaryProvider(valobj, internal_dict):
    if valobj.GetNumChildren() == 1:
        return valobj.GetChildAtIndex(0).GetSummary()
    return ""


class PointerIntPairSynthProvider:
    def __init__(self, valobj, internal_dict):
        self.valobj = valobj
        self.update()

    def num_children(self):
        return 2

    def get_child_index(self, name):
        if name == "Pointer":
            return 0
        if name == "Int":
            return 1
        return None

    def _get_raw_value(self):
        data: SBData = self.value.GetData()
        error = lldb.SBError()
        raw_bytes = data.ReadRawData(error, 0, self.ptr_size)
        if error.Fail():
            return None

        return raw_bytes

    def _get_pointer(self, pointer_bit_mask: int, pointer_ty: SBType):
        raw_bytes = self._get_raw_value()
        if raw_bytes is None:
            return

        unmasked_pointer = int.from_bytes(raw_bytes, self.byteorder)
        pointer_value = unmasked_pointer & pointer_bit_mask

        data = lldb.SBData()
        data.SetDataFromUInt64Array([pointer_value])
        return self.valobj.CreateValueFromData("Pointer", data, pointer_ty)

    def _get_int(self, int_shift: int, int_mask: int, int_ty: SBType):
        raw_bytes = self._get_raw_value()
        if raw_bytes is None:
            return

        unmasked_pointer = int.from_bytes(raw_bytes, self.byteorder)
        int_value = (unmasked_pointer >> int_shift) & int_mask

        data = lldb.SBData()
        data.SetDataFromUInt64Array([int_value])
        return self.valobj.CreateValueFromData("Int", data, int_ty)

    def get_child_at_index(self, index):
        if index == 0:
            return self.pointer_valobj
        if index == 1:
            return self.int_valobj
        return None

    def update(self):
        self.byteorder = (
            "big"
            if self.valobj.target.GetByteOrder() == lldb.eByteOrderBig
            else "little"
        )
        self.ptr_size = self.valobj.target.GetAddressByteSize()
        self.value: SBValue = self.valobj.GetChildMemberWithName("Value")
        if not self.value:
            return

        valobj_type = self.valobj.GetType()

        pointer_ty: SBType = valobj_type.GetTemplateArgumentType(0)
        if not pointer_ty:
            return

        int_ty: SBType = valobj_type.GetTemplateArgumentType(2)
        if not int_ty:
            return

        pointer_info = valobj_type.GetTemplateArgumentType(4)
        if not pointer_info:
            return

        mask_and_shift_constants = pointer_info.FindDirectNestedType(
            "MaskAndShiftConstants"
        ).GetEnumMembers()

        # FIXME: SBAPI should provide a way to retrieve an enum member
        # by name.
        pointer_bit_mask: SBTypeEnumMember = (
            mask_and_shift_constants.GetTypeEnumMemberAtIndex(0)
        )
        if pointer_bit_mask.name != "PointerBitMask":
            return

        int_shift: SBTypeEnumMember = mask_and_shift_constants.GetTypeEnumMemberAtIndex(
            1
        )
        if int_shift.name != "IntShift":
            return

        int_mask: SBTypeEnumMember = mask_and_shift_constants.GetTypeEnumMemberAtIndex(
            2
        )
        if int_mask.name != "IntMask":
            return

        self.pointer_valobj = self._get_pointer(
            pointer_bit_mask.GetValueAsUnsigned(), pointer_ty
        )
        self.int_valobj = self._get_int(
            int_shift.GetValueAsUnsigned(), int_mask.GetValueAsUnsigned(), int_ty
        )


class PointerUnionSynthProvider:
    def __init__(self, valobj, internal_dict):
        self.valobj = valobj
        self.update()

    def num_children(self):
        return 1

    def get_child_index(self, name):
        if name == "Pointer":
            return 0
        return None

    def get_child_at_index(self, index):
        if index != 0:
            return None

        return self.pointer_valobj

    def update(self):
        pointer_int_pair: SBValue = self.valobj.GetChildMemberWithName(
            "Val"
        ).GetSyntheticValue()

        if not pointer_int_pair:
            return

        pointer: SBValue = pointer_int_pair.GetChildAtIndex(0)
        if not pointer:
            return

        active_tag: SBValue = pointer_int_pair.GetChildAtIndex(1)
        if not active_tag:
            return

        # Index into the parameter pack of llvm::PointerUnion to find the active type.
        active_type: SBType = self.valobj.GetType().GetTemplateArgumentType(
            active_tag.GetValueAsUnsigned()
        )
        if not active_type:
            return

        data = lldb.SBData()
        data.SetDataFromUInt64Array([pointer.GetValueAsUnsigned()])

        self.pointer_valobj = self.valobj.CreateValueFromData(
            "Pointer", data, active_type
        )


def DenseMapSummary(valobj: lldb.SBValue, _) -> str:
    raw_value = valobj.GetNonSyntheticValue()
    num_entries = raw_value.GetChildMemberWithName("NumEntries").unsigned
    num_tombstones = raw_value.GetChildMemberWithName("NumTombstones").unsigned

    summary = f"size={num_entries}"
    if num_tombstones == 1:
        # The heuristic to identify valid entries does not handle the case of a
        # single tombstone. The summary calls attention to this.
        summary = f"tombstones=1, {summary}"
    return summary


class DenseMapSynthetic:
    valobj: lldb.SBValue

    # The indexes into `Buckets` that contain valid map entries.
    child_buckets: list[int]

    def __init__(self, valobj: lldb.SBValue, _) -> None:
        self.valobj = valobj

    def num_children(self) -> int:
        return len(self.child_buckets)

    def get_child_at_index(self, child_index: int) -> lldb.SBValue:
        bucket_index = self.child_buckets[child_index]
        entry = self.valobj.GetValueForExpressionPath(f".Buckets[{bucket_index}]")

        # By default, DenseMap instances use DenseMapPair to hold key-value
        # entries. When the entry is a DenseMapPair, unwrap it to expose the
        # children as simple std::pair values.
        #
        # This entry type is customizable (a template parameter). For other
        # types, expose the entry type as is.
        if entry.type.name.startswith("llvm::detail::DenseMapPair<"):
            entry = entry.GetChildAtIndex(0)

        return entry.Clone(f"[{child_index}]")

    def update(self):
        self.child_buckets = []

        num_entries = self.valobj.GetChildMemberWithName("NumEntries").unsigned
        if num_entries == 0:
            return

        buckets = self.valobj.GetChildMemberWithName("Buckets")
        num_buckets = self.valobj.GetChildMemberWithName("NumBuckets").unsigned

        # Bucket entries contain one of the following:
        #   1. Valid key-value
        #   2. Empty key
        #   3. Tombstone key (a deleted entry)
        #
        # NumBuckets is always greater than NumEntries. The empty key, and
        # potentially the tombstone key, will occur multiple times. A key that
        # is repeated is either the empty key or the tombstone key.

        # For each key, collect a list of buckets it appears in.
        key_buckets: dict[str, list[int]] = collections.defaultdict(list)
        for index in range(num_buckets):
            bucket = buckets.GetValueForExpressionPath(f"[{index}]")
            key = bucket.GetChildAtIndex(0)
            key_buckets[str(key.data)].append(index)

        # Heuristic: This is not a multi-map, any repeated (non-unique) keys are
        # either the the empty key or the tombstone key. Populate child_buckets
        # with the indexes of entries containing unique keys.
        for indexes in key_buckets.values():
            if len(indexes) == 1:
                self.child_buckets.append(indexes[0])


class DenseSetSynthetic:
    valobj: lldb.SBValue
    map: lldb.SBValue

    def __init__(self, valobj: lldb.SBValue, _) -> None:
        self.valobj = valobj

    def num_children(self) -> int:
        return self.map.num_children

    def get_child_at_index(self, idx: int) -> lldb.SBValue:
        map_entry = self.map.child[idx]
        set_entry = map_entry.GetChildAtIndex(0)
        return set_entry.Clone(f"[{idx}]")

    def update(self):
        raw_map = self.valobj.GetChildMemberWithName("TheMap")
        self.map = raw_map.GetSyntheticValue()


class ExpectedSynthetic:
    # The llvm::Expected<T> value.
    expected: lldb.SBValue
    # The stored success value or error value.
    stored_value: lldb.SBValue

    def __init__(self, valobj: lldb.SBValue, _) -> None:
        self.expected = valobj

    def update(self) -> None:
        has_error = self.expected.GetChildMemberWithName("HasError").unsigned
        if not has_error:
            name = "value"
            member = "TStorage"
        else:
            name = "error"
            member = "ErrorStorage"
        # Anonymous union.
        union = self.expected.child[0]
        storage = union.GetChildMemberWithName(member)
        stored_type = storage.type.template_args[0]
        self.stored_value = storage.Cast(stored_type).Clone(name)

    def num_children(self) -> int:
        return 1

    def get_child_index(self, name: str) -> int:
        if name == self.stored_value.name:
            return 0
        # Allow dereferencing for values, not errors.
        if name == "$$dereference$$" and self.stored_value.name == "value":
            return 0
        return -1

    def get_child_at_index(self, idx: int) -> lldb.SBValue:
        if idx == 0:
            return self.stored_value
        return lldb.SBValue()


def SmallBitVectorSummary(valobj, _):
    underlyingValue = valobj.GetChildMemberWithName("X").unsigned
    numBaseBits = valobj.target.addr_size * 8
    smallNumRawBits = numBaseBits - 1
    smallNumSizeBits = None
    if numBaseBits == 32:
        smallNumSizeBits = 5
    elif numBaseBits == 64:
        smallNumSizeBits = 6
    else:
        smallNumSizeBits = smallNumRawBits
    smallNumDataBits = smallNumRawBits - smallNumSizeBits

    # If our underlying value is not small, print we can not dump large values.
    isSmallMask = 1
    if underlyingValue & isSmallMask == 0:
        return "<can not read large SmallBitVector>"

    smallRawBits = underlyingValue >> 1
    smallSize = smallRawBits >> smallNumDataBits
    bits = smallRawBits & ((1 << (smallSize + 1)) - 1)
    # format `bits` in binary (b), with 0 padding, of width `smallSize`, and left aligned (>)
    return f"[{bits:0>{smallSize}b}]"