from typing import Callable

from lldb import *
from renderers.jb_lldb_primitive_type import PrimitiveTypeSpec, PrimitiveType
from renderers.jb_lldb_format_specs import *
from renderers.jb_lldb_utils import *
from renderers.jb_lldb_string_utils import *
from renderers.jb_lldb_declarative_formatters_options import *
from renderers.jb_lldb_format import update_value_dynamic_state, get_array_size_format


class StructChildrenProvider(AbstractChildrenProvider):
    def __init__(self, value_non_synth: lldb.SBValue):
        self.value = value_non_synth
        self.format_spec = value_non_synth.GetFormat()
        self._indirect_children_name_to_index: dict[str, int] = {}
        self._indirect_children_index_to_name: dict[int, str] = {}

    def num_children(self) -> int:
        return self.value.GetNumChildren()

    def get_child_index(self, name) -> int:
        direct_child_index = self.value.GetIndexOfChildWithName(name)
        if direct_child_index != INVALID_CHILD_INDEX:
            return direct_child_index
        indirect_child_index = self._indirect_children_name_to_index.get(name)
        if indirect_child_index is not None:
            return indirect_child_index
        # GetIndexOfChildWithName() only searches direct children.
        # Fall back to GetChildMemberWithName() which recurses into anonymous structs and base classes.
        # This is needed because custom and built-in LLDB formatters may call GetChildMemberWithName() on the ValueObject,
        # which gets routed through our synthetic provider's get_child_index().
        indirect_child_member: lldb.SBValue = self.value.GetChildMemberWithName(name)
        if indirect_child_member.IsValid():
            indirect_child_index = self.value.GetNumChildren() + len(self._indirect_children_index_to_name)
            RENDER_LOG.debug("Indirect member '%s' of type '%s' mapped to virtual index %d",
                             name, self.value.GetTypeName(), indirect_child_index)
            self._indirect_children_name_to_index[name] = indirect_child_index
            self._indirect_children_index_to_name[indirect_child_index] = name
            return indirect_child_index
        return INVALID_CHILD_INDEX

    def get_child_at_index(self, index):
        child = self._get_child_by_index(index)
        # non_synth children haven't any formatting
        # so apply parent value formatting without non-inheritable flags
        child_root = get_root_value(child)
        child_root.SetFormat(self.format_spec & eFormatInheritedFlagsMask)

        ItemExpression.update_struct_child_item_expression(child, self.value)
        update_value_dynamic_state(child)

        return child

    def _get_child_by_index(self, index: int) -> lldb.SBValue:
        indirect_child_name = self._indirect_children_index_to_name.get(index)
        if indirect_child_name is not None:
            return self.value.GetChildMemberWithName(indirect_child_name)
        return self.value.GetChildAtIndex(index)


class WrappedChildrenProvider(AbstractChildrenProvider):
    def __init__(self, value: lldb.SBValue):
        self.value = value

    def num_children(self) -> int:
        return self.value.GetNumChildren()

    def get_child_index(self, name) -> int:
        return self.value.GetIndexOfChildWithName(name)

    def get_child_at_index(self, index):
        child = self.value.GetChildAtIndex(index)
        update_value_dynamic_state(child)
        return child


class PointerChildrenProvider(AbstractChildrenProvider):
    def __init__(self, value_non_synth: lldb.SBValue):
        # no need to manipulate formats here
        # pointees inherit all formatting from pointer parents (?)
        self.pointee = ItemExpression.dereference(value_non_synth)

    def num_children(self) -> int:
        if not self.pointee.IsValid():
            return 0
        return 1

    def get_child_index(self, name) -> int:
        return INVALID_CHILD_INDEX

    def get_child_at_index(self, index):
        return self.pointee


