import sys, os

sys.path.append(os.path.join(os.path.dirname(__file__), ".."))

from sc2.game_data import GameData
from sc2.game_info import GameInfo
from sc2.game_info import Ramp
from sc2.game_state import GameState
from sc2.bot_ai import BotAI
from sc2.units import Units
from sc2.unit import Unit
from sc2.position import Point2, Point3, Size, Rect
from sc2.game_data import Cost

from sc2.ids.unit_typeid import UnitTypeId
from sc2.ids.upgrade_id import UpgradeId
from sc2.ids.ability_id import AbilityId
from sc2.ids.buff_id import BuffId
from sc2.ids.effect_id import EffectId

from sc2.data import Race

import pickle, pytest, random, math, lzma
from hypothesis import given, event, settings, strategies as st

from typing import Iterable


"""
You can execute this test running the following command from the root python-sc2 folder:
pipenv run pytest test/test_pickled_data.py

This test/script uses the pickle files located in "python-sc2/test/pickle_data" generated by "generate_pickle_files_bot.py" file, which is a bot that starts a game on each of the maps defined in the main function.

It will load the pickle files, recreate the bot object from scratch and tests most of the bot properties and functions.
All functions that require some kind of query or interaction with the API directly will have to be tested in the "autotest_bot.py" in a live game.
"""


def get_map_specific_bots() -> Iterable[BotAI]:
    folder = os.path.dirname(__file__)
    subfolder_name = "pickle_data"
    pickle_folder_path = os.path.join(folder, subfolder_name)
    files = os.listdir(pickle_folder_path)
    for file in (f for f in files if f.endswith(".xz")):
        with lzma.open(os.path.join(folder, subfolder_name, file), "rb") as f:
            raw_game_data, raw_game_info, raw_observation = pickle.load(f)

        # Build fresh bot object, and load the pickle'd data into the bot object
        bot = BotAI()
        game_data = GameData(raw_game_data.data)
        game_info = GameInfo(raw_game_info.game_info)
        game_state = GameState(raw_observation)
        bot._initialize_variables()
        bot._prepare_start(client=None, player_id=1, game_info=game_info, game_data=game_data)
        bot._prepare_step(state=game_state, proto_game_info=raw_game_info)

        yield bot


# From https://docs.pytest.org/en/latest/example/parametrize.html#a-quick-port-of-testscenarios
def pytest_generate_tests(metafunc):
    idlist = []
    argvalues = []
    for scenario in metafunc.cls.scenarios:
        idlist.append(scenario[0])
        items = scenario[1].items()
        argnames = [x[0] for x in items]
        argvalues.append(([x[1] for x in items]))
    metafunc.parametrize(argnames, argvalues, ids=idlist, scope="class")


class TestClass:
    # Load all pickle files and convert them into bot objects from raw data (game_data, game_info, game_state)
    scenarios = [(bot_obj.game_info.map_name, {"bot": bot_obj}) for bot_obj in get_map_specific_bots()]

    def test_main_base_ramp(self, bot: BotAI):
        bot._game_info.map_ramps, bot._game_info.vision_blockers = bot._game_info._find_ramps_and_vision_blockers()
        # Test if main ramp works for all spawns
        game_info: GameInfo = bot._game_info

        for spawn in bot._game_info.start_locations + [bot.townhalls.random.position]:
            # Remove cached ramp
            if hasattr(bot, "cached_main_base_ramp"):
                del bot.cached_main_base_ramp
            # Set start location as one of the opponent spawns
            bot._game_info.player_start_location = spawn
            # Find main base ramp for opponent
            ramp: Ramp = bot.main_base_ramp
            # On the map HonorgroundsLE, the main base is large and it would take a bit of effort to fix, so it returns None or empty set
            if len(ramp.upper) in {2, 5}:
                assert ramp.upper2_for_ramp_wall
                # Check if terran wall was found
                assert ramp.barracks_correct_placement
                assert ramp.barracks_in_middle
                assert ramp.depot_in_middle
                assert len(ramp.corner_depots) == 2
                # Check if protoss wall was found
                assert ramp.protoss_wall_pylon
                assert len(ramp.protoss_wall_buildings) == 2
                assert ramp.protoss_wall_warpin
            else:
                # On maps it is unable to find valid wall positions (Honorgrounds LE) it should return None, empty sets or empty lists
                assert ramp.barracks_correct_placement is None
                assert ramp.barracks_in_middle is None
                assert ramp.depot_in_middle is None
                assert ramp.corner_depots == set()
                assert ramp.protoss_wall_pylon is None
                assert ramp.protoss_wall_buildings == []
                assert ramp.protoss_wall_warpin is None
            assert ramp.top_center
            assert ramp.bottom_center
            assert ramp.size
            assert ramp.points
            assert ramp.upper
            assert ramp.lower
            # Test if ramp was detected far away
            distance = ramp.top_center.distance_to(bot._game_info.player_start_location)
            assert (
                distance < 30
            ), f"Distance from spawn to main ramp was detected as {distance:.2f}, which is too far. Spawn: {spawn}, Ramp: {ramp.top_center}"
