
from collections import namedtuple
import logging
from ..tools.recursivedescent import RecursiveDescentParser
from . import nodes


class LlvmIrParser(RecursiveDescentParser):
    """ Recursive descent parser for llvm-ir

    Closely modeled after LLParser.cpp at:
    https://github.com/llvm-mirror/llvm/blob/master/lib/AsmParser/LLParser.cpp
    """
    logger = logging.getLogger('llvm-parse')

    def __init__(self, context):
        super().__init__()
        self.context = context
        self._pfs = None
        self.module = None
        self.numbered_vals = []

    def parse_module(self):
        """ Parse a module """
        self.logger.debug('Parsing module')
        self.module = nodes.Module(self.context)
        self.forward_refs = {}
        while not self.at_end:
            if self.peak == 'define':
                self.parse_define()
            elif self.peak == 'declare':
                self.parse_declare()
            elif self.peak == 'target':
                self.parse_target_definition()
            elif self.peak == 'LID':
                self.parse_unnamed_type()
            elif self.peak == 'LID':
                self.parse_named_type()
            elif self.peak == 'GID':
                self.parse_named_global()
            elif self.peak == 'attributes':
                self.parse_unnamed_attr_group()
            elif self.peak == 'MDVAR':
                self.parse_named_metadata()
            elif self.peak == '!':
                self.parse_standalone_metadata()
            else:  # pragma: no cover
                self.error('TODO')
                raise NotImplementedError(str(self.peak))
        return self.module

    def parse_unnamed_type(self):
        """ Parse LocalVarID '=' 'type' type """
        name = self.consume('LID')
        print(name)
        self.consume('=')
        self.consume('type')
        self.parse_struct_definition()

    def parse_named_type(self):
        """ Parse LocalVar '=' 'type' type """
        self.consume('=')
        self.consume('type')
        self.parse_struct()

    def parse_define(self):
        """ Parse a function """
        self.consume('define')
        function = self.parse_function_header()
        self.parse_function_body(function)

    def parse_named_global(self):
        self.consume('GID')
        self.consume('=')
        self.parse_optional_linkage()
        self.parse_global()

    def parse_global(self):
        is_constant = self.parse_global_type()
        ty = self.parse_type()
        self.parse_global_value(ty)
        print(is_constant)

    def parse_global_type(self):
        if self.has_consumed('constant'):
            return True
        elif self.has_consumed('global'):
            return False
        else:
            self.error("Expected 'global' or 'constant'")

    def parse_function_body(self, function):
        self._pfs = PerFunctionState(self, function)
        self.consume('{')
        while self.peak != '}':
            self.parse_basic_block()
        self.consume('}')
        self._pfs = None

    def parse_declare(self):
        """ Parse a function declaration """
        self.consume('declare')
        function = self.parse_function_header()
        print(function)

    def parse_function_header(self):
        return_type = self.parse_type()
        name = self.parse_name(global_name=True)
        arg_list = self.parse_argument_list()
        attributes = self.parse_fn_attribute_value_pairs()
        print(attributes)
        param_types = [a.type for a in arg_list]
        if self.peak == 'ATTRID':
            self.consume('ATTRID')
        function_type = nodes.FunctionType.get(return_type, param_types)
        function = nodes.Function.create(function_type, name, self.module)
        for argument, arg_info in zip(function.arguments, arg_list):
            if arg_info.name:
                argument.set_name(arg_info.name)
            else:
                # The parameter had no name, so number it?
                name = '%{}'.format(len(self.numbered_vals))
                argument.set_name(name)
                self.numbered_vals.append((name, argument))
        return function

    def parse_target_definition(self):
        """ Parse a target top level entity """
        self.consume('target')
        if self.peak == 'triple':
            self.consume('triple')
            self.consume('=')
            val = self.parse_string_constant()
        else:
            self.consume('datalayout')
            self.consume('=')
            val = self.parse_string_constant()
            self.module.data_layout = nodes.DataLayout.from_string(val)
        return val

    def parse_unnamed_attr_group(self):
        self.consume('attributes')
        attr_id = self.consume('ATTRID')
        self.consume('=')
        self.consume('{')
        attributes = self.parse_fn_attribute_value_pairs()
        print(attributes)
        self.consume('}')
        return attr_id

    def parse_fn_attribute_value_pairs(self):
        attributes = []
        while True:
            if self.peak in ['nounwind', 'sspstrong', 'uwtable', 'readonly']:
                a = self.consume().val
                attributes.append(a)
            elif self.peak == 'STR':
                key = self.parse_string_constant()
                if self.has_consumed('='):
                    val = self.parse_string_constant()
                else:
                    val = None
                attributes.append((key, val))
            else:
                break
        return attributes

    def create_global_fwd_ref(self, module, pty, name):
        if isinstance(pty.el_type, nodes.FunctionType):
            self.not_impl()
        else:
            return nodes.GlobalVariable(pty.el_type, name, module=module)

    def get_global_val(self, name, ty):
        """ Get a global with the given id """
        if not isinstance(ty, nodes.PointerType):
            self.error('global variable reference must have pointer type')

        if name in self.module.vmap:
            val = self.module.vmap[name]
        elif name in self.forward_refs:
            val = self.forward_refs[name]
        else:
            print(self.module, ty, name)
            val = self.create_global_fwd_ref(self.module, ty, name)
            self.forward_refs[name] = val
        return val

    def parse_named_metadata(self):
        """ Parse meta data starting with '!my.var'  """
        self.consume('MDVAR')
        self.consume('=')
        self.consume('!')
        self.consume('{')
        while self.peak != '}':
            self.consume()
        self.consume('}')

    def parse_standalone_metadata(self):
        """ Parse meta data starting with '!' 'num' """
        self.consume('!')
        self.parse_number()
        self.consume('=')
        self.consume('!')
        self.parse_md_node_vector()

    def parse_md_node_vector(self):
        self.consume('{')
        elements = []
        while self.peak != '}':
            self.consume()
            # TODO: parse meta data!
            # if elements:
            #    self.consume(',')
            # metadata = self.parse_metadata()
            # elements.append(metadata)
        self.consume('}')
        return elements

    def parse_metadata(self):
        """ Parse meta data

        !42
        !{...}
        !"string"

        """
        if self.peak == '!':
            self.consume('!')
        else:
            md = 1
        return md

    def parse_string_constant(self):
        return self.consume('STR').val

    def parse_basic_block(self):
        if self.peak == 'LBL':
            label = self.consume('LBL').val
        else:
            label = None
        bb = self._pfs.define_bb(label)

        # Parse instructions until terminator
        while True:
            if self.peak == 'LID':
                name = self.parse_name()
                self.consume('=')
            else:
                name = None

            instruction = self.parse_instruction()
            bb.instructions.append(instruction)
            if name:
                self._pfs.set_instruction_name(name, instruction)
            if instruction.is_terminator:
                break
        return bb

    def parse_argument_list(self):
        """ Parse '(' ... ')' """
        self.consume('(')
        args = []
        if self.peak == ')':
            pass
        else:
            args.append(self.parse_arg())
            while self.has_consumed(','):
                args.append(self.parse_arg())
        self.consume(')')
        return args

    def parse_arg(self):
        ty = self.parse_type()
        attrs = self.parse_optional_param_attrs()
        print(attrs)
        if self.peak == 'LID':
            name = self.consume('LID').val
        else:
            name = None
        return ArgInfo(ty, name)

    def parse_optional_param_attrs(self):
        attrs = []
        while True:
            if self.peak in ['nocapture', 'nonnull']:
                attr = self.consume().val
                attrs.append(attr)
            else:
                break
        return attrs

    def parse_optional_linkage(self):
        if self.has_consumed('private'):
            pass
        else:
            pass

    def parse_type(self):
        if self.peak == 'type':
            typ = self.consume('type').val
        elif self.peak == '{':
            typ = self.parse_anon_struct_type(packed=False)
        elif self.peak == '[':
            typ = self.parse_array_vector_type(False)
        elif self.peak == '<':
            typ = self.parse_array_vector_type(True)
        elif self.peak == 'LID':
            lid = self.consume('LID')
            # named_types[lid]
            print(lid)
            typ = nodes.StructType.get(self.context, [], False)
        else:  # pragma: no cover
            self.not_impl()

        assert isinstance(typ, nodes.Type), str(typ)

        # Parse suffix:
        while True:
            if self.peak == '*':
                self.consume('*')
                typ = nodes.PointerType.get_unequal(typ)
            elif self.peak == '(':
                typ = self.parse_function_type(typ)
            else:
                break
        return typ

    def parse_function_type(self, result):
        arg_list = self.parse_argument_list()
        return nodes.FunctionType.get(result, arg_list)

    def parse_anon_struct_type(self, packed=False):
        elts = self.parse_struct_body()
        return nodes.StructType.get(self.context, elts, packed)

    def parse_struct_definition(self):
        body = self.parse_struct_body()
        print(body)

    def parse_struct_body(self):
        """ Parse struct body.

        Can be either of:
        '{' '}'
        '{' type (',' type)* '}'
        """
        body = []
        self.consume('{')
        if self.has_consumed('}'):
            return body
        ty = self.parse_type()
        body.append(ty)
        while self.has_consumed(','):
            ty = self.parse_type()
            body.append(ty)
        self.consume('}')
        return body

    def parse_array_vector_type(self, is_vector):
        if is_vector:
            self.consume('<')
        else:
            self.consume('[')
        size = self.parse_number()
        self.consume('x')
        ty = self.parse_type()
        if is_vector:
            self.consume('>')
            return nodes.VectorType.get(ty, size)
        else:
            self.consume(']')
            return nodes.ArrayType.get(ty, size)

    def parse_number(self):
        return self.consume('NUMBER').val

    def parse_name(self, global_name=False):
        if global_name:
            return self.consume('GID').val
        else:
            return self.consume('LID').val

    def parse_global_value(self, ty):
        v = self.parse_val_id()
        v2 = self.convert_val_id_to_value(ty, v)
        return v2

    def parse_value(self, ty):
        """ Parse a value with a certain type """
        v = self.parse_val_id()
        v2 = self.convert_val_id_to_value(ty, v)
        return v2

    def parse_val_id(self):
        """ Some more abstract value parsing """
        if self.peak == 'zeroinitializer':
            self.consume('zeroinitializer')
            v = ValId('zero', None)
        elif self.peak == 'LID':
            v = ValId('local', self.parse_name())
        elif self.peak == 'GID':
            v = ValId('global', self.parse_name(global_name=True))
        elif self.peak == 'NUMBER':
            v = ValId('int', self.consume('NUMBER').val)
        elif self.peak == 'HEXDOUBLE':
            v = ValId('float', self.consume('HEXDOUBLE').val)
        elif self.peak == 'undef':
            self.consume('undef')
            v = ValId('undef', None)
        elif self.peak == 'true':
            self.consume('true')
            v = ValId('constant', nodes.ConstantInt.get_true(self.context))
        elif self.peak == 'false':
            self.consume('false')
            v = ValId('constant', nodes.ConstantInt.get_false(self.context))
        elif self.peak == '{':
            self.consume('{')
            elts = self.parse_global_value_vector()
            self.consume('}')
            v = ValId('constant_struct', elts)
        elif self.peak == '<':
            # '<' constvect '>'
            self.consume('<')
            elts = self.parse_global_value_vector()
            self.consume('>')
            v = ValId('constant', nodes.ConstantVector.get(elts))
        elif self.peak == 'getelementptr':
            self.consume('getelementptr')
            in_bounds = self.has_consumed('inbounds')
            print(in_bounds)
            self.consume('(')
            ty = self.parse_type()
            print(ty)
            self.consume(',')
            elts = self.parse_global_value_vector()
            self.consume(')')
            v = ValId('constant_gep', elts)
        else:  # pragma: no cover
            self.not_impl()
        assert isinstance(v, ValId), str(v)
        return v

    def convert_val_id_to_value(self, ty, val_id):
        if val_id.kind == 'local':
            v = self._pfs.get_val(val_id.val, ty)
        elif val_id.kind == 'global':
            v = self.get_global_val(val_id.val, ty)
        elif val_id.kind == 'int':
            if not ty.is_integer:
                self.error('integer constant must have integer type')
            v = nodes.ConstantInt.get(ty, val_id.val)
        elif val_id.kind == 'float':
            if not ty.is_floating_point:
                self.error('Floating point constant invalid for type')
            v = nodes.ConstantFP.get(ty, val_id.val)
        elif val_id.kind == 'zero':
            v = nodes.Constant.get_null_value(ty)
        elif val_id.kind == 'undef':
            v = nodes.UndefValue.get(ty)
        elif val_id.kind == 'constant':
            v = val_id.val
            if v.ty is not ty:
                self.error('constant expression type mismatch')
        else:
            raise NotImplementedError(str(val_id))
        return v

    def parse_type_and_value(self):
        ty = self.parse_type()
        val = self.parse_value(ty)
        return val

    def parse_type_and_basic_block(self):
        value = self.parse_type_and_value()
        if not isinstance(value, nodes.BasicBlock):
            self.error('Expected a basic block')
        return value

    def parse_global_value_vector(self):
        """ """
        v = []
        v.append(self.parse_global_type_and_value())
        while self.has_consumed(','):
            v.append(self.parse_global_type_and_value())
        return v

    def parse_global_type_and_value(self):
        ty = self.parse_type()
        val = self.parse_global_value(ty)
        return val

    # Instructions
    def parse_instruction(self):
        """ Parse a single instruction """
        if self.peak == 'store':
            instruction = self.parse_store()
        elif self.peak == 'call':
            instruction = self.parse_call()
        elif self.peak == 'ret':
            instruction = self.parse_ret()
        elif self.peak == 'br':
            instruction = self.parse_br()
        elif self.peak == 'alloca':
            instruction = self.parse_alloca()
        elif self.peak == 'load':
            instruction = self.parse_load()
        elif self.peak == 'extractelement':
            instruction = self.parse_extract_element()
        elif self.peak == 'insertelement':
            instruction = self.parse_insert_element()
        elif self.peak == 'shufflevector':
            instruction = self.parse_shuffle_vector()
        elif self.peak in [
                'add', 'fadd', 'sub', 'fsub', 'mul', 'fmul',
                'udiv', 'sdiv', 'fdiv', 'urem', 'srem',
                'shl', 'lshr', 'ashr']:
            instruction = self.parse_arithmatic()
        elif self.peak in ['and', 'or']:
            instruction = self.parse_arithmatic()
        elif self.peak in [
                'sext', 'zext',
                'trunc', 'fptrunc',
                'uitofp', 'fptoui', 'sitofp', 'fptosi',
                'ptrtoint', 'inttoptr']:
            instruction = self.parse_cast()
        elif self.peak == 'select':
            instruction = self.parse_select()
        elif self.peak in ['icmp', 'fcmp']:
            instruction = self.parse_compare()
        elif self.peak == 'getelementptr':
            instruction = self.parse_getelementptr()
        else:  # pragma: no cover
            self.not_impl()
        return instruction

    def parse_arithmatic(self):
        op = self.consume().val
        lhs = self.parse_type_and_value()
        self.consume(',')
        rhs = self.parse_value(lhs.ty)
        return nodes.BinaryOperator.create(op, lhs, rhs)

    def parse_cast(self):
        op = self.consume()
        val = self.parse_type_and_value()
        self.consume('to')
        dest_ty = self.parse_type()
        return nodes.CastInst.create(op, val, dest_ty)

    def parse_select(self):
        self.consume('select')
        op0 = self.parse_type_and_value()
        self.consume(',')
        op1 = self.parse_type_and_value()
        self.consume(',')
        op2 = self.parse_type_and_value()
        return nodes.SelectInst.create(op0, op1, op2)

    def parse_compare(self):
        icmp = self.peak == 'icmp'
        if icmp:
            self.consume('icmp')
        else:
            self.consume('fcmp')
        pred = self.parse_cmp_predicate(icmp)
        lhs = self.parse_type_and_value()
        self.consume(',')
        rhs = self.parse_value(lhs.ty)
        return nodes.ICmpInst(pred, lhs, rhs)

    def parse_cmp_predicate(self, icmp):
        if icmp:
            predicates = {
                'eq': nodes.CmpInst.ICMP_EQ,
                'ne': nodes.CmpInst.ICMP_NE,
                'slt': nodes.CmpInst.ICMP_SLT,
                'sgt': nodes.CmpInst.ICMP_SGT,
                'sle': nodes.CmpInst.ICMP_SLE,
                'sge': nodes.CmpInst.ICMP_SGE,
                'ult': nodes.CmpInst.ICMP_ULT,
                'ugt': nodes.CmpInst.ICMP_UGT,
                'ule': nodes.CmpInst.ICMP_ULE,
                'uge': nodes.CmpInst.ICMP_UGE,
            }
        else:
            predicates = {
                'ueq': nodes.CmpInst.FCMP_UEQ,
            }
        if self.peak in predicates:
            v = self.consume().val
            return predicates[v]
        else:
            self.error('Expected one of {}'.format(predicates.keys()))

    def parse_alloca(self):
        self.consume('alloca')
        ty = self.parse_type()
        size = 0
        alignment = 0
        if self.peak == ',':
            self.consume(',')
            self.consume('align')
            self.parse_number()
        return nodes.AllocaInst(ty, size, alignment)

    def parse_extract_element(self):
        self.consume('extractelement')
        op1 = self.parse_type_and_value()
        self.consume(',')
        op2 = self.parse_type_and_value()
        return nodes.ExtractElementInst(op1, op2)

    def parse_insert_element(self):
        self.consume('insertelement')
        op1 = self.parse_type_and_value()
        self.consume(',')
        op2 = self.parse_type_and_value()
        self.consume(',')
        op3 = self.parse_type_and_value()
        return nodes.InsertElementInst(op1, op2, op3)

    def parse_shuffle_vector(self):
        """ Parse the shuffle vector """
        self.consume('shufflevector')
        op1 = self.parse_type_and_value()
        self.consume(',')
        op2 = self.parse_type_and_value()
        self.consume(',')
        op3 = self.parse_type_and_value()
        return nodes.ShuffleVectorInst(op1, op2, op3)

    def parse_load(self):
        """ Parse load instruction

        'load' 'volatile'? typeandvalue (',' 'align' i32)?
        """
        self.consume('load')
        atomic = self.has_consumed('atomic')
        volatile = self.has_consumed('volatile')
        print(atomic, volatile)
        ty = self.parse_type()
        self.consume(',')
        val = self.parse_type_and_value()
        if self.has_consumed(','):
            self.consume('align')
            self.parse_number()
        return nodes.LoadInst(ty, val)

    def parse_store(self):
        """ Parse store instruction """
        self.consume('store')
        val = self.parse_type_and_value()
        self.consume(',')
        ptr = self.parse_type_and_value()
        if self.has_consumed(','):
            self.consume('align')
            self.parse_number()
        return nodes.StoreInst(val, ptr)

    def parse_getelementptr(self):
        """ parse the get element ptr (GEP) """
        self.consume('getelementptr')
        inbounds = self.has_consumed('inbounds')
        ty = self.parse_type()
        self.consume(',')
        ptr = self.parse_type_and_value()
        indices = []
        while self.has_consumed(','):
            val = self.parse_type_and_value()
            indices.append(val)
        gep = nodes.GetElementPtrInst(ty, ptr, indices)
        gep.inbounds = inbounds
        return gep

    def parse_call(self):
        self.consume('call')
        ret_type = self.parse_type()
        name = self.parse_name(global_name=True)
        args = self.parse_parameter_list()
        return nodes.CallInst(ret_type, name, args)

    def parse_parameter_list(self):
        self.consume('(')
        args = []
        while self.peak != ')':
            if args:
                self.consume(',')
            ty = self.parse_type()
            v = self.parse_value(ty)
            args.append(v)
        self.consume(')')
        return args

    def parse_ret(self):
        """ Parse return instruction """
        self.consume('ret')
        ty = self.parse_type()
        if ty.is_void:
            return nodes.ReturnInst(ty)
        else:
            val = self.parse_value(ty)
            return nodes.ReturnInst(ty, val)

    def parse_br(self):
        """ Parse a branch instruction """
        self.consume('br')
        op0 = self.parse_type_and_value()
        if isinstance(op0, nodes.BasicBlock):
            return nodes.BranchInst(op0)
        else:
            self.consume(',')
            op1 = self.parse_type_and_basic_block()
            self.consume(',')
            op2 = self.parse_type_and_basic_block()
            return nodes.BranchInst(op1, op2, op0)

    def parse_phi(self):
        """ Parse a phi instruction.

        For example 'phi' Type '[' Value ',' Value ']' (',' ...) *
        """
        self.parse_type()
        self.consume('[')
        self.parse_value()
        self.consume(',')
        self.parse_value()
        self.consume(']')
        while self.has_consumed(','):
            self.consume('[')
            self.parse_value()
            self.consume(',')
            self.parse_value()
            self.consume(']')