class AbstractNumberVisDescriptor(AbstractVisDescriptor):
    formats = {
        eFormatHex, eFormatHexNoPrefix, eFormatHexUppercase, eFormatHexUppercaseNoPrefix,
        eFormatOctal, eFormatBinary, eFormatBinaryNoPrefix, eFormatDecimal, eFormatChar
    }
    eFormatHexShowBoth = eFormatBasicSpecsMask + 1

    def __init__(self):
        self.presenters: dict[int, Callable[[lldb.SBValue, Stream], None]] = {}
        self.default_presenter: Callable[[lldb.SBValue, Stream], None] = lambda x, s: NumberVisDescriptor.output_integer(x, s)

    @classmethod
    def make_presenters(cls, size: int, signed: bool) -> dict[int, Callable[[lldb.SBValue, Stream], None]]:
        format_hex = "0x{{:0{}x}}".format(size * 2)
        format_no_prefix = "{{:0{}x}}".format(size * 2)
        format_hex_uppercase = "0x{{:0{}X}}".format(size * 2)
        format_hex_uppercase_no_prefix = "{{:0{}X}}".format(size * 2)
        format_octal = "{{:0{}o}}".format(size * 3)
        format_binary = "0b{{:0{}b}}".format(size * 8)
        format_binary_no_prefix = "{{:0{}b}}".format(size * 8)

        def output_number_with_hex(x, s: Stream):
            NumberVisDescriptor.output_integer(x, s)
            s.output(" [")
            s.output_number(format_hex.format(x.GetValueAsUnsigned()))
            s.output("]")

        return {
            eFormatHex: lambda x, s: s.output_number(format_hex.format(x.GetValueAsUnsigned())),
            eFormatHexNoPrefix: lambda x, s: s.output_number(format_no_prefix.format(x.GetValueAsUnsigned())),
            eFormatHexUppercase: lambda x, s: s.output_number(format_hex_uppercase.format(x.GetValueAsUnsigned())),
            eFormatHexUppercaseNoPrefix: lambda x, s: s.output_number(format_hex_uppercase_no_prefix.format(x.GetValueAsUnsigned())),
            eFormatOctal: lambda x, s: s.output_number(format_octal.format(x.GetValueAsUnsigned())),
            eFormatBinary: lambda x, s: s.output_number(format_binary.format(x.GetValueAsUnsigned())),
            eFormatBinaryNoPrefix: lambda x, s: s.output_number(format_binary_no_prefix.format(x.GetValueAsUnsigned())),
            eFormatDecimal: lambda x, s: NumberVisDescriptor.output_decimal(x, signed, s),
            eFormatChar: lambda x, s: NumberVisDescriptor.output_char(x.GetValueAsUnsigned(), s),
            cls.eFormatHexShowBoth: output_number_with_hex
        }

    @staticmethod
    def output_char(code: int, stream: Stream):
        # char format specifier interprets integers as shorts in utf-16 encoding
        code = code & 0xffff
        enc = 'utf-16'
        stream.output_number(str(code))
        stream.output(" ")
        stream.output_string("'" + escape_char(code, 2, enc) + "'")

    @staticmethod
    def output_integer(value_non_synth: lldb.SBValue, stream: Stream):
        # no need to get static value for numeric values
        value_non_synth.SetFormat(eFormatDefault)
        stream.output_number(value_non_synth.GetValue())

    @staticmethod
    def output_decimal(value_non_synth: lldb.SBValue, signed: bool, stream: Stream):
        fmt = signed and eFormatDecimal or eFormatUnsigned
        # no need to get static value for numeric values
        value_non_synth.SetFormat(fmt)
        stream.output_number(value_non_synth.GetValue())

    @staticmethod
    def output_float(value_non_synth: lldb.SBValue, stream: Stream):
        # no need to get static value for numeric values
        value_non_synth.SetFormat(eFormatFloat)
        stream.output_number(value_non_synth.GetValue())

    def output_summary(self, value_non_synth: lldb.SBValue, stream: Stream):
        if not value_non_synth.IsValid():
            stream.output('???')
            return
        err: SBError = value_non_synth.GetError()
        if err is not None and err.Fail():
            stream.output('???')
            return

        spec = value_non_synth.GetFormat() & eFormatBasicSpecsMask
        # check for global hex option if format specifier was not explicitly set
        if is_global_hex() and spec not in self.formats:
            spec = eFormatHex
            if is_global_hex_show_both():
                spec = self.eFormatHexShowBoth
        presenter = self.presenters.get(spec, self.default_presenter)
        presenter(value_non_synth, stream)


