import collections
import copy
import os
import re
import sys

import cargos
import global_vars as g
import nml_blocks
import nml_props

cur_dir = os.path.dirname(__file__)

class ConsistDb:
    def __init__(self):
        self.db = {}
        self.make_list = []

    def __contains__(self, id_):
        if id_ in self.db:
            return True
        return False

    def __getitem__(self, id_):
        return self.db[id_]

    def __iter__(self):
        for id_ in sorted(self.db):
            yield self[id_]

    def add_consist(self, id_, **kwargs):
        assert id_ not in self.db
        self.db[id_] = Consist(id_, **kwargs)

    def output_lang(self):
        out = []
        for v in self:
            for s in v.strings:
                out.append('{} :{}'.format(s[0], s[1]))
        return '\n'.join(sorted(out))

    def output_make_rules(self):
        return '\n'.join(self.make_list)

    def output_nml(self):
        return '\n'.join([v.output_nml() for v in self if v.output_nml()])

    def output_sort_order(self):
        out = [c.id_ for c in sorted(self, 
                                     key=lambda x: x.props.get('sort_pos'))]
        return 'sort (FEAT_TRAINS, [{}]);'.format(', '.join(out))


class Consist:
    def __init__(self, id_, prop_map):
        self.id_ = id_
        self.props = nml_props.NmlProps(id_, prop_map)

        self.block_stack = []
        self.nml = ''
        self.resolved = False
        self.parts = []

    @property
    def is_arti(self):
        """Helper function to check if consist is articulated."""
        return True if len(self.parts) > 1 \
                    or any(int(p['length']) > 8 for p in self.parts) \
          else False

    @property
    def is_hacked_arti(self):
        return any(int(p["length"]) > 8 for p in self.parts)

    @property
    def block_names(self):
        return {b.name: b for b in self.block_stack}

    @property
    def strings(self):
        stringl = []
        if self.props.get('long_name'):
            pur_str = re.sub(r'\bSTR_PUR_\w+', '{STRING}',
                             self.props.get('long_name'))
            stringl.append(('STR_PUR_{}'.format(self.id_.upper()), pur_str))
        if self.props.get('xtra_text'):
            xtra_str = re.sub(r'\bSTR_ADD_\w+', '{STRING}',
                              self.props.get('xtra_text'))
            xtra_str = re.sub(r'\n', '{}', xtra_str)
            stringl.append(('STR_ADD_{}'.format(self.id_.upper()), xtra_str))
        return stringl

    def add_part(self, cbs):
        # Set some default values, length especially is mandatory
        part = {'cargo_capacity': self.props.get('capacity', 0),
                'colour_mapping': 'switch_recolour_lights', 'length': 8}
        if self.props.get('xtra_text') and len(self.parts) == 0:
            xtras = re.findall(r'\bSTR_ADD_\w+', self.props.get('xtra_text'))
            if xtras:
                xs = ', ' + ', '.join(['string({})'.format(s) for s in xtras])
            else:
                xs = ''
            part['additional_text'] = 'string(STR_ADD_{}{})'.format(
                self.id_.upper(), xs)
        part.update(cbs)
        self.parts.append(part)

    def add_block(self, block):
        # Make sure we didn't mess up and try add a block already created
        assert block not in self.block_stack
        assert block.name not in {b.name for b in self.block_stack}
        self.block_stack.append(block)

    def add_spritegroup(self, **kwargs):
        self.add_block(nml_blocks.SpriteGroup(**kwargs))

    def add_spriteset(self, **kwargs):
        self.add_block(nml_blocks.SpriteSet(**kwargs))

    def add_switch(self, **kwargs):
        self.add_block(nml_blocks.Switch(**kwargs))

    def output_nml(self):
        if not self.nml: # Stub props, no gfx
            out = []
            out.append(
                'item (FEAT_TRAINS, {}) {{ property {{'.format(self.id_))
            for p, v in sorted(self.props.get_props(nogfx=True)):
                out.append('{}{}: {};'.format(g.ind, p, v))
            out.append('}}')
            return '\n'.join(out)
        else:
            return self.nml

    def resolve(self):
        assert not self.resolved
        # Copy the original stack to prevent adding duplicate switches 
        # if this method is called multiple times on the same run
        # Item blocks

        # Hack the flipped gfx
        if len(self.parts) == 1 and int(self.parts[0]["length"]) < 8:
            part = self.parts[0]
            flipcp = self.copy_from(part["default"])
            for b in flipcp:
                if b in self.block_names and \
                        isinstance(self.block_names[b],
                            nml_blocks.SpriteSet):
                    self.block_names[b].flipped = True
            self.add_switch(name="gfx_flipped", var="var[0xC8]",
                            range_vals=[
                                ("0xFE..0xFF", flipcp[0]),
                                ("default", part["default"])])
            part["default"] = "gfx_flipped"

        items = []
        def add_item(id_, cbs):
            # If length is over 0, then it's articulated, and use
            # articulated prop mapping for the rest of the prop blocks
            props = self.props.get_props(len(items))
            items.append({'id': id_, 'cbs': cbs, 'props': props})

        invis_cbs = {'cargo_capacity': 0, 'default': 'spriteset_empty',
                     'length': 4, 'visual_effect_and_powered':
                     ('visual_effect_and_powered(\n'
                     'VISUAL_EFFECT_DISABLE, 0, DISABLE_WAGON_POWER)')}

        arti_part_sw = 'articulated'
        for n, part in enumerate(self.parts):
            if n > 0:
                item_name = '{}_part_{}'.format(self.id_, n)
            else:
                item_name = self.id_
            if int(part['length']) == 10:
                # Leading part
                cbs = part.copy()
                cbs.update(invis_cbs)
                if n == 0:
                    # Hack for depot gfx
                    depot_sw = 'gfx_depot_1'
                    depot_cp = self.copy_from(part["default"])
                    for b in depot_cp:
                        if b in self.block_names and \
                                isinstance(self.block_names[b],
                                    nml_blocks.SpriteSet):
                            self.block_names[b].depot = True
                    self.add_switch(name=depot_sw, var='extra_callback_info1',
                                    range_vals=[
                                        ('0x10', depot_cp[0]),
                                        ('default', 'spriteset_empty')])
                    cbs['default'] = depot_sw
                    # Arti part
                    cbs['articulated_part'] = arti_part_sw
                add_item(item_name, cbs)
                # Middle part
                part_name = item_name + '_fake_arti_1'
                cbs = part.copy()
                cbs['length'] = 2
                if n == 0:
                    # Hack for depot gfx
                    depot_sw = 'gfx_depot_2'
                    self.add_switch(name=depot_sw, var='extra_callback_info1',
                                    range_vals=[
                                        ('0x10', 'spriteset_empty'),
                                        ('default', part['default'])])
                    cbs['default'] = depot_sw
                add_item(part_name, cbs)
                # Trailing part
                part_name = item_name + '_fake_arti_2'
                cbs = part.copy()
                cbs.update(invis_cbs)
                add_item(part_name, cbs)
            else:
                cbs = part.copy()
                if n == 0 and self.is_arti:
                    cbs['articulated_part'] = arti_part_sw
                add_item(item_name, cbs)

        if len(items) > 1:
            arti_rvals = []
            for n, it in enumerate(items[1:]):
                arti_rvals.append((n+1, it['id']))
            arti_rvals.append((n+2, 'CB_RESULT_NO_MORE_ARTICULATED_PARTS'))
            self.add_switch(name=arti_part_sw, var='extra_callback_info1',
                            range_vals = arti_rvals)

        # Reference cb rvals to block objects, if possible
        for cbs in [i['cbs'] for i in items]:
            for cb, rval in cbs.items():
                if rval in self.block_names:
                    cbs[cb] = self.block_names[rval]

        # Reference return values of global blocks
        self.resolve_block_stack()
        self.fix_switches()

        # Switches, spritesets, etc. blocks that are global to consist
        out = []
        for block in self.block_stack:
            # Make sure there are no name collisions in final nml
            block.name = '{}_{}'.format(self.id_, block.name)
            out.append(block.output())

        gfx_blocks = []
        prop_blocks = []
        for item in items:
            gfx_blocks.append(
                'item (FEAT_TRAINS, {}) {{ graphics {{'.format(item['id']))
            for cb, rval in sorted(item['cbs'].items()):
                gfx_blocks.append('{}{}: {};'.format(g.ind, cb, rval))
            gfx_blocks.append('}}')

            prop_blocks.append(
                'item (FEAT_TRAINS, {}) {{ property {{'.format(item['id']))
            for p, v in sorted(item['props']):
                prop_blocks.append('{}{}: {};'.format(g.ind, p, v))
            prop_blocks.append('}}')

        out.extend(['\n'.join(gfx_blocks), '\n'.join(prop_blocks)])

        self.nml = '\n\n'.join(out) if any(out) else None

        self.resolved = True

    def copy_from(self, block):
        out = []
        stack = collections.deque([self.block_names[block]])
        seen = set()
        while stack:
            cur = stack.pop()
            cp = copy.deepcopy(cur)
            cp.name += "_copy"
            if cp.name not in out:
                self.add_block(cp)
                out.append(cp.name)
                if isinstance(cp, nml_blocks.ReturnBlock):
                    for rval in cp.rvals:
                        if rval.val in self.block_names:
                            stack.append(self.block_names[rval.val])
                            rval.val += "_copy"
        return out

    def fix_switches(self):
        """
        Some hacks to "fix" blocks if the vehicle is a longhack
        articulated veh.
        """
        # Fix position_in_consist from end
        if not self.is_arti:
            return
        for b in [b for b in self.block_stack if isinstance(
                b, nml_blocks.Switch)]:
            if 'position_in_consist_from_end' in b.var:
                for r, v in list(b.range_vals):
                    if str(r) == '0':
                        oldval = v
                        b.range_vals.remove((r, v))
                b.range_vals.append(('1', oldval))

    def resolve_block_stack(self):
        """
        Resolve connections from block return values and callbacks to blocks.
        
        """
        # Replace return values that point to other
        # blocks with references to block objects
        #
        # Find immediate links between blocks
        # (ie. block1 has a the name of block2 in return values
        # = block1 has block2 in it's links)
        links = {}
        for block in [b for b in self.block_stack \
                      if isinstance(b, nml_blocks.ReturnBlock)]:
            links[block] = set()
            for rval in block.rvals:
                if rval.val in self.block_names:
                    links[block].add(self.block_names[rval.val])
                    rval.val = self.block_names[rval.val]

        # Resolve block order
        n = len(self.block_stack)
        maxrec = (n*(n+1))/2 # Max recursion in worst case scenario
                             # avoid looping forever if we screw up
        temp_stack = []
        n = 0
        while self.block_stack:
            for b in list(self.block_stack): # Work on a copy
                if b not in links or links[b].issubset(set(temp_stack)):
                    temp_stack.append(b)
                    self.block_stack.remove(b)
            n += 1
            assert n <= maxrec

        # Replace old stack with ordered
        self.block_stack = temp_stack
