from rpython.jit.metainterp.history import AbstractDescr, getkind
from rpython.jit.codewriter.flatten import Register, Label, TLabel, KINDS
from rpython.jit.codewriter.flatten import ListOfKind, IndirectCallTargets
from rpython.jit.codewriter.format import format_assembler
from rpython.jit.codewriter.jitcode import SwitchDictDescr, JitCode
from rpython.jit.codewriter import heaptracker, longlong
from rpython.rlib.objectmodel import ComputedIntSymbolic
from rpython.flowspace.model import Constant
from rpython.rtyper.lltypesystem import lltype, llmemory, rffi
from rpython.rtyper import rclass


class AssemblerError(Exception):
    pass


class Assembler(object):

    def __init__(self):
        self.insns = {}
        self.descrs = []
        self.indirectcalltargets = set()    # set of JitCodes
        self.list_of_addr2name = []
        self._descr_dict = {}
        self._count_jitcodes = 0
        self._seen_raw_objects = set()

    def assemble(self, ssarepr, jitcode=None):
        """Take the 'ssarepr' representation of the code and assemble
        it inside the 'jitcode'.  If jitcode is None, make a new one.
        """
        self.setup(ssarepr.name)
        ssarepr._insns_pos = []
        for insn in ssarepr.insns:
            ssarepr._insns_pos.append(len(self.code))
            self.write_insn(insn)
        self.fix_labels()
        self.check_result()
        if jitcode is None:
            jitcode = JitCode(ssarepr.name)
        jitcode._ssarepr = ssarepr
        self.make_jitcode(jitcode)
        if self._count_jitcodes < 20:    # stop if we have a lot of them
            jitcode._dump = format_assembler(ssarepr)
        self._count_jitcodes += 1
        return jitcode

    def setup(self, name):
        self.code = []
        self.constants_dict = {}
        self.constants_i = []
        self.constants_r = []
        self.constants_f = []
        self.label_positions = {}
        self.tlabel_positions = []
        self.switchdictdescrs = []
        self.count_regs = dict.fromkeys(KINDS, 0)
        self.liveness = {}
        self.startpoints = set()
        self.alllabels = set()
        self.resulttypes = {}
        self.ssareprname = name

    def emit_reg(self, reg):
        if reg.index >= self.count_regs[reg.kind]:
            self.count_regs[reg.kind] = reg.index + 1
        self.code.append(chr(reg.index))

    def emit_const(self, const, kind, allow_short=False):
        value = const.value
        if kind == 'int':
            TYPE = const.concretetype
            if isinstance(TYPE, lltype.Ptr):
                assert TYPE.TO._gckind == 'raw'
                self.see_raw_object(value)
                value = llmemory.cast_ptr_to_adr(value)
                TYPE = llmemory.Address
            if TYPE == llmemory.Address:
                value = heaptracker.adr2int(value)
            if TYPE is lltype.SingleFloat:
                value = longlong.singlefloat2int(value)
            if not isinstance(value, (llmemory.AddressAsInt,
                                      ComputedIntSymbolic)):
                value = lltype.cast_primitive(lltype.Signed, value)
                if allow_short:
                    try:
                        short_num = -128 <= value <= 127
                    except TypeError:    # "Symbolics cannot be compared!"
                        short_num = False
                    if short_num:
                        # emit the constant as a small integer
                        self.code.append(chr(value & 0xFF))
                        return True
            constants = self.constants_i
        elif kind == 'ref':
            value = lltype.cast_opaque_ptr(llmemory.GCREF, value)
            constants = self.constants_r
        elif kind == 'float':
            if const.concretetype == lltype.Float:
                value = longlong.getfloatstorage(value)
            else:
                assert longlong.is_longlong(const.concretetype)
                value = rffi.cast(lltype.SignedLongLong, value)
            constants = self.constants_f
        else:
            raise AssemblerError('unimplemented %r in %r' %
                                 (const, self.ssareprname))
        key = (kind, Constant(value))
        if key not in self.constants_dict:
            constants.append(value)
            val = 256 - len(constants)
            assert val >= 0, "too many constants"
            self.constants_dict[key] = val
        # emit the constant normally, as one byte that is an index in the
        # list of constants
        self.code.append(chr(self.constants_dict[key]))
        return False

    def write_insn(self, insn):
        if insn[0] == '---':
            return
        if isinstance(insn[0], Label):
            self.label_positions[insn[0].name] = len(self.code)
            return
        if insn[0] == '-live-':
            key = len(self.code)
            live_i, live_r, live_f = self.liveness.get(key, ("", "", ""))
            live_i = self.get_liveness_info(live_i, insn[1:], 'int')
            live_r = self.get_liveness_info(live_r, insn[1:], 'ref')
            live_f = self.get_liveness_info(live_f, insn[1:], 'float')
            self.liveness[key] = live_i, live_r, live_f
            return
        startposition = len(self.code)
        self.code.append("temporary placeholder")
        #
        argcodes = []
        allow_short = (insn[0] in USE_C_FORM)
        for x in insn[1:]:
            if isinstance(x, Register):
                self.emit_reg(x)
                argcodes.append(x.kind[0])
            elif isinstance(x, Constant):
                kind = getkind(x.concretetype)
                is_short = self.emit_const(x, kind, allow_short=allow_short)
                if is_short:
                    argcodes.append('c')
                else:
                    argcodes.append(kind[0])
            elif isinstance(x, TLabel):
                self.alllabels.add(len(self.code))
                self.tlabel_positions.append((x.name, len(self.code)))
                self.code.append("temp 1")
                self.code.append("temp 2")
                argcodes.append('L')
            elif isinstance(x, ListOfKind):
                itemkind = x.kind
                lst = list(x)
                assert len(lst) <= 255, "list too long!"
                self.code.append(chr(len(lst)))
                for item in lst:
                    if isinstance(item, Register):
                        assert itemkind == item.kind
                        self.emit_reg(item)
                    elif isinstance(item, Constant):
                        assert itemkind == getkind(item.concretetype)
                        self.emit_const(item, itemkind)
                    else:
                        raise NotImplementedError("found in ListOfKind(): %r"
                                                  % (item,))
                argcodes.append(itemkind[0].upper())
            elif isinstance(x, AbstractDescr):
                if x not in self._descr_dict:
                    self._descr_dict[x] = len(self.descrs)
                    self.descrs.append(x)
                if isinstance(x, SwitchDictDescr):
                    self.switchdictdescrs.append(x)
                num = self._descr_dict[x]
                assert 0 <= num <= 0xFFFF, "too many AbstractDescrs!"
                self.code.append(chr(num & 0xFF))
                self.code.append(chr(num >> 8))
                argcodes.append('d')
            elif isinstance(x, IndirectCallTargets):
                self.indirectcalltargets.update(x.lst)
            elif x == '->':
                assert '>' not in argcodes
                argcodes.append('>')
            else:
                raise NotImplementedError(x)
        #
        opname = insn[0]
        if '>' in argcodes:
            assert argcodes.index('>') == len(argcodes) - 2
            self.resulttypes[len(self.code)] = argcodes[-1]
        key = opname + '/' + ''.join(argcodes)
        num = self.insns.setdefault(key, len(self.insns))
        self.code[startposition] = chr(num)
        self.startpoints.add(startposition)

    def get_liveness_info(self, prevlives, args, kind):
        """Return a string whose characters are register numbers.
        We sort the numbers, too, to increase the chances of duplicate
        strings (which are collapsed into a single string during translation).
        """
        lives = set(prevlives)    # set of characters
        for reg in args:
            if isinstance(reg, Register) and reg.kind == kind:
                lives.add(chr(reg.index))
        return lives

    def fix_labels(self):
        for name, pos in self.tlabel_positions:
            assert self.code[pos  ] == "temp 1"
            assert self.code[pos+1] == "temp 2"
            target = self.label_positions[name]
            assert 0 <= target <= 0xFFFF
            self.code[pos  ] = chr(target & 0xFF)
            self.code[pos+1] = chr(target >> 8)
        for descr in self.switchdictdescrs:
            as_dict = {}
            for key, switchlabel in descr._labels:
                target = self.label_positions[switchlabel.name]
                as_dict[key] = target
            descr.attach(as_dict)

    def check_result(self):
        # Limitation of the number of registers, from the single-byte encoding
        assert self.count_regs['int'] + len(self.constants_i) <= 256
        assert self.count_regs['ref'] + len(self.constants_r) <= 256
        assert self.count_regs['float'] + len(self.constants_f) <= 256

    def make_jitcode(self, jitcode):
        jitcode.setup(''.join(self.code),
                      self.constants_i,
                      self.constants_r,
                      self.constants_f,
                      self.count_regs['int'],
                      self.count_regs['ref'],
                      self.count_regs['float'],
                      liveness=self.liveness,
                      startpoints=self.startpoints,
                      alllabels=self.alllabels,
                      resulttypes=self.resulttypes)

    def see_raw_object(self, value):
        if value._obj not in self._seen_raw_objects:
            self._seen_raw_objects.add(value._obj)
            if not value:    # filter out NULL pointers
                return
            TYPE = lltype.typeOf(value).TO
            if isinstance(TYPE, lltype.FuncType):
                name = value._obj._name
            elif TYPE == rclass.OBJECT_VTABLE:
                if not value.name:    # this is really the "dummy" class
                    return            #   pointer from some dict
                name = ''.join(value.name.chars)
            else:
                return
            addr = llmemory.cast_ptr_to_adr(value)
            self.list_of_addr2name.append((addr, name))

    def finished(self, callinfocollection):
        # Helper called at the end of assembling.  Registers the extra
        # functions shown in _callinfo_for_oopspec.
        for func in callinfocollection.all_function_addresses_as_int():
            func = heaptracker.int2adr(func)
            self.see_raw_object(func.ptr)


# A set of instructions that use the 'c' encoding for small constants.
# Allowing it anywhere causes the number of instruction variants to
# expode, growing past 256.  So we list here only the most common
# instructions where the 'c' variant might be useful.
USE_C_FORM = set([
    'copystrcontent',
    'getarrayitem_gc_pure_i',
    'getarrayitem_gc_pure_r',
    'getarrayitem_gc_i',
    'getarrayitem_gc_r',
    'goto_if_not_int_eq',
    'goto_if_not_int_ge',
    'goto_if_not_int_gt',
    'goto_if_not_int_le',
    'goto_if_not_int_lt',
    'goto_if_not_int_ne',
    'int_add',
    'int_and',
    'int_copy',
    'int_eq',
    'int_ge',
    'int_gt',
    'int_le',
    'int_lt',
    'int_ne',
    'int_return',
    'int_sub',
    'jit_merge_point',
    'new_array',
    'new_array_clear',
    'newstr',
    'setarrayitem_gc_i',
    'setarrayitem_gc_r',
    'setfield_gc_i',
    'strgetitem',
    'strsetitem',

    'foobar', 'baz',    # for tests
])