class NumberVisDescriptor(AbstractNumberVisDescriptor):
    integer_types = {
        "char": PrimitiveType.Char, "signed char": PrimitiveType.SignedChar, "unsigned char": PrimitiveType.UnsignedChar,
        "wchar_t": PrimitiveType.WChar, "char8_t": PrimitiveType.Char8, "char16_t": PrimitiveType.Char16, "char32_t": PrimitiveType.Char32,
        "short": PrimitiveType.Short, "unsigned short": PrimitiveType.UnsignedShort,
        "int": PrimitiveType.Int, "unsigned int": PrimitiveType.UnsignedInt, "unsigned": PrimitiveType.UnsignedInt,
        "long": PrimitiveType.Long, "long int": PrimitiveType.Long,
        "unsigned long": PrimitiveType.UnsignedLong, "unsigned long int": PrimitiveType.UnsignedLong,
        "long long": PrimitiveType.LongLong, "long long int": PrimitiveType.LongLong,
        "unsigned long long": PrimitiveType.UnsignedLongLong, "unsigned long long int": PrimitiveType.UnsignedLongLong,
        "__int128": PrimitiveType.Int128, "unsigned __int128": PrimitiveType.UnsignedInt128
    }
    float_types = {"half", "float", "double", "long double"}

    @classmethod
    def is_number_type(cls, type_name: str) -> bool:
        return type_name == "bool" or type_name in cls.integer_types or type_name in cls.float_types

    def __init__(self, value_type: lldb.SBType, value_type_name: str):
        super().__init__()
        if value_type_name == "bool":
            self.default_presenter = lambda x, s: s.output_keyword("true" if x.GetValueAsUnsigned() != 0 else "false")
        elif primitive_type := self.integer_types.get(value_type_name):
            primitive_type_spec = primitive_type.resolve_spec(None, value_type)
            self.presenters = self.make_presenters(primitive_type_spec.byte_size, primitive_type_spec.signed)
        elif value_type_name in self.float_types:
            self.default_presenter = lambda x, s: NumberVisDescriptor.output_float(x, s)


class EnumVisDescriptor(AbstractVisDescriptor):
    def output_summary(self, value_non_synth: lldb.SBValue, stream: Stream):
        original_fmt = value_non_synth.GetFormat()
        basic_fmt = original_fmt & eFormatBasicSpecsMask
        value_non_synth.SetFormat(basic_fmt)
        summary_value = value_non_synth.GetValue() or ''
        stream.output(summary_value)


class PointerAsIntegerVisDescriptor(AbstractNumberVisDescriptor):
    def __init__(self, pointer_size: int):
        super().__init__()
        self.presenters = self.make_presenters(pointer_size, False)


class CharVisDescriptor(AbstractVisDescriptor):
    char_types = {
        "char": PrimitiveType.Char,
        "unsigned char": PrimitiveType.UnsignedChar,
        "signed char": PrimitiveType.SignedChar,
        "wchar_t": PrimitiveType.WChar,
        "char8_t": PrimitiveType.Char8,
        "char16_t": PrimitiveType.Char16,
        "char32_t": PrimitiveType.Char32
    }

    def __init__(self, char_spec: PrimitiveTypeSpec):
        self.char_spec = char_spec

    def output_summary(self, value_non_synth: lldb.SBValue, stream: Stream):
        enc = self.char_spec.encoding
        if enc == '__locale__':
            enc = get_locale()
        err = SBError()
        code = value_non_synth.GetValueAsSigned(err)
        if err.Fail():
            stream.output('<error>')
            return
        ordinal = code
        if code < 0:
            # convert signed to unsigned
            if code >= -0x80:
                ordinal = code & 0xff
            elif code >= -0x8000:
                ordinal = code & 0xffff
            elif code >= -0x80000000:
                ordinal = code & 0xffffffff
        if is_global_hex():
            if is_global_hex_show_both():
                stream.output_number(str(code))
                stream.output(" [")
                stream.output_number("{0:#x}".format(ordinal))
                stream.output("]")
            else:
                stream.output_number("{0:#x}".format(ordinal))
        else:
            stream.output_number(str(code))
        stream.output(" ")
        # Unlike Visual Studio's behavior we present wchar_t as L'x' (for consistency with wide strings L"xxx")
        escaped_char = escape_char(ordinal, self.char_spec.byte_size, enc)
        stream.output_string(f"{self.char_spec.literal_prefix}'{escaped_char}'")


