#!/usr/bin/env python3
import argparse
import time
from typing import Callable, Iterable, List, Optional, Tuple
from client import get_state, send_turn, set_log_level, TIME_BEFORE_RETRY


parser = argparse.ArgumentParser()
parser.add_argument("--server", type=str, default="http://localhost:5000", help="Server address.")
parser.add_argument("--game", type=str, default="main", help="'main' or 'test_#'.")
parser.add_argument("--token", type=str, default=None, help="API token.")
parser.add_argument("--log_level", type=str, default="WARNING", choices=["ERROR", "WARNING", "INFO"])
# you can add your own arguments, the strategy function will get them


Coords = Tuple[int, int]


class Member:
    def __init__(self, team: int, id: int, remaining_rounds: int):
        self.team = team
        self.id = id
        self.remaining_rounds = remaining_rounds
        self.action: Optional[str] = None


class Field:
    def __init__(self, home_for_team: Optional[int],
                 occupied_by_team: Optional[int],
                 hill: bool, members: List[Member]):
        self.home_for_team = home_for_team
        self.occupied_by_team = occupied_by_team
        self.hill = hill
        self.members = members


def strategy(team_id: int, state: dict, args: argparse.Namespace) -> dict:
    """Finds the best move for the given state.

    This function is called every round.
    It should return a dict with actions for each of team's soldiers.

    State format:
    ```yaml
    {
        map: [[{
            home_for_team: Optional[int],
            occupied_by_team: Optional[int],
            hill: bool,
            members: [{
                type: "soldier",
                team: int,
                id: int,
                remaining_rounds: int
            }]
        }]]
    }
    ```

    Turn format:
    ```yaml
    {
        members: [{
            id: int,
            action: Optional[str]
            # None, "left", "right", "up", "down"
        }]
    }
    ```
    """

    world = parse_world(state["map"])
    home = find_fields(world, lambda f: f.home_for_team == team_id)[0][1]
    members = find_members(world, lambda m: m.team == team_id)

    # TODO: implement your strategy here
    for member, coords in members:
        member.action = "up"
    return build_turn(map(lambda x: x[0], members))


def find_fields(
    world: List[List[Field]],
    predicate: Callable[[Field], bool]
) -> List[Tuple[Field, Coords]]:
    """Finds all fields that satisfy the given predicate."""

    result = []
    for y, row in enumerate(world):
        for x, field in enumerate(row):
            if predicate(field):
                result.append((field, (x, y)))
    return result


def find_members(
    world: List[List[Field]],
    predicate: Callable[[Member], bool]
) -> List[Tuple[Member, Coords]]:
    """Finds all members that satisfy the given predicate."""

    result = []
    for y, row in enumerate(world):
        for x, field in enumerate(row):
            for member in field.members:
                if predicate(member):
                    result.append((member, (x, y)))
    return result


def build_turn(members: Iterable[Member]) -> dict:
    return {
        "members": [
            {"id": member.id, "action": member.action}
            for member in members
        ]
    }


def parse_world(world: dict) -> List[List[Field]]:
    fields = []
    for row in world:
        fields_row = []
        for field in row:
            members = []
            for member in field["members"]:
                members.append(Member(
                    member["team"],
                    member["id"],
                    member["remaining_rounds"]
                ))
            fields_row.append(Field(
                field["home_for_team"],
                field["occupied_by_team"],
                field["hill"],
                members
            ))
        fields.append(fields_row)
    return fields


def main(args: argparse.Namespace):
    min_round = 0
    set_log_level(args.log_level)

    while True:
        state, wait_time = get_state(min_round, args)
        if state is None:
            # retry later
            time.sleep(wait_time)
            continue

        turn = strategy(state["team_id"], state["state"], args)

        round = state["round"]
        while not send_turn(turn, round, args):
            # if there was a connection error, retry later
            time.sleep(TIME_BEFORE_RETRY)
        min_round = round + 1


if __name__ == '__main__':
    args = parser.parse_args()
    main(args)