import logging
import uuid

from cooptools.timeTracker import TimeTracker
# from cooptools import taskProcessing as tsk
from cooptools import graphs as grph
from cooptools.physics import waypointFollower as wf
from cooptools.sectors import SectorTree
from voroni import *
from wave_collapse import *
from typing import Iterable
import math
import pprint
from cooptools import os_manip as osm
from cooptools.coopEnum import CoopEnum
from cooptools.toggles import BooleanToggleable
from examples.voroni_world.models import *
from behaviors import *
from input_controller import InputController
from behaviorTypes import BehaviorType
import cooptools.typeProviders as tp
import coopstorage.storage2 as store
from cooptools import randoms as crnd

logger = logging.getLogger(__name__)


class ClientType(CoopEnum):
    VILLAGER = auto()
    REGION = auto()


class RoadType(CoopEnum):
    DIRT_PATH = .9
    COBBLESTONE_ROAD = .8
    PAVED_ROAD = .5
    RAIL = .25


@timer(logger=logger, log_level=0, time_tracking_class=tt)
def _generate_regions(n, bounds) -> Tuple[List[Region], Voronoi, Dict[int, List[int]]]:
    pts = [pt for pt in vec.random_point_generator(n, bounds)]

    # generate voronoi diagram
    vor, vn_map, points, polys, border_map = generate_voronoi(points=pts)

    terrains = wave_collapse_terrains(vn_map)

    regions = []
    neighbor_cnxn_node_ids: Dict[frozenset, str] = {}


    for ii, pt in enumerate(pts):
        logging.debug(f"Adding region {ii + 1} / {len(pts)}")

        neigh_cnxns = {}
        for neigh in vn_map[ii]:
            if (id_neigh := frozenset([ii, neigh])) not in neighbor_cnxn_node_ids:
                neighbor_cnxn_node_ids[id_neigh] = f"{ii}<>{neigh}"

            neigh_cnxns[neigh] = neighbor_cnxn_node_ids[id_neigh]

        regions.append(
            Region(
                name=ii,
                pt=pt,
                terrain=terrains[ii],
                boundary=polys[ii],
                vert_node_ids=vor.regions[ii],
                neighbors={
                    x: NeighborAttrs(
                        shared_verts=border_map[frozenset([ii, x])],
                        cnxn_node_id=f"C:{neigh_cnxns[x]}",
                    ) for x in vn_map[ii]
                }
            )
        )

    return regions, vor, vn_map