def _cast_value_to_array(value_non_synth: lldb.SBValue, is_array: bool, fmt: int, array_size: int):
    new_val = ItemExpression.cast_value_to_array(value_non_synth, is_array, array_size)

    # reset format as array flag
    new_val_fmt = fmt & ~eFormatAsArray
    set_value_format(new_val, new_val_fmt)
    return new_val


class CharArrayOrPointerVisDescriptor(AbstractVisDescriptor):
    _basic_types_convertable_to_pointer = {
        eBasicTypeInt, eBasicTypeUnsignedInt, eBasicTypeLong, eBasicTypeUnsignedLong, eBasicTypeLongLong, eBasicTypeUnsignedLongLong
    }

    @classmethod
    def can_type_be_used_as_char_pointer(cls, value_type: lldb.SBType):
        return value_type.GetBasicType() in cls._basic_types_convertable_to_pointer

    def __init__(self, char_spec: PrimitiveTypeSpec, is_array: bool, array_size=None, formatted_array=False):
        self.char_spec = char_spec
        self.is_array = is_array
        self.array_size = array_size
        self.formatted_array = formatted_array

    def output_summary(self, value_non_synth: lldb.SBValue, stream: Stream):
        fmt = value_non_synth.GetFormat()

        if fmt & eFormatAsArray != 0:
            provider, new_val = self._format_as_array(value_non_synth, fmt)
            return provider.output_summary(new_val.GetNonSyntheticValue(), stream)

        basic_fmt = fmt & eFormatBasicSpecsMask
        suppress_address = (fmt & eFormatNoAddress != 0) or (basic_fmt in FMT_STRING_SET_ALL)
        if self.is_array:
            address = value_non_synth.GetLoadAddress()
        else:
            address = value_non_synth.GetValueAsUnsigned()

        if not suppress_address:
            stream.output_address(address)
            stream.output_comment(" ") # print space as a comment to skip it during text extraction (copy value)

        zero_required = not self.formatted_array

        err = lldb.SBError()
        max_size = self.array_size * self.char_spec.byte_size if self.array_size is not None else None
        content, zero_found = extract_string(value_non_synth, self.is_array, address, self.char_spec.byte_size, max_size, err)
        if err.Fail():
            stream.output('<error>')
            return

        enc = self.char_spec.encoding
        if enc == '__locale__':
            enc = get_locale()
        s = escape_bytes(content, enc)

        quote = basic_fmt not in FMT_STRING_NOQUOTES_SET
        result = StringIO()

        if quote:
            result.write(self.char_spec.literal_prefix)
            result.write('"')
        result.write(s)
        if zero_found or not zero_required:
            if quote:
                result.write('"')
        else:
            result.write('...')

        # quoted: STRING token (frontend adds quotes and the needed coloring)
        # unquoted: natvis specifiers (sb/s8b/sub/s32b) -> CHAR value token (rendered as it is)
        if quote:
            stream.output_string(result.getvalue())
        else:
            stream.output_value(result.getvalue())


    def prepare_children(self, value_non_synth: lldb.SBValue):
        fmt = value_non_synth.GetFormat()

        if fmt & eFormatAsArray != 0:
            provider, new_val = self._format_as_array(value_non_synth, fmt)
            return provider.prepare_children(new_val.GetNonSyntheticValue())

        return StructChildrenProvider(value_non_synth)

    def _format_as_array(self, value_non_synth: lldb.SBValue, fmt: int) -> tuple[AbstractVisDescriptor, lldb.SBValue]:
        fmt_array_size = get_array_size_format(value_non_synth)
        # TODO: check fmt_array_size

        provider = CharArrayOrPointerVisDescriptor(self.char_spec, True, fmt_array_size, True)
        if self.can_type_be_used_as_char_pointer(value_non_synth.GetType()):
            value_non_synth = ItemExpression.cast_value_to_basic_type_pointer(value_non_synth, self.char_spec.lldb_basic_type)
        new_val = _cast_value_to_array(value_non_synth, self.is_array, fmt, fmt_array_size)
        return provider, new_val


