from collections import defaultdict, deque
from contextlib import contextmanager
from copy import deepcopy
from random import random
from string import ascii_lowercase as lowercase

from robots.constants import Action, City, Energy, ENERGY_BY_ACTION
from robots.state import GameState
from robots.utils import iter_board

@contextmanager
def protected(value, name):
    value_copy = deepcopy(value)
    try:
        yield value_copy
    finally:
        if value_copy != value:
            raise RuntimeError('Protected value %s was modified.' % name)

def extract_spawn_points(map_):
    points = []

    for x, y, cell in iter_board(map_):
        if cell in lowercase:
            cell = lowercase.index(cell)
        try:
            n = int(cell)
        except ValueError:
            pass
        else:
            points.append((n, (x, y)))
            map_[y][x] = City.NORMAL

    points.sort()
    return [point for n, point in points]

class Game:
    def __init__(self, map_, victory_by_combat=True):
        self.bots = {}
        self.time = 0
        self.victory_by_combat = victory_by_combat

        map_ = deepcopy(map_)
        self.spawn_points = deque(extract_spawn_points(map_))

        self.state = GameState(map_)

    def add_bot(self, bot, name=None):
        if name is None:
            name = bot.__name__

        if name in self.bots:
            raise KeyError(
                'There is already a bot named %r in this game!' %
                (name,)
            )

        self.bots[name] = bot

        # Place the initial robot for this player.
        try:
            x, y = self.spawn_points.popleft()
        except IndexError:
            raise IndexError('Not enough spawn points in the map.')
        self.state.robots_by_player[name] = [
            (x, y, Energy.INITIAL),
        ]

    @property
    def finished(self):
        target = self.state.n_cities_to_win

        winners = []
        roboteers = []

        n_allegiances_by_player = self.state.n_allegiances_by_player

        for name, robots in self.state.robots_by_player.items():
            n_cities = n_allegiances_by_player[name]
            if n_cities >= target:
                winners.append(name)

            if robots:
                roboteers.append(name)

        won = bool(winners)
        if len(roboteers) <= 1 and self.victory_by_combat:
            won = True
            winners.extend(roboteers)

        return winners or won

    def call_bots(self):
        kill_players = []
        for whoami, bot in self.bots.items():
            try:
                with protected(self.state, 'state') as state:
                    result = bot(whoami, state)

                assert isinstance(result, (str, list, tuple)), \
                    'Bot did not return a str/list/tuple, got %r' % (result,)
                result = str(''.join(result))

                my_robots = self.state.robots_by_player[whoami]
                assert len(result) == len(my_robots), (
                    'Bot did not return an action for all robots '
                    '(%d actions != %d robots)' % (len(result), len(my_robots))
                )

                for action in result:
                    assert action in Action.all, (
                        'Got unexpected action %r' % (action,)
                    )

            except Exception:
                # TODO: log this exception
                kill_players.append(whoami)
#                traceback.print_exc()
#                time.sleep(1)

            else:
                for bot, action in zip(my_robots, result):
                    yield whoami, bot, action

        for player in kill_players:
            self.state.robots_by_player[player][:] = []

    def apply_actions(self, actions):
        robots_by_player = defaultdict(list)

        for whoami, bot, action in actions:
            x, y, energy = bot

            if energy < ENERGY_BY_ACTION[action]:
                action = Action.NOTHING

            if action == Action.PROMOTE:
                if self.state.allegiances.get((x, y)) != whoami:
                    self.state.allegiances[x, y] = whoami
                else:
                    action = Action.NOTHING

            elif action != Action.NOTHING:
                nx, ny = self.state.expected_position(x, y, action)

                if x == nx and y == ny:
                    action = Action.NOTHING
                x, y = nx, ny

            energy -= ENERGY_BY_ACTION[action]

            robots_by_player[whoami].append((x, y, energy))

        for name, robots in self.state.robots_by_player.items():
            robots[:] = robots_by_player[name]

    def factories_produce_robots(self):
        state_robots = self.state.robots

        for x, y in self.state.factories:
            if state_robots[(x, y)]:
                continue

            owner = self.state.allegiances.get((x, y))
            if owner is None:
                continue

            if random() < 0.2:
                self.state.robots_by_player[owner].append(
                    (x, y, Energy.INITIAL)
                )

    def robots_gain_energy(self):
        for robots in self.state.robots_by_player.values():
            robots[:] = [
                (x, y, min(Energy.MAXIMUM, energy + Energy.GAIN))
                for x, y, energy in robots
            ]

    def resolve_combats(self):
        to_remove = set()
        to_append = defaultdict(list)

        for (x, y), robots in self.state.robots.items():
            if len(robots) <= 1:
                # No combat here.
                continue

            to_remove.add((x, y))

            # each player's strength is the sum of their robots' energy
            strength_by_player = defaultdict(int)
            total = 0
            for player, energy in robots:
                strength_by_player[player] += energy
                total += energy

            winner_count = 0
            for player, strength in strength_by_player.items():
                # subtract the sum of the other players' strength
                strength -= total - strength
                # remove with strength players <= 0
                if strength > 0:
                    energy = min(strength, Energy.MAXIMUM)
                    to_append[player].append((x, y, energy))
                    winner_count += 1

            # we should end up with one player.
            # (and I wrote a proof of that, so something's seriously wrong if
            # this fails)
            assert winner_count <= 1

        for player, robots in self.state.robots_by_player.items():
            extra_robots = to_append[player]
            robots[:] = [
                (x, y, energy)
                for x, y, energy in robots
                if (x, y) not in to_remove
            ] + extra_robots

    def __iter__(self):
        while not self.finished:
            yield self.state
            self.next()
            if self.time > 400:
                break
        yield self.state

    def next(self):
        # 1. You give commands you your robots.
        actions = self.call_bots()
        # 2. They perform the requested commands.
        self.apply_actions(actions)
        # 3. Factories produce new robots.
        self.factories_produce_robots()
        # 4. Combats are resolved.
        self.resolve_combats()
        # 5. Each robot gains energy, up to the maximum.
        self.robots_gain_energy()

        self.time += 1