class GameModel:

    def __init__(self,
                 input_controller: InputController,
                 max_villagers: int,
                 init_villagers: int,
                 villager_spawn_ms: int,
                 sector_tree_depth: int = 3,
                 seconds_per_hr: float = 1):
        self.input_controller = input_controller

        self.regions: Dict[int, Region] = None
        self.world_sector_map: SectorTree = None
        self.village_id: int = None
        self.world_graph: grph.Graph = None
        self.world_clock: TimeTracker = TimeTracker(as_async=False, start_on_init=True).start()
        self._seconds_per_hr = seconds_per_hr
        self._storage = store.Storage.from_meta(
            location_type_counts=[(store.dcs.LocationMeta(
                dims=(10, 10, 10),
                channel_processor=store.FIFOFlowChannelProcessor(),
                capacity=3
            ), 100)],
            naming_provider=lambda x: crnd.a_string(k=7)
        )

        self._last_villager_name = -1

        self.villager_spawn_timer = TimeTracker(timeout_ms=villager_spawn_ms,
                                                start_on_init=True,
                                                as_async=False,
                                                id="Villager Spawn Timer")
        self._villagers: Dict[int, Villager] = {}

        self._max_villagers = max_villagers
        self._n_init_villagers = init_villagers

        self._villager_waypoint_cache = {}

        self.crash_toggle: BooleanToggleable = BooleanToggleable(default=False)
        self.crash_regions: Tuple[Region, Region] = None

        self.rivers: grph.Graph = None
        # self._tasks = List[tsk.Task]
        self._sector_tree_depth = sector_tree_depth

    def choose_village_region(self) -> Region:
        village = next(x for x in self.shuffled_region_generator() if x.terrain in [
            TerrainType.PLAINS,
            TerrainType.GRASS,
            TerrainType.PLAINSHILLS,
            TerrainType.GRASSHILLS,
            TerrainType.BEACH
        ])

        return reveal_a_region(village, regions=self.regions, force=True)

    def create_villager(self):
        name = self._get_villager_name()
        self._villagers[name] = Villager(
            name=name).set_pos(self.Village.pt)

        self.world_sector_map.add_update_clients({self._villagers[name].id: ClientType.VILLAGER})
        self.villager_spawn_timer.reset(start=True)

    def init_villagers(self, n):
        logger.info(f"Initializing {n} villagers")
        self._villagers = {}
        self._last_villager_name = -1
        for ii in range(n):
            self.create_villager()

    def _build_boundary_graph(
            self,
            regions: Dict[int, Region]
    ) -> grph.Graph:
        logger.debug(f"Building Boundary Graph...")
        graph_dict = {}

        for id, region in regions.items():
            me = grph.Node(region.node_id, region.pt)

            for neigh, attrs in region.neighbors.items():
                if attrs.Boundary is not None:
                    cnxn_pos = attrs.Boundary.Midpoint.as_tuple()
                else:
                    cnxn_pos = tuple(float(x) for x in list(attrs.shared_verts.values())[0])

                shared_vert_nodes = []
                for shared_vert, vert_pos in attrs.shared_verts.items():
                    shared_vert_node = grph.Node(
                        name=f"V{shared_vert}",
                        pos=tuple(x for x in vert_pos)
                    )
                    shared_vert_nodes.append(shared_vert_node)

                cnxn_pos_node = grph.Node(
                    name=region.neighbors[neigh].cnxn_node_id,
                    pos=cnxn_pos
                )


                # Connect first shared vert to the in-between cnxn node
                graph_dict.setdefault(shared_vert_nodes[0], []).append(cnxn_pos_node)

                # Connect region center to the cnxn node
                graph_dict.setdefault(me, []).append(cnxn_pos_node)

                # connect cnxn node back to region and the first shared node
                graph_dict.setdefault(cnxn_pos_node, [])
                graph_dict[cnxn_pos_node] += [me, shared_vert_nodes[0]]

                if len(shared_vert_nodes) > 1:
                    # connect the second shared vert to the in-between cnxn node
                    graph_dict.setdefault(shared_vert_nodes[1], []).append(cnxn_pos_node)

                    # connect the in-between cnxn_node back to the second shared vert
                    graph_dict[cnxn_pos_node] += [shared_vert_nodes[1]]

        return grph.Graph(
            graph_dict=graph_dict
        )

    def _build_region_graph(
            self,
            regions: Dict[int, Region]
    ) -> grph.Graph:
        logger.debug(f"Building Region Graph...")
        graph_dict = {}

        for id, region in regions.items():
            me = grph.Node(region.node_id, region.pt)

            for neigh, attrs in region.neighbors.items():
                if attrs.Boundary is not None:
                    boundary_pos = attrs.Boundary.Midpoint.as_tuple()
                else:
                    boundary_pos = tuple(float(x) for x in list(attrs.shared_verts.values())[0])

                boundary_node = grph.Node(
                    name=region.neighbors[neigh].cnxn_node_id,
                    pos=boundary_pos
                )

                graph_dict.setdefault(me, []).append(boundary_node)
                graph_dict.setdefault(boundary_node, []).append(me)

        return grph.Graph(
            graph_dict=graph_dict
        )

    def redefine(self,
                 world_bounds: Tuple[int, int],
                 n_regions: int,
                 seed=None):
        if seed is None:
            seed = rnd.randrange(sys.maxsize)
        self._seed = seed
        rnd.seed(self._seed)
        regions, vor, vn_map = _generate_regions(n_regions,
                                                 bounds=[(0, world_bounds[0]),
                                                         (0, world_bounds[1])])
        self.regions: Dict[int, Region] = {x.name: x for x in regions}

        self.village_id = self.choose_village_region().name
        # self.world_graph = self._build_region_graph(self.regions)
        self.world_graph = self._build_boundary_graph(regions=self.regions)

        def __get_villager_or_region_by_id(id: uuid.UUID):
            if id in self.VillagerIdNameMap.keys():
                return self.Villagers[self.VillagerIdNameMap[id]].Position
            if id in self.RegionIdNameMap.keys():
                return self.regions[self.RegionIdNameMap[id]].boundary.exterior.coords

            raise Exception(f'Id: {id} not found in villagers or regions')

        logger.debug(f"Building World SectorTree...")
        self.world_sector_map = SectorTree(
            area_rect=(0, 0, world_bounds[0], world_bounds[1]),
            obj_collider_provider=__get_villager_or_region_by_id,
            shape=(3, 3),
            capacity=1,
            max_lvls=self._sector_tree_depth,
            id='WORLD').add_update_clients(
            {x.id: ClientType.VILLAGER for x in self.Villagers.values()}
        ).add_update_clients({x.id: ClientType.REGION for x in self.regions.values()})

        self._define_rivers()

        self.init_villagers(self._n_init_villagers)

    def _define_rivers(self):
        logger.info(f"Defining Rivers...")
        water_boundaries_node_ids = self._boundaries_between_land_and_water()
        snowpeak_boundaries_node_ids = self._boundaries_between_mountain_and_snowpeak()

        # eval_graph = self.region_graph
        eval_graph = self.world_graph

        if len(water_boundaries_node_ids) == 0 or len(snowpeak_boundaries_node_ids) == 0:
            return

        mouths = [eval_graph.node_by_name(rnd.choice(water_boundaries_node_ids))]
        heads = [eval_graph.node_by_name(rnd.choice(snowpeak_boundaries_node_ids)) for x in range(3)]
        grnd_lvls = {}
        for head in heads:
            grnd_lvls[head] = self.closest_region(
                pt = head.pos,
                excluded_terrain_types=[
                    TerrainType.MOUNTAINS,
                    TerrainType.SNOWPEAK,
                    TerrainType.DEEPSEAWATER,
                    TerrainType.SHALLOWSEAWATER
                ]
            )

        edges_to_add = []

        grnd_cnxn_nodes = []
        for head, clst_grnd in grnd_lvls.items():
            path = eval_graph.astar(
                start=head,
                end=eval_graph.node_by_name(clst_grnd.node_id)
            )
            for e in path.edges[:-1]:
                edges_to_add.append(e)
            grnd_cnxn_nodes.append(edges_to_add[-1].end)

        edges_to_add = set(edges_to_add)


        # Exclude edges that have both neighbor regions ineligible for rivers
        ineligible_cnxns = []
        for id, region in self.regions.items():
            for n_id, attrs in region.neighbors.items():
                if not region.RiverEligible and not self.regions[n_id].RiverEligible:
                    ineligible_cnxns.append(attrs.cnxn_node_id)

        # Also exclude the nodes that define the region's center
        disabled = ineligible_cnxns + [x.name for x in self.world_graph.Nodes.values() if x.name[0] == 'R']

        river_graph = eval_graph.branching_target_seek(
            sources=mouths,
            targets=grnd_cnxn_nodes,
            disabled_node_ids=frozenset(disabled)
        ).reverse()


        river_graph.add_edges(
            edges=list(edges_to_add)
        )

        self.rivers = river_graph


    @timer(logger=logger, log_level=0, time_tracking_class=tt)
    def _handle_villager_reached_waypoint(self,
                                          v: Villager,
                                          w: wf.Waypoint):
        results = None
        if (w is not None and w.is_destination is True):
            results = self._finish_villager_behavior(v, waypoint=w)
            self._set_villager_behavior(v)
            success = self._handle_villager_behavior(v)

            # if not success:
            #     """CRASH IN"""
            #     self.crash_toggle.toggle()
            #     self.crash_regions = (current_region, region)
            #     """CRASH OUT"""

        return results

    def _set_villager_behavior(self, v: Villager):
        current_behaviors = self.CurrentBehaviors

        if v.Behavior is not None and len(current_behaviors[v.Behavior]) > tp.resolve(self.input_controller.villager_counts[v.Behavior]):
            v.set_behavior(None)

        if v.Behavior is None:
            elgbls = [v for v in BehaviorType if len(current_behaviors[v]) < tp.resolve(self.input_controller.villager_counts[v])]
            if len(elgbls) != 0:
                v.set_behavior(rnd.choice(elgbls))

    def _handle_villager_behavior(self, v: Villager):
        # no behavior, send back to village
        if v.Behavior is None:
            cmn.send_villager_to_region(
                v,
                self.Village,
                villager_current_region_provider=self.get_villager_current_region,
                route_to_waypoints_provider=self.route_to_waypoints,
                route_provider=self.route_between_regions,
            )

            return

        behavior_map = {
            BehaviorType.SCOUT: lambda: scout(
                villager=v,
                villager_current_region_provider=self.get_villager_current_region,
                route_to_waypoints_provider=self.route_to_waypoints,
                route_provider=self.route_between_regions,
                village_region=self.Village,
                regions=self.regions
            ),
            BehaviorType.FARM: lambda: farm(
                villager=v,
                villager_current_region_provider=self.get_villager_current_region,
                route_to_waypoints_provider=self.route_to_waypoints,
                route_provider=self.route_between_regions,
                village_region=self.Village,
                regions=self.regions
            )
        }

        return behavior_map[v.Behavior]()


    def reveal_region_at_waypoint(self, villager: Villager, waypoint: wf.Waypoint) -> BehaviorResults:
        revealed_regions = []

        # Find which regions are connected to the waypoint that was just reached
        cnx_nodes = self.world_graph.adjacent_nodes(
            node=self.world_graph.node_by_name(waypoint.id)
        )
        cnct_regions = {id: r for id, r in self.regions.items() if r.node_id in [x.name for x in cnx_nodes]}

        revealed = next(self.regions[x] for x in cnct_regions
                    if self.regions[x].visibility_level == VisibilityLevel.REVEALING)

        reveal_a_region(r=revealed, regions=self.regions)

        revealed_regions.append(revealed)

        return BehaviorResults(
            villager=villager,
            revealed_regions=revealed_regions
        )


    def _finish_villager_behavior(self, v: Villager, waypoint: wf.Waypoint) -> BehaviorResults:
        if v.Behavior is None:
            return

        behavior_map = {
            BehaviorType.SCOUT: lambda: self.reveal_region_at_waypoint(v, waypoint),
            BehaviorType.FARM: lambda: None
        }

        return behavior_map[v.Behavior]()

    @timer(logger=logger, log_level=0, time_tracking_class=tt)
    def _update_villagers(self, delta_ms: int):
        # Check and add villager
        self.villager_spawn_timer.update(delta_ms=delta_ms)
        if self.villager_spawn_timer.Finished and len(self._villagers) < self._max_villagers:
            self.create_villager()
            self.villager_spawn_timer.reset(start=True)

        ret = []

        # update villagers pos and add new routes if necessary
        for id, v in self._villagers.items():
            reached = v.waypoint_follower.update(delta_time_ms=delta_ms,
                                                 close_enough_tolerance=1,
                                                 slow_enough_tolerance=.1,
                                                 )
            if reached is not None:
                results = self._handle_villager_reached_waypoint(v, reached)
                if results is not None:
                    ret.append(results)

            # if v.waypoint_follower.CurrentTarget is None:
            if v.Behavior is None:
                self._set_villager_behavior(v)
                success = self._handle_villager_behavior(v)

        self.world_sector_map.add_update_clients({x.id: ClientType.VILLAGER for x in self.Villagers.values()})
        return ret

    @timer(logger=logger, log_level=0, time_tracking_class=tt)
    def __capsule_nearby_clients(self, nearby_range: float, pt: vec.FloatVec):
        return self.world_sector_map.nearby_clients(
            nearby_range,
            pt=pt,
            client_types=[ClientType.REGION]
        )

    @timer(logger=logger, log_level=0, time_tracking_class=tt)
    def _get_clst_rgn_lgbls(self,
                            pt: vec.FloatVec = None,
                            client_id: UniqueIdentifier = None,
                            nearby_range: int = 50
                            ):
        if pt is None and client_id is None:
            raise Exception(f"at least one of pt and client_id must by non-None")

        # find collocations
        if client_id is not None:
            nearby_region_ids = [x for x in self.world_sector_map.CoLocatedClients[client_id]
                                 if self.world_sector_map.ClientTypeMap[x] == ClientType.REGION]
        else:
            nearby_region_ids = self.__capsule_nearby_clients(
                nearby_range,
                pt=pt,
            )

        return nearby_region_ids

    @timer(logger=logger, log_level=0, time_tracking_class=tt)
    def closest_region(self,
                       pt: vec.FloatVec = None,
                       client_id: UniqueIdentifier = None,
                       nearby_range: int = 50,
                       excluded_terrain_types: Iterable[TerrainType] = None) -> Region:

        if client_id is not None:
            pt = self.world_sector_map.ClientsPos[client_id]

        nearby_region_ids = self._get_clst_rgn_lgbls(
            pt=pt,
            client_id=client_id,
            nearby_range=nearby_range
        )

        # do formal comparisons and return
        region_candidates: List[Region] = [self.regions[self.RegionIdNameMap[x]] for x in nearby_region_ids
                                           if x in self.RegionIdNameMap.keys()]

        # exclude regions that dont match terrain type
        if excluded_terrain_types is not None:
            region_candidates = [x for x in region_candidates if x.terrain not in excluded_terrain_types]

        min_l = None
        nrst = None
        for r in region_candidates:
            dist = vec.distance_between(pt, r.pt)
            if min_l is None or dist < min_l:
                min_l = dist
                nrst = r

        return nrst

    def update(self, delta_ms: int) -> List[BehaviorResults]:
        if self.crash_toggle.value is False:
            results = self._update_villagers(delta_ms=delta_ms)
            self.world_clock.update(delta_ms)
            return results

    def shuffled_region_generator(self) -> Iterator[Region]:
        regs = list(self.regions.values())
        rnd.shuffle(regs)
        for x in regs:
            yield x

    def _get_villager_name(self) -> int:
        self._last_villager_name += 1
        return self._last_villager_name

    def route_to_waypoints(self,
                           path: Iterable[grph.Node]) -> List[wf.Waypoint]:
        wpts = []
        for ii, seg in enumerate(path):
            if ii == len(path) - 1 and seg.name[0] == 'C' and len(path) > 1:
                pos = vec.interpolate(path[ii-1].pos, path[ii].pos, 0.95)
            else:
                pos = seg.pos

            wpts.append(
                wf.Waypoint(
                    end_pos=pos,
                    id=seg.name,
                    is_destination=True if ii == len(path) - 1 else False
                )
            )

        return wpts

    @timer(logger=logger, log_level=0, time_tracking_class=tt)
    def route_between_regions(self,
                              start: Region,
                              end: Region,
                              ) -> grph.AStarResults:

        disbld = [x.node_id for x in self.regions.values()
                  if x.node_id not in [y.node_id for y in self.NavigateableRegions.values()]]

        disbld += [x.name for x in self.world_graph.Nodes.values() if x.name[0] == 'V']

        if str(end.node_id) in disbld:
            disbld.remove(end.node_id)

        if str(start.node_id) in disbld:
            disbld.remove(start.node_id)

        route = self.world_graph.astar(
            start=self.world_graph.node_by_name(start.node_id),
            end=self.world_graph.node_by_name(end.node_id),
            disabled_node_ids=frozenset(disbld)
        )
        return route

    def get_villager_current_region(self, villager: Villager) -> Region:
        return self.closest_region(client_id=villager.id)

    def _boundaries_between_terrain_types(self,
                                          target_group: Iterable[TerrainType]) -> List[str]:
        ret = []
        for id, region in self.regions.items():
            for n_id in region.neighbors.keys():
                if (self.regions[n_id].terrain in target_group and
                        region.terrain in target_group):
                    # Both are target_group
                    continue

                if (self.regions[n_id].terrain not in target_group and
                        region.terrain not in target_group):
                    # neither are in target_group
                    continue

                if (cnxn := region.neighbors[n_id].cnxn_node_id) not in ret:
                    ret.append(cnxn)

        return ret

    def _boundaries_between_mountain_and_snowpeak(self) -> List[str]:
        return self._boundaries_between_terrain_types([TerrainType.SNOWPEAK])

    def _boundaries_between_land_and_water(self) -> List[str]:
        return self._boundaries_between_terrain_types(
            [TerrainType.DEEPSEAWATER,
              TerrainType.SHALLOWSEAWATER])

    def world_time(self):
        s = self.world_clock.AccumulatedS
        hrs = s / self._seconds_per_hr
        days = int(hrs / 24)
        hrs = int(hrs - 24 * days)

        return days, hrs



    @property
    def Village(self) -> Region:
        return self.regions[self.village_id]

    @property
    def Villagers(self) -> Dict[int, Villager]:
        return self._villagers

    @property
    def NVillagers(self) -> int:
        return len(self.Villagers)

    @property
    def VillagerIdNameMap(self) -> Dict[uuid.UUID, int]:
        return {x.id: x.name for x in self.Villagers.values()}

    @property
    def VillagerNameIdMap(self) -> Dict[int, uuid.UUID]:
        return {x.name: x.id for x in self.Villagers.values()}

    @property
    def RegionIdNameMap(self) -> Dict[uuid.UUID, int]:
        return {x.id: x.name for x in self.regions.values()}

    @property
    def RegionNameIdMap(self) -> Dict[int, uuid.UUID]:
        return {x.name: x.id for x in self.regions.values()}

    @property
    def NavigateableRegions(self) -> Dict[int, Region]:
        return {
            k: v for k, v in self.regions.items() if v.visibility_level in [
                VisibilityLevel.REVEALED
            ] and v.terrain not in [
                 TerrainType.MOUNTAINS,
                 TerrainType.SNOWPEAK,
                 TerrainType.SHALLOWSEAWATER,
                 TerrainType.DEEPSEAWATER
             ]
        }

    @property
    def ReachableRegions(self) -> Dict[int, Region]:
        return {
            k: v for k, v in self.regions.items() if v.visibility_level in [
                VisibilityLevel.REVEALED,
                VisibilityLevel.REVEALABLE,
                VisibilityLevel.REVEALING
            ]
        }

    @property
    def CurrentBehaviors(self) -> Dict[BehaviorType, List[Villager]]:
        ret = {b: [] for b in BehaviorType}

        for id, v in self.Villagers.items():
            ret.setdefault(v.Behavior, []).append(v)
        return ret