class GenericArrayVisDescriptor(AbstractVisDescriptor):
    def output_summary(self, value_non_synth: lldb.SBValue, stream: Stream):
        fmt = value_non_synth.GetFormat()
        if fmt & eFormatAsArray != 0:
            provider, new_val = self._format_as_array(value_non_synth, fmt)
            return provider.output_summary(new_val.GetNonSyntheticValue(), stream)

        basic_fmt = fmt & eFormatBasicSpecsMask
        if char_kind := FMT_STRING_SET_ALL.get(basic_fmt):
            char_spec = char_kind.resolve_spec(value_non_synth)
            provider = CharArrayOrPointerVisDescriptor(char_spec, True, None)
            return provider.output_summary(value_non_synth, stream)

        address = value_non_synth.GetLoadAddress()
        stream.output_address(address)
        stream.output(" {")

        if stream.level >= g_max_recursion_level:
            stream.output('...')
        else:
            for child_index in range(value_non_synth.GetNumChildren()):
                if child_index != 0:
                    stream.output(", ")
                if stream.length > get_max_string_length():
                    stream.output("...")
                    break

                child: lldb.SBValue = value_non_synth.GetChildAtIndex(child_index)
                update_value_dynamic_state(child)
                stream.output_object(child.GetNonSyntheticValue())

        stream.output("}")

    def prepare_children(self, value_non_synth: lldb.SBValue):
        fmt = value_non_synth.GetFormat()

        if fmt & eFormatAsArray != 0:
            provider, new_val = self._format_as_array(value_non_synth, fmt)
            return provider.prepare_children(new_val.GetNonSyntheticValue())

        return StructChildrenProvider(value_non_synth)

    @staticmethod
    def _format_as_array(value_non_synth: lldb.SBValue, fmt: int) -> tuple[AbstractVisDescriptor, lldb.SBValue]:
        fmt_array_size = get_array_size_format(value_non_synth)
        # TODO: check fmt_array_size
        provider = GenericArrayVisDescriptor()
        new_val = _cast_value_to_array(value_non_synth, True, fmt, fmt_array_size)
        return provider, new_val


RESOLVE_CONTEXT_MASK = lldb.eSymbolContextModule | lldb.eSymbolContextFunction | lldb.eSymbolContextSymbol

class GenericPointerVisDescriptor(AbstractVisDescriptor):
    def __init__(self, pointee_expands: bool, pointee_has_empty_description: bool):
        self.pointee_expands = pointee_expands
        self.pointee_has_empty_description = pointee_has_empty_description

    def output_summary(self, value_non_synth: lldb.SBValue, stream: Stream):
        fmt = value_non_synth.GetFormat()
        if fmt & eFormatAsArray != 0:
            provider, new_val = self._format_as_array(value_non_synth, fmt)
            return provider.output_summary(new_val.GetNonSyntheticValue(), stream)

        basic_fmt = fmt & eFormatBasicSpecsMask
        if char_kind := FMT_STRING_SET_ALL.get(basic_fmt):
            char_spec = char_kind.resolve_spec(value_non_synth)
            provider = CharArrayOrPointerVisDescriptor(char_spec, False, None)
            return provider.output_summary(value_non_synth, stream)

        err = SBError()
        address = value_non_synth.GetValueAsUnsigned(err)
        if err.Fail():
            stream.output('???')
            return

        need_separator = False
        if value_non_synth.GetFormat() & eFormatNoAddress == 0:
            stream.output_address(address)
            need_separator = True

        sb_addr = SBAddress(address, value_non_synth.GetTarget())
        symbol_context: lldb.SBSymbolContext = sb_addr.GetSymbolContext(RESOLVE_CONTEXT_MASK)
        if symbol_context is not None:
            module: lldb.SBModule = symbol_context.GetModule()

            function: lldb.SBFunction = symbol_context.GetFunction()
            symbol: lldb.SBSymbol = symbol_context.GetSymbol()
            name = None
            if function is not None and function.IsValid():
                name = function.GetDisplayName() or function.GetName()
            elif symbol is not None and symbol.IsValid():
                name = symbol.GetDisplayName() or symbol.GetName()

            if name is not None:
                if need_separator:
                    stream.output(" ")
                stream.output("{")

                if module is not None and module.IsValid():
                    file: lldb.SBFileSpec = module.GetFileSpec()
                    file_name = file.GetFilename()
                    stream.output(file_name)
                    stream.output("!")

                stream.output(name)
                stream.output("}")

                need_separator = True

        if self.pointee_has_empty_description:
            return

        if need_separator:
            stream.output(" ")

        if stream.level >= g_max_recursion_level:
            stream.output('{...}')
        elif self.pointee_expands:
            stream.output_object(ItemExpression.dereference(value_non_synth).GetNonSyntheticValue())
        else:
            stream.output('{')
            if stream.length > get_max_string_length():
                stream.output("...")
            else:
                stream.output_object(ItemExpression.dereference(value_non_synth).GetNonSyntheticValue())
            stream.output('}')

    def prepare_children(self, value_non_synth: lldb.SBValue):
        fmt = value_non_synth.GetFormat()
        if fmt & eFormatAsArray != 0:
            provider, new_val = self._format_as_array(value_non_synth, fmt)
            return provider.prepare_children(new_val.GetNonSyntheticValue())

        err = SBError()
        address = value_non_synth.GetValueAsUnsigned(err)
        if err.Fail():
            return AbstractChildrenProvider()

        if self.pointee_expands:
            return WrappedChildrenProvider(ItemExpression.dereference(value_non_synth))
        return PointerChildrenProvider(value_non_synth)

    @staticmethod
    def _format_as_array(value_non_synth: lldb.SBValue, fmt: int) -> tuple[AbstractVisDescriptor, lldb.SBValue]:
        fmt_array_size = get_array_size_format(value_non_synth)
        # TODO: check fmt_array_size

        provider = GenericArrayVisDescriptor()
        new_val = _cast_value_to_array(value_non_synth, False, fmt, fmt_array_size)
        return provider, new_val


class GenericReferenceVisDescriptor(AbstractVisDescriptor):
    def output_summary(self, value_non_synth: lldb.SBValue, stream: Stream):
        if stream.level >= g_max_recursion_level:
            stream.output('{...}')
        else:
            stream.output_object(ItemExpression.dereference(value_non_synth).GetNonSyntheticValue())

    def prepare_children(self, value_non_synth: lldb.SBValue):
        return WrappedChildrenProvider(ItemExpression.dereference(value_non_synth))


class StructVisDescriptor(AbstractVisDescriptor):
    def __init__(self, value_type: lldb.SBType):
        self.value_type = value_type

    def output_summary(self, value_non_synth: lldb.SBValue, stream: Stream):
        provider = self.prepare_children(value_non_synth)
        num_children = provider.num_children()

        # Skip base structs in summary presentation
        base_classes_count = self.value_type.GetNumberOfDirectBaseClasses()
        # TODO: what about virtual bases?
        # TODO: there is bug with empty bases that are not presented

        stream.output("{")
        if stream.level >= g_max_recursion_level or stream.length > get_max_string_length():
            stream.output('...')
        elif num_children == base_classes_count:
            stream.output('...')
        else:
            for child_index in range(base_classes_count, num_children):
                if child_index != base_classes_count:
                    stream.output(", ")

                if child_index > base_classes_count + 2 or stream.length > get_max_string_length():
                    stream.output("...")
                    break

                child: lldb.SBValue = provider.get_child_at_index(child_index)
                child_non_synth = child.GetNonSyntheticValue()
                child_name = child_non_synth.GetName() or ''
                stream.output(child_name)
                stream.output("=")
                if stream.length > get_max_string_length():
                    stream.output("...")
                    break

                stream.output_object(child_non_synth)

        stream.output("}")

    def prepare_children(self, value_non_synth: lldb.SBValue):
        if value_non_synth.MightHaveChildren():
            return StructChildrenProvider(value_non_synth)
        return super().prepare_children(value_non_synth)


class LambdaVisDescriptor(StructVisDescriptor):
    def __init__(self, value_type: lldb.SBType, lambda_name: str):
        super(LambdaVisDescriptor, self).__init__(value_type)
        self.lambda_name = lambda_name

    def output_summary(self, value_non_synth: lldb.SBValue, stream: Stream):
        value_type: SBType = value_non_synth.GetType()
        num = value_type.GetNumberOfMemberFunctions()
        for i in range(num):
            func: SBTypeMemberFunction = value_type.GetMemberFunctionAtIndex(i)
            func_name = func.GetName()
            if func_name != 'operator()':
                continue
            func_type: SBType = func.GetType()
            signature = func_type.GetName()
            if signature:
                stream.output('<lambda> ')
                stream.output(signature)
                stream.output(' ')

        super(LambdaVisDescriptor, self).output_summary(value_non_synth, stream)
