import collections
import copy
import re

from common import FrissError
import nml_blocks

class Node:

    def __init__(self, name, parent=None):
        self.name = name
        self.parent = parent
        self.children = []
        self.transitions = []

    def __eq__(self, other):
        return (self.name, self.parent) == (other.name, other.parent)

    def __hash__(self):
        return hash((self.name, self.parent))

    def add(self, parent, child, start_y, end_y, prob=None):
        for n in self.bfs():
            if n.name == parent:
                n.children.append(Node(child, parent))
                n.transitions.append(
                    {"start_y": start_y, "end_y": end_y, "prob": prob})
                return
        raise FrissError

    def bfs(self):
        yield self
        q = collections.deque(self.children)
        while q:
            node = q.popleft()
            yield node
            for c in node.children:
                q.append(c)


class RailtypeDb:

    def __init__(self):
        self.railtypes = {}
        # These trings will never be translated
        self.strings = {}

    def blocks(self):
        return nml_blocks.switch_stack + nml_blocks.item_stack

    def add_railtype(self, **kwargs):
        label = kwargs["label"]
        self.railtypes[label] = Railtype(**kwargs)

    def __getitem__(self, key):
        try:
            return self.railtypes[key]
        except KeyError:
            raise FrissError("No railtype with label {}".format(key))

    def __iter__(self):
        return self.railtypes.__iter__()


