
import logging
import os
import random
import subprocess
import sys
import traceback
from io import StringIO
from pathlib import Path

import zmq

from . import layout
from .exceptions import PlayerDisconnected, PlayerTimeout
from .layout import layout_as_str, BOT_I2N
from .network import ZMQClientError, ZMQConnection, ZMQReplyTimeout, ZMQUnreachablePeer

_logger = logging.getLogger(__name__)


def _ensure_list_tuples(list):
    """ Ensures that an iterable is a list of position tuples. """
    return [tuple(item) for item in list]

def _ensure_tuple_tuples(set):
    """ Ensures that an iterable is a tuple of position tuples. """
    return tuple(sorted(tuple(item) for item in set))


def create_homezones(shape, walls):
    width, height = shape
    return [
        tuple((x, y) for x in range(0, width // 2)
                     for y in range(0, height) if (x, y) not in walls),
        tuple((x, y) for x in range(width // 2, width)
                     for y in range(0, height) if (x, y) not in walls)
    ]


class Team:
    """
    Wraps a move function and forwards it the `set_initial`
    and `get_move` requests.

    The Team class holds the team’s state between turns, the team’s
    random number generator and the bot track (resets every time a bot
    is killed).

    Parameters
    ----------
    team_move : function with (bot, state) -> position
        the team’s move function
    team_name :
        the name of the team (optional)
    """
    def __init__(self, team_move, *, team_name=""):
        self._team_move = team_move
        self.team_name = team_name

        #: Storage for the team state
        self._state = {}

        #: The team’s random number generator
        self._random = None

        #: The history of bot positions
        self._bot_track = [[], []]


    def set_initial(self, team_id, game_state):
        """ Sets the bot indices for the team and returns the team name.
        Currently, we do not call _set_initial on the user side.

        Parameters
        ----------
        team_id : int
            The id of the team
        game_state : dict
            The initial game state

        Returns
        -------
        Team name : string
            The name of the team

        """
        # Reset the team state
        self._state.clear()

        # Initialize the random number generator
        # with the seed that we received from game
        self._random = random.Random(game_state['seed'])

        # Reset the bot tracks
        self._bot_track = [[], []]

        # Store the walls, which are only transmitted once
        self._walls = _ensure_tuple_tuples(game_state['walls'])

        # Store the shape, which is only transmitted once
        self._shape = tuple(game_state['shape'])

        # Cache the initial positions so that we don’t have to calculate them at each step
        self._initial_positions = layout.initial_positions(self._walls, self._shape)

        # Cache the homezone so that we don’t have to create it at each step
        self._homezone = create_homezones(self._shape, self._walls)

        return self.team_name

    def get_move(self, game_state):
        """ Requests a move from the Player who controls the Bot with id `bot_id`.

        This method returns a dict with a key `move` and a value specifying the direction
        in a tuple. Additionally, a key `say` can be added with a textual value.

        Parameters
        ----------
        game_state : dict
            The initial game state

        Returns
        -------
        move : dict
        """
        me = make_bots(walls=self._walls,
                       shape=self._shape,
                       initial_positions=self._initial_positions,
                       homezone=self._homezone,
                       team=game_state['team'],
                       enemy=game_state['enemy'],
                       round=game_state['round'],
                       bot_turn=game_state['bot_turn'],
                       rng=self._random)

        team = me._team

        for idx, mybot in enumerate(team):
            # If a bot has been killed, we reset its bot track
            if mybot.was_killed:
                self._bot_track[idx] = []

        # Add our track
        if len(self._bot_track[me._bot_turn]) == 0:
            self._bot_track[me._bot_turn] = [me.position]

        for idx, mybot in enumerate(team):
            # If the track of any bot is empty,
            # add its current position
            if me._bot_turn != idx:
                self._bot_track[idx].append(mybot.position)

            mybot.track = self._bot_track[idx][:]

        try:
            # request a move from the current bot
            move = self._team_move(team[me._bot_turn], self._state)

            # check that the returned value is a position tuple
            try:
                if len(move) != 2:
                    raise ValueError(f"Function move did not return a valid position: got {move} instead.")
            except TypeError:
                # Convert to ValueError
                raise ValueError(f"Function move did not return a valid position: got {move} instead.") from None
        except Exception as e:
            # Our client had an exception. We print a traceback and
            # return the type of the exception to the server.
            # If this is a remote player, then this will be detected in pelita_player
            # and pelita_player will close the connection automatically.
            traceback.print_exc()
            return {
                "error": (type(e).__name__, str(e)),
            }

        return {
            "move": move,
            "say": me._say
        }

    def _exit(self):
        """ Dummy function. Only needed for `RemoteTeam`. """
        pass

    def __repr__(self):
        return f'Team({self._team_move!r}, {self.team_name!r})'


class RemoteTeam:
    """ Start a child process with the given `team_spec` and handle
    communication with it through a zmq.PAIR connection.

    It also does some basic checks for correct return values and tries to
    terminate the child process once it is not needed anymore.

    Parameters
    ----------
    team_spec
        The string to pass as a command line argument to pelita_player
        or the address of a remote player
    team_name
        Overrides the team name
    zmq_context
        A zmq_context (if None, a new one will be created)
    idx
        The team index (currently only used to specify the team’s color)
    store_output
        If store_output is a string it will be interpreted as a path to a
        directory where to store stdout and stderr for the client processes.
        It helps in debugging issues with the clients.
        In the special case of store_output==subprocess.DEVNULL, stdout of
        the remote clients will be suppressed.
    """
    def __init__(self, team_spec, *, team_name=None, zmq_context=None, idx=None, store_output=False):
        if zmq_context is None:
            zmq_context = zmq.Context()

        self._team_spec = team_spec
        self._team_name = team_name

        #: Default timeout for a request, unless specified in the game_state
        self._request_timeout = 3

        if team_spec.startswith('remote:tcp://'):
            # We connect to a remote player that is listening
            # on the given team_spec address.
            # We create a new DEALER socket and send a single
            # REQUEST message to the remote address.
            # The remote player will then create a new instance
            # of a player and forward all of our zmq traffic
            # to that player.

            # remove the "remote:" part from the address
            send_addr = team_spec[len("remote:"):]
            address = "tcp://*"
            self.bound_to_address = address

            socket = zmq_context.socket(zmq.DEALER)
            socket.setsockopt(zmq.LINGER, 0)
            socket.connect(send_addr)
            _logger.info("Connecting zmq.DEALER to remote player at {}.".format(send_addr))
            socket.send_json({"REQUEST": address})
            self.proc = None

        else:
            # We bind to a local tcp port with a zmq PAIR socket
            # and start a new subprocess of pelita_player.py
            # that includes the address of that socket and the
            # team_spec as command line arguments.
            # The subprocess will then connect to this address
            # and load the team.

            socket = zmq_context.socket(zmq.PAIR)
            port = socket.bind_to_random_port('tcp://*')
            self.bound_to_address = f"tcp://localhost:{port}"
            if idx == 0:
                color='blue'
            elif idx == 1:
                color='red'
            else:
                color=''
            self.proc = self._call_pelita_player(team_spec, self.bound_to_address,
                                                 color=color, store_output=store_output)

        self.zmqconnection = ZMQConnection(socket)

    def _call_pelita_player(self, team_spec, address, color='', store_output=False):
        """ Starts another process with the same Python executable,
        the same start script (pelitagame) and runs `team_spec`
        as a standalone client on URL `addr`.
        """
        player = 'pelita.scripts.pelita_player'
        external_call = [sys.executable,
                         '-m',
                         player,
                         team_spec,
                         address,
                         '--color',
                         color]

        _logger.debug("Executing: %r", external_call)
        if store_output == subprocess.DEVNULL:
            return (subprocess.Popen(external_call, stdout=store_output), None, None)
        elif store_output:
            store_path = Path(store_output)
            stdout = (store_path / f"{color or team_spec}.out").open('w')
            stderr = (store_path / f"{color or team_spec}.err").open('w')

            # We must run in unbuffered mode to enforce flushing of stdout/stderr,
            # otherwise we may lose some of what is printed
            proc = subprocess.Popen(external_call, stdout=stdout, stderr=stderr,
                                    env=dict(os.environ, PYTHONUNBUFFERED='x'))
            return (proc, stdout, stderr)
        else:
            return (subprocess.Popen(external_call), None, None)

    @property
    def team_name(self):
        if self._team_name is not None:
            return self._team_name

        try:
            msg_id = self.zmqconnection.send("team_name", {})
            team_name = self.zmqconnection.recv_timeout(msg_id, self._request_timeout)
            if team_name:
                self._team_name = team_name
            return team_name
        except ZMQReplyTimeout:
            _logger.info("Detected a timeout, returning a string nonetheless.")
            return "%error%"
        except ZMQUnreachablePeer:
            _logger.info("Detected a DeadConnection, returning a string nonetheless.")
            return "%error%"

    def set_initial(self, team_id, game_state):
        timeout_length = game_state['timeout_length']
        try:
            msg_id = self.zmqconnection.send("set_initial", {"team_id": team_id,
                                                    "game_state": game_state})
            team_name = self.zmqconnection.recv_timeout(msg_id, timeout_length)
            if team_name:
                self._team_name = team_name
            return team_name
        except ZMQReplyTimeout:
            # answer did not arrive in time
            raise PlayerTimeout()
        except ZMQUnreachablePeer:
            _logger.info("Could not properly send the message. Maybe just a slow client. Ignoring in set_initial.")
        except ZMQClientError as e:
            error_message = e.message
            error_type = e.error_type
            _logger.warning(f"Client connection failed ({error_type}): {error_message}")
            raise PlayerDisconnected(*e.args) from None

    def get_move(self, game_state):
        timeout_length = game_state['timeout_length']
        try:
            msg_id = self.zmqconnection.send("get_move", {"game_state": game_state})
            reply = self.zmqconnection.recv_timeout(msg_id, timeout_length)
            # make sure it is a dict
            reply = dict(reply)
            if "error" in reply:
                return reply
            # make sure that the move is a tuple
            reply["move"] = tuple(reply.get("move"))
            return reply
        except ZMQReplyTimeout:
            # answer did not arrive in time
            raise PlayerTimeout()
        except TypeError:
            # if we could not convert into a tuple or dict (e.g. bad reply)
            return None
        except ZMQUnreachablePeer:
            # if the remote connection is closed
            raise PlayerDisconnected()
        except ZMQClientError:
            raise

    def _exit(self):
        # We only want to exit once.
        if getattr(self, '_sent_exit', False):
            return
        try:
            # TODO: make zmqconnection stateful. set flag when already disconnected
            # For now, we simply check the state of the socket so that we do not send
            # over an already closed socket.
            if self.zmqconnection.socket.closed:
                return
            self.zmqconnection.send("exit", {}, timeout=1)
            self._sent_exit = True
        except ZMQUnreachablePeer:
            _logger.info("Remote Player %r is already dead during exit. Ignoring.", self)

    def __del__(self):
        try:
            self._exit()
            if self.proc:
                self.proc[0].terminate()
        except AttributeError:
            # in case we exit before self.proc or self.zmqconnection have been set
            pass

    def __repr__(self):
        team_name = f" ({self._team_name})" if self._team_name else ""
        return f"RemoteTeam<{self._team_spec}{team_name} on {self.bound_to_address}>"


def make_team(team_spec, team_name=None, zmq_context=None, idx=None, store_output=False):
    """ Creates a Team object for the given team_spec.

    If no zmq_context is passed for a remote team, then a new context
    will be automatically created and returned. Otherwise, the same
    zmq_context (or None) will be returned.

    Parameters
    ----------
    team_spec : callable or str
        A move function or a team_spec that is passed on to pelita_player

    team_name : str, optional
        Optional team name for a local team

    zmq_context : zmq context, optional
        ZMQ context to avoid having to create a new context for every team

    Returns
    -------
    team_player, zmq_context : tuple
        The team class to interact with
        The new ZMQ context

    """
    if callable(team_spec):
        _logger.info("Making a local team for %s", team_spec)
        # wrap the move function in a Team
        if team_name is None:
            team_name = f'local-team ({team_spec.__name__})'
        team_player = Team(team_spec, team_name=team_name)
    elif isinstance(team_spec, str):
        _logger.info("Making a remote team for %s", team_spec)
        # set up the zmq connections and build a RemoteTeam
        if not zmq_context:
            zmq_context = zmq.Context()
        team_player = RemoteTeam(team_spec=team_spec, zmq_context=zmq_context, idx=idx, store_output=store_output)
    else:
        raise TypeError(f"Not possible to create team from {team_spec} (wrong type).")

    return team_player, zmq_context


class Bot:
    def __init__(self, *, bot_index,
                          is_on_team,
                          position,
                          initial_position,
                          walls,
                          shape,
                          homezone,
                          food,
                          score,
                          kills,
                          deaths,
                          was_killed,
                          random,
                          round,
                          bot_char,
                          is_blue,
                          team_name,
                          error_count,
                          is_noisy,
                          bot_turn=None):
        self._bots = None
        self._say = None

        #: The previous positions of this bot including the current one.
        self.track = []

        #: Is this a friendly bot?
        self._is_on_team = is_on_team

        self.random = random
        self.position = tuple(position)
        self._initial_position = tuple(initial_position)
        self.walls = walls

        self.homezone = homezone
        self.food = food
        self.shape = shape
        self.score  = score
        self.kills = kills
        self.deaths = deaths
        self.was_killed = was_killed
        self._bot_index  = bot_index
        self.round = round
        self.char = bot_char
        self.is_blue = is_blue
        self.team_name = team_name
        self.error_count = error_count
        self.is_noisy = is_noisy

        # Attributes for Bot
        if self._is_on_team:
            assert bot_turn is not None
            self._bot_turn = bot_turn

    # @property
    # def legal_directions(self):
        # """ The legal moves that the bot can make from its current position,
        # including no move at all.
        # """
        # legal_directions = [(0, 0)]

        # for move in [(-1, 0), (1, 0), (0, 1), (0, -1)]:
            # new_pos = (self.position[0] + move[0], self.position[1] + move[1])
            # if not new_pos in self.walls:
                # legal_directions.append(move)

        # return legal_directions

    @property
    def legal_positions(self):
        """ The legal positions that the bot can reach from its current position,
        including the current position.
        """
        legal_positions = []

        for direction in [(0, 0), (-1, 0), (1, 0), (0, 1), (0, -1)]:
            new_pos = (self.position[0] + direction[0], self.position[1] + direction[1])
            if not new_pos in self.walls:
                legal_positions.append(new_pos)

        return legal_positions

    @property
    def _team(self):
       """ Both of our bots.
       """
       if self._is_on_team:
           return self._bots['team']
       else:
           return self._bots['enemy']

    @property
    def turn(self):
        """ The turn of our bot. """
        return self._bot_index % 2

    @property
    def other(self):
        """ The other bot in our team. """
        return self._team[1 - self.turn]

    @property
    def enemy(self):
        """ The list of enemy bots
        """
        if not self._is_on_team:
            return self._bots['team']
        else:
            return self._bots['enemy']

    def say(self, text):
        """ Print some text in the graphical interface. """
        self._say = text

    # def get_direction(self, position):
        # """ Return the direction needed to get to the given position.

        # Raises
        # ======
        # ValueError
            # If the position cannot be reached by a legal move
        # """
        # direction = (position[0] - self.position[0], position[1] - self.position[1])
        # if direction not in self.legal_directions:
            # raise ValueError("Cannot reach position %s (would have been: %s)." % (position, direction))
        # return direction

    # def get_position(self, direction):
        # """ Return the position reached with the given direction

        # Raises
        # ======
        # ValueError
            # If the direction is not legal.
        # """
        # if direction not in self.legal_directions:
            # raise ValueError(f"Direction {direction} is not legal.")
        # position = (direction[0] + self.position[0], direction[1] + self.position[1])
        # return position


    def _repr_html_(self):
        """ Jupyter-friendly representation. """
        bot = self
        width, height = bot.shape

        with StringIO() as out:
            out.write("<table>")
            for y in range(height):
                out.write("<tr>")
                for x in range(width):
                    if (x, y) in bot.walls:
                        bg = 'style="background-color: {}"'.format(
                            "rgb(94, 158, 217)" if x < width // 2 else
                            "rgb(235, 90, 90)")
                    else:
                        bg = ""
                    out.write("<td %s>" % bg)
                    if (x, y) in bot.walls: out.write("#")
                    if (x, y) in bot.food: out.write('<span style="color: rgb(247, 150, 213)">●</span>')
                    if (x, y) in bot.enemy[0].food: out.write('<span style="color: rgb(247, 150, 213)">●</span>')
                    for idx in range(2):
                        if bot._team[idx].position == (x, y):
                            if idx == bot.turn:
                                out.write('<b>' + str(idx) + '</b>')
                            else:
                                out.write(str(idx))
                    for idx in range(2):
                        if bot.enemy[idx].position == (x, y):
                            out.write(str(idx))

                    out.write("</td>")
                out.write("</tr>")
            out.write("</table>")
            return out.getvalue()

    def __str__(self):
        bot = self

        if bot.is_blue:
            blue = bot if not bot.turn else bot.other
            red = bot.enemy[0]
        else:
            blue = bot.enemy[0]
            red = bot if not bot.turn else bot.other

        bot_positions = [blue.position, red.position, blue.other.position, red.other.position]
        bot_noise = [blue.is_noisy, red.is_noisy, blue.other.is_noisy, red.other.is_noisy]

        header = ("{blue}{you_blue} vs {red}{you_red}.\n" +
                  "Playing on {col} side. Current turn: {turn}. "+
                  "Bot: {bot_char}. Round: {round}, score: {blue_score}:{red_score}. " +
                  "timeouts: {blue_timeouts}:{red_timeouts}\n").format(
            blue=blue.team_name,
            red=red.team_name,
            turn=bot.turn,
            round=bot.round,
            bot_char=bot.char,
            blue_score=blue.score,
            red_score=red.score,
            col="blue" if bot.is_blue else "red",
            you_blue=" (you)" if bot.is_blue else "",
            you_red=" (you)" if not bot.is_blue else "",
            blue_timeouts=blue.error_count,
            red_timeouts=red.error_count,
        )

        footer = ("Bots: {bots}\nNoisy: {noise}\nFood: {food}\n").format(
                  bots={BOT_I2N[idx]:pos for idx, pos in enumerate(bot_positions)},
                  noise={BOT_I2N[idx]:state for idx, state in enumerate(bot_noise)},
                  food = bot.food + bot.enemy[0].food)

        with StringIO() as out:
            out.write(header)

            layout = layout_as_str(walls=bot.walls,
                                   food=bot.food + bot.enemy[0].food,
                                   bots=bot_positions,
                                   shape=bot.shape)

            out.write(str(layout))
            out.write(footer)
            return out.getvalue()


# def __init__(self, *, bot_index, position, initial_position, walls, homezone, food, is_noisy, score, random, round, is_blue):
def make_bots(*, walls, shape, initial_positions, homezone, team, enemy, round, bot_turn, rng):
    bots = {}

    team_index = team['team_index']
    enemy_index = enemy['team_index']

    team_initial_positions = initial_positions[team_index::2]
    enemy_initial_positions = initial_positions[enemy_index::2]

    team_bots = []
    for idx, position in enumerate(team['bot_positions']):
        b = Bot(bot_index=idx,
            is_on_team=True,
            score=team['score'],
            deaths=team['deaths'][idx],
            kills=team['kills'][idx],
            was_killed=team['bot_was_killed'][idx],
            is_noisy=False,
            error_count=team['error_count'],
            food=_ensure_list_tuples(team['food']),
            walls=walls,
            shape=shape,
            round=round,
            bot_turn=bot_turn,
            bot_char=BOT_I2N[team_index + idx*2],
            random=rng,
            position=team['bot_positions'][idx],
            initial_position=team_initial_positions[idx],
            is_blue=team_index % 2 == 0,
            homezone=homezone[team_index],
            team_name=team['name'])
        b._bots = bots
        team_bots.append(b)

    enemy_bots = []
    for idx, position in enumerate(enemy['bot_positions']):
        b = Bot(bot_index=idx,
            is_on_team=False,
            score=enemy['score'],
            kills=enemy['kills'][idx],
            deaths=enemy['deaths'][idx],
            was_killed=enemy['bot_was_killed'][idx],
            is_noisy=enemy['is_noisy'][idx],
            error_count=enemy['error_count'],
            food=_ensure_list_tuples(enemy['food']),
            walls=walls,
            shape=shape,
            round=round,
            bot_char = BOT_I2N[team_index + idx*2],
            random=rng,
            position=enemy['bot_positions'][idx],
            initial_position=enemy_initial_positions[idx],
            is_blue=enemy_index % 2 == 0,
            homezone=homezone[enemy_index],
            team_name=enemy['name'])
        b._bots = bots
        enemy_bots.append(b)

    bots['team'] = team_bots
    bots['enemy'] = enemy_bots
    return team_bots[bot_turn]