class PerFunctionState:
    """ Keep track of parsing stuff in a function """
    def __init__(self, parser, function):
        self.parser = parser
        self.function = function
        self.forward = {}

    def get_val(self, name, ty):
        """ Get a value with a name and a given type """
        # Lookup the name:
        if name in self.function.vmap:
            val = self.function.vmap[name]
        elif name in self.forward:
            val = self.forward[name]
        else:
            if ty.is_label:
                val = nodes.BasicBlock.create(
                    self.function.context, name, self.function)
            else:
                raise NotImplementedError(str(ty))
            self.forward[name] = val

        if ty is not val.ty:
            self.parser.error('{} != {}'.format(ty, val.ty))
        return val

    def get_bb(self, name):
        return self.get_val(name, self.function.context.label_ty)

    def define_bb(self, name):
        bb = self.get_bb(name)

        # Move basic block from forward to function:
        self.function.basic_blocks.append(bb)
        self.forward.pop(name)
        return bb

    def set_instruction_name(self, name, instruction):
        # This will update the symbol table:
        instruction.set_name(name)

        # Remove placeholder:
        if name in self.forward:
            sentinel = self.forward.pop(name)
            if sentinel.ty is not instruction.ty:
                self.parser.error(
                    'Instruction forward referenced with type {}'.format(
                        sentinel.ty))
            sentinel.replace_all_uses_with(instruction)


ValId = namedtuple('ValId', ['kind', 'val'])
ArgInfo = namedtuple('ArgInfo', ['type', 'name'])