class Railtype:
    prop_strings = [
        "name", "toolbar_caption", "menu_text", "build_window_caption",
        "autoreplace_text", "new_engine_text"
    ]
    constant_railtypes = {
        "unelectrified": ["MONO", "RAIL", "SAAN", "SBAN", "SDAN", "SFAN",
                          "SGAN", "SABN", "SACN", "SADN", "SAEN"],
        "electrified": ["ELRL", "SAAA", "SBAA", "SCAA", "SEAA", "SHAA", "SIAA",
                       "SJAA", "SABA", "SACA", "SADA", "SAEA"],
        "third_rail": ["MONO", "3RDR", "MTRO", "SAA3", "SAB3", "SAC3",
                       "SAD3", "SAE3"],
    }
    for k, v in constant_railtypes.items():
        constant_railtypes[k] = ['"{}"'.format(label) for label in v]
    constant_railtypes["all"] = \
        [railtype for cat in constant_railtypes.values() for railtype in cat]

    def __getitem__(self, key):
        try:
            return self.data[key]
        except KeyError:
            raise FrissError("Railtype {} has no data with key {}".format(key))

    def __init__(self, **kwargs):
        self.data = kwargs
        # Create item block and necessary graphic blocks
        item_block = self.gen_item_block()
        item_block.add_graphics("fences", "fences")
        if self["ballast"] or self["sleepers"]:
            item_block.add_graphics("track_overlay",
                                    self.track_stack("overlay"))
            item_block.add_graphics("underlay",
                                    self.track_stack("underlay"))
            item_block.add_graphics("tunnels",
                                    self.track_stack("tunnel_underlay"))
            item_block.add_graphics("bridge_surfaces",
                                    self.track_stack("bridge"))
            item_block.add_graphics("tunnel_overlay", "tracks_tunnel_overlay")
        if self.data["catenary"]:
            item_block.add_graphics("catenary_pylons", self.catenary_stack())
            item_block.add_graphics("catenary_wire", "catenary_wires")
        if self.data["lc_rural"]:
            item_block.add_graphics("level_crossings", self.lc_stack())
        if self.data["signals"]:
            item_block.add_graphics("signals", self.signal_stack())

    def get(self, key, default=None):
        try:
            return self.data[key]
        except KeyError:
            return default

    def catenary_stack(self):
        nodes = copy.deepcopy(self.data["catenary"])
        for n in nodes.travel():
            n.name = "catenary_pylons_{}".format(n.name)
        return gen_switches(nodes)

    def gen_item_block(self):
        # This is full of terrible hacks to handle some parameters and
        # metro. I don't really know what else to do with this. It's
        # always the same when you try to cleverly parse CSVs but
        # end up writing 50% of the code for special cases anyways.
        # I guess it's better than manually writing thousands of lines
        # of unmaintainable NML.
        mode_check = "if (param_mode == 0) {{\n{}\n}}"
        item_name = "METR" if self["label"] == "3RDR" else self["label"]
        item = nml_blocks.Item(feature="FEAT_RAILTYPES", name=item_name)
        if item_name != "METR":
            item.add_decor(mode_check)
        if "replaced_by_metro" in self["extra"]:
            item.add_decor("if (param_metro == 0) {{\n{}\n}}")
        if item_name == "METR":
            item.add_decor("if (param_metro != 0) {{\n{}\n}}")
            flags = []
            if self["catenary"]:
                flags.append("RAILTYPE_FLAG_CATENARY")
            if not self["lc_rural"]:
                flags.append("RAILTYPE_FLAG_NO_LEVEL_CROSSING")
            item.add_prop("railtype_flags",
                          "bitmask({})".format('|'.join(flags)))
            for p in range(1, 3):
                dummy = nml_blocks.Item(feature="FEAT_RAILTYPES",
                                        name=item_name)
                dummy.add_decor(
                    "if (param_metro == {}) {{{{\n{{}}\n}}}}".format(p))
                if p == 1:
                    dummy.add_prop(
                        "speed_limit", "{} km/h".format(self["speed"]))
                else:
                    dummy.add_prop("speed_limit", "0")
        elif "high_speed_param" in self["extra"]:
            for p in range(0, 4):
                dummy = nml_blocks.Item(feature="FEAT_RAILTYPES",
                                        name=item_name)
                dummy.add_decor(
                    "if (param_high_speed == {}) {{{{\n{{}}\n}}}}".format(p))
                dummy.add_decor(mode_check)
                if p in {0, 1}:
                    dummy.add_prop(
                        "speed_limit", "{} km/h".format(self["speed"]))
                else:
                    dummy.add_prop("speed_limit", "0")
                flags = []
                if self["catenary"]:
                    flags.append("RAILTYPE_FLAG_CATENARY")
                if p in {1, 3}:
                    flags.append("RAILTYPE_FLAG_NO_LEVEL_CROSSING")
                if flags:
                    dummy.add_prop("railtype_flags",
                                   "bitmask({})".format('|'.join(flags)))
        else:
            item.add_prop("speed_limit", "{} km/h".format(self["speed"]))
            flags = []
            if self["catenary"]:
                flags.append("RAILTYPE_FLAG_CATENARY")
            if not self["lc_rural"]:
                flags.append("RAILTYPE_FLAG_NO_LEVEL_CROSSING")
            item.add_prop("railtype_flags",
                          "bitmask({})".format('|'.join(flags)))
        if item_name == "METR":
            item.add_prop("sort_order", "26")
        else:
            item.add_prop("sort_order", self["sort"])
        item.add_prop("label", '"{}"'.format(self["label"]))
        item.add_prop("introduction_date", "date({},01,01)".format(
            self["available"]))
        for k in self.prop_strings:
            item.add_prop(k, "string({})".format(self.get(k, "STR_FIXME")))
        item.add_prop("compatible_railtype_list", "[{}]".format(
            ','.join(self.constant_railtypes["all"])))
        if self["label"] == "3RDR":
            item.add_prop("powered_railtype_list", "[{}]".format(
                ','.join(self.constant_railtypes["third_rail"])))
            item.add_prop("alternative_railtype_list", "[{}]".format(
                ','.join(self.constant_railtypes["third_rail"])))
        elif self.data["catenary"]:
            item.add_prop("powered_railtype_list", "[{}]".format(
                ','.join(self.constant_railtypes["electrified"])))
            item.add_prop("alternative_railtype_list", "[{}]".format(
                ','.join(self.constant_railtypes["electrified"])))
        else:
            item.add_prop("powered_railtype_list", "[{}]".format(
                ','.join(self.constant_railtypes["all"])))
            item.add_prop("alternative_railtype_list", "[{}]".format(
                ','.join(self.constant_railtypes["unelectrified"])))
        item.add_prop("curve_speed_multiplier", "0")
        item.add_prop("station_graphics", "RAILTYPE_STATION_NORMAL")
        item.add_prop("construction_cost", self["pur_cost"])
        item.add_prop("acceleration_model", "ACC_MODEL_RAIL")
        # Don't use because ottd map shows legend for only one color
        #item.add_prop("map_colour", self.get("color", self["sort"]))
        #item.add_prop("requires_railtype_list", something)
        if self["introduces"]:
            item.add_prop("introduces_railtype_list", "[{}]".format(
                ','.join(['"{}"'.format(r) for r in self["introduces"]])))
        item.add_prop("maintenance_cost", self["maint_cost"])
        return item

    def lc_stack(self):
        lc_fmt = "lc_{}_{}"
        ballast_nodes = copy.deepcopy(self.data["ballast"])
        for bn in ballast_nodes.travel():
            rural_nodes = copy.deepcopy(self.data["lc_rural"])
            for n in rural_nodes.travel():
                n.name = lc_fmt.format(n.name, bn.name)
            city_nodes = copy.deepcopy(self.data["lc_city"])
            for n in city_nodes.travel():
                n.name = lc_fmt.format(n.name, bn.name)
            lc_root_rural = gen_switches(rural_nodes)
            lc_root_city = gen_switches(city_nodes)
            lc_switch = nml_blocks.Switch(expression="town_zone")
            lc_switch.add("TOWNZONE_EDGE", lc_root_rural)
            lc_switch.add("default", lc_root_city)
            bn.name = lc_switch
        return gen_switches(ballast_nodes)

    def signal_stack(self):
        nodes = copy.deepcopy(self.data["signals"])
        for n in nodes.travel():
            n.name = "signals_{}".format(n.name)
        stack = []
        stack.append(gen_switches(nodes))  # Map switches
        stack.append(gen_switches(nodes, phased=False))  # GUI switches
        stack.append(nml_blocks.Switch(  # Test if map or GUI
            expression="getbits(extra_callback_info1, 0, 8)"))
        stack[-1].add("0", stack[-3])
        stack[-1].add("0x10", stack[-2])
        if self["signal_group"]:
            stack.append(nml_blocks.Switch(
                expression="param_signal_group_{}_light".format(
                    self["signal_group"])))
            stack[-1].add("0", stack[-2])
            stack[-1].add("default", "signal_group_{}_light".format(
                self["signal_group"]))
            stack.append(nml_blocks.Switch(
                expression="getbits(extra_callback_info2, 8, 8)"))
            stack[-1].add("0", stack[-2])
            stack[-1].add("1", "signal_group_{}_sema".format(
                self["signal_group"]))
        return stack[-1]

    def track_stack(self, type_):
        # If no ballast or sleeper transition, then no need to generate
        # multiple transition switch stacks
        other = None
        if self["ballast"] and not self["sleepers"]:
            other = self.data["ballast"]
        elif not self["ballast"] and self["sleepers"]:
            other = self.data["sleepers"]
        if other:
            nodes = copy.deepcopy(other)
            for n in nodes.travel():
                n.name = "tracks_{}_{}".format(other.name, type_)
            return gen_switches(nodes)
        block_name = "tracks_{}_{}_{}"
        if len(list(self.data["ballast"].travel())) == 1:
            # Prepend all sleeper names with ballast name
            sleeper_nodes = copy.deepcopy(self.data["sleepers"])
            for n in sleeper_nodes.travel():
                n.name = block_name.format(self.data["ballast"].name,
                                           n.name, type_)
            return gen_switches(sleeper_nodes)
        if len(list(self.data["sleepers"].travel())) == 1:
            # Append all ballast names with sleeper name
            ballast_nodes = copy.deepcopy(self.data["ballast"])
            for n in ballast_nodes.travel():
                n.name = block_name.format(
                    n.name, self.data["sleepers"].name, type_)
            return gen_switches(ballast_nodes)
        # If both sleepers and ballast have transitions, then for each
        # ballast type, a corresponding transition switch stack for
        # sleepers must be generated
        ballast_nodes = copy.deepcopy(self.data["ballast"])
        for bn in ballast_nodes.travel():
            sleeper_nodes = copy.deepcopy(self.data["sleepers"])
            for sn in sleeper_nodes.travel():
                sn.name = block_name.format(bn.name, sn.name, type_)
            switches = gen_switches(sleeper_nodes)
            bn.name = switches
        return gen_switches(ballast_nodes)


def gen_switches(tree, phased=True):
    stack = {}
    for node in reversed(list(tree.travel())):
        if node.leaf:
            stack[node] = node.name
        elif not phased:
            assert len(node.children) == 1
            stack[node] = nml_blocks.Switch(expression="current_year")
            stack[node].add(
                "0..{}".format(node.transitions[0]["start_y"]-1),
                node.name)
            stack[node].add("default", stack[node.children[0]])
        else:
            prob_tot = 0
            random_stack = []
            for c, t in node:
                prob = t["prob"] // 25
                prob_tot += prob
                random_stack.append((prob, phased_transition(
                    start=node.name, end=stack[c],
                    start_y=t["start_y"], end_y=t["end_y"])))
            if prob_tot != 4:
                random_stack.append((4-prob_tot, node.name))
            # Only one 100% transition
            if len(random_stack) == 1:
                stack[node] = random_stack[0][1]
            else:
                stack[node] = nml_blocks.Switch(random=True)
                for p, s in random_stack:
                    stack[node].add(p, s)
    return stack[tree]  # "tree" is the root node


def even_year_ranges(start_y, end_y):
    if start_y == end_y:
        return []
    if 3 >= end_y - start_y >= 1:
        return [str(start_y + i) for i in range(end_y - start_y)]
    out = []
    quot, rem = divmod((end_y-start_y), 3)
    prev = start_y
    for _ in range(3):
        this_interval = quot
        if rem:
            this_interval += 1
            rem -= 1
        if this_interval > 1:
            out.append("{}..{}".format(prev, (prev+this_interval) - 1))
        else:
            out.append(str(prev))
        prev += this_interval
    return out


def phased_transition(start, end, start_y, end_y):
    years = even_year_ranges(start_y, end_y)
    if len(years) == 0:
        ret_sw = nml_blocks.Switch(expression="current_year")
        ret_sw.add("0..{}".format(end_y-1), start)
        ret_sw.add("default", end)
        return ret_sw
    if len(years) == 1:
        phases = [(2, 2)]
    elif len(years) == 2:
        phases = [(1, 3), (3, 1)]
    else:
        phases = [(1, 3), (2, 2), (3, 1)]
    temp = collections.deque()
    for y, p in zip(years, phases):
        temp.append(nml_blocks.Switch(random=True))
        temp[-1].add(p[0], end)
        temp[-1].add(p[1], start)
    ret_sw = nml_blocks.Switch(expression="current_year")
    ret_sw.add("0..{}".format(start_y-1), start)
    for y, p in zip(years, phases):
        ret_sw.add(y, temp.popleft())
    ret_sw.add("default", end)
    return ret_sw
