NetUtils.py

from __future__ import annotations

from collections.abc import Mapping, Sequence
import typing
import enum
import warnings
from json import JSONEncoder, JSONDecoder

if typing.TYPE_CHECKING:
    from websockets import WebSocketServerProtocol as ServerConnection

from BaseUtils import ByValue, Version

class HintStatus(ByValue, enum.IntEnum):
    HINT_UNSPECIFIED = 0
    HINT_NO_PRIORITY = 10
    HINT_AVOID = 20
    HINT_PRIORITY = 30
    HINT_FOUND = 40

class MWGGUIHintStatus(ByValue, enum.IntFlag):
    """
    Shop Item, Goal-required item, and 'this item is what is keeping me in BK_MODE'
    BK_MODE items will be shown as the highest priority.
    """
    HINT_UNSPECIFIED = 0b000
    HINT_SHOP = 0b001
    HINT_GOAL = 0b010
    HINT_BK_MODE = 0b100


class JSONMessagePart(typing.TypedDict, total=False):
    text: str
    # optional
    type: str
    color: str
    # owning player for location/item
    player: int
    # if type == item indicates item flags
    flags: int
    # if type == hint_status
    hint_status: HintStatus


class ClientStatus(ByValue, enum.IntEnum):
    CLIENT_UNKNOWN = 0
    CLIENT_CONNECTED = 5
    CLIENT_READY = 10
    CLIENT_PLAYING = 20
    CLIENT_GOAL = 30


class SlotType(ByValue, enum.IntFlag):
    spectator = 0b00
    player = 0b01
    group = 0b10

    @property
    def always_goal(self) -> bool:
        """Mark this slot as having reached its goal instantly."""
        return self.value != 0b01


class Permission(ByValue, enum.IntFlag):
    disabled = 0b000  # 0, completely disables access
    enabled = 0b001  # 1, allows manual use
    goal = 0b010  # 2, allows manual use after goal completion
    auto = 0b110  # 6, forces use after goal completion, only works for release
    auto_enabled = 0b111  # 7, forces use after goal completion, allows manual use any time

    @staticmethod
    def from_text(text: str):
        data = 0
        if "auto" in text:
            data |= 0b110
        elif "goal" in text:
            data |= 0b010
        if "enabled" in text:
            data |= 0b001
        return Permission(data)


class NetworkPlayer(typing.NamedTuple):
    """Represents a particular player on a particular team."""
    team: int
    slot: int
    alias: str
    name: str


class NetworkSlot(typing.NamedTuple):
    """Represents a particular slot across teams."""
    name: str
    game: str
    type: SlotType
    group_members: Sequence[int] = ()  # only populated if type == group


class NetworkItem(typing.NamedTuple):
    item: int
    location: int
    player: int
    """ Sending player, except in LocationInfo (from LocationScouts), where it is the receiving player. """
    flags: int = 0

def escape_markup(text: str) -> str:
    return text.replace('&', '&').replace('[', '&bl;').replace(']', '&br;')

def _scan_for_TypedTuples(obj: typing.Any) -> typing.Any:
    if isinstance(obj, tuple) and hasattr(obj, "_fields"):  # NamedTuple is not actually a parent class
        data = obj._asdict()
        data["class"] = obj.__class__.__name__
        return data
    if isinstance(obj, (tuple, list, set, frozenset)):
        return tuple(_scan_for_TypedTuples(o) for o in obj)
    if isinstance(obj, dict):
        return {key: _scan_for_TypedTuples(value) for key, value in obj.items()}
    return obj


_base_types = str | int | bool | float | None | tuple["_base_types", ...] | dict["_base_types", "base_types"]


def convert_to_base_types(obj: typing.Any) -> _base_types:
    if isinstance(obj, (tuple, list, set, frozenset)):
        return tuple(convert_to_base_types(o) for o in obj)
    elif isinstance(obj, dict):
        return {convert_to_base_types(key): convert_to_base_types(value) for key, value in obj.items()}
    elif obj is None or type(obj) in (str, int, float, bool):
        return obj
    # unwrap simple types to their base, such as StrEnum
    elif isinstance(obj, str):
        return str(obj)
    elif isinstance(obj, int):
        return int(obj)
    elif isinstance(obj, float):
        return float(obj)
    else:
        raise Exception(f"Cannot handle {type(obj)}")


_encode = JSONEncoder(
    ensure_ascii=False,
    check_circular=False,
    separators=(',', ':'),
).encode


def encode(obj: typing.Any) -> str:
    return _encode(_scan_for_TypedTuples(obj))


def get_any_version(data: dict) -> Version:
    data = {key.lower(): value for key, value in data.items()}  # .NET version classes have capitalized keys
    return Version(int(data["major"]), int(data["minor"]), int(data["build"]))


allowlist = {
    "NetworkPlayer": NetworkPlayer,
    "NetworkItem": NetworkItem,
    "NetworkSlot": NetworkSlot
}

custom_hooks = {
    "Version": get_any_version
}


def _object_hook(o: typing.Any) -> typing.Any:
    if isinstance(o, dict):
        hook = custom_hooks.get(o.get("class", None), None)
        if hook:
            return hook(o)
        cls = allowlist.get(o.get("class", None), None)
        if cls:
            for key in tuple(o):
                if key not in cls._fields:
                    del (o[key])
            return cls(**o)

    return o


decode = JSONDecoder(object_hook=_object_hook).decode


class Endpoint:
    __slots__ = ("socket",)

    socket: "ServerConnection"

    def __init__(self, socket):
        self.socket = socket


class HandlerMeta(type):
    def __new__(mcs, name, bases, attrs):
        handlers = attrs["handlers"] = {}
        trigger: str = "_handle_"
        for base in bases:
            handlers.update(base.handlers)
        handlers.update({handler_name[len(trigger):]: method for handler_name, method in attrs.items() if
                         handler_name.startswith(trigger)})

        orig_init = attrs.get('__init__', None)
        if not orig_init:
            for base in bases:
                orig_init = getattr(base, '__init__', None)
                if orig_init:
                    break

        def __init__(self, *args, **kwargs):
            if orig_init:
                orig_init(self, *args, **kwargs)
            # turn functions into bound methods
            self.handlers = {name: method.__get__(self, type(self)) for name, method in
                             handlers.items()}

        attrs['__init__'] = __init__
        return super(HandlerMeta, mcs).__new__(mcs, name, bases, attrs)


class JSONTypes(str, enum.Enum):
    color = "color"
    text = "text"
    player_id = "player_id"
    player_name = "player_name"
    item_name = "item_name"
    item_id = "item_id"
    location_name = "location_name"
    location_id = "location_id"
    entrance_name = "entrance_name"

# Default color definitions - these should be imported by GUI themes
# [Dark, Light]
TEXT_COLORS = {
    "default_color": "cdcdcd",
    "command_echo_color": "ff9334",
    "player1_color": "ff87d7",
    "player2_color": "5fafff",
    "progression_goal_item_color": "ffa700",
    "progression_item_color": "ffbe00",
    "progression_deprioritized_item_color": "d2ff49",
    "useful_item_color": "6EC471",
    "regular_item_color": "b2b2b2",
    "trap_item_color": "d75f5f",
    "location_color": "00c51b",
    "entrance_color": "60b7e8",
}

class JSONtoTextParser(metaclass=HandlerMeta):
    # add *all* of the colors, to prevent crashes where colors are expected.
    from kivy.utils import hex_colormap as color_codes
    # then add the custom ones
    for key,value in TEXT_COLORS.items():
        color_codes[key] = value

    def __init__(self, ctx):
        self.ctx = ctx

    def __call__(self, input_object: typing.List[JSONMessagePart]) -> str:
        return "".join(self.handle_node(section) for section in input_object)

    def handle_node(self, node: JSONMessagePart):
        node_type = node.get("type", None)
        handler = self.handlers.get(node_type, self.handlers["plaintext"])
        return handler(node)

    def _handle_color(self, node: JSONMessagePart):
        codes = node["color"].split(";")
        buffer = "".join(color_code(code) for code in codes if code in color_codes)
        return buffer + self._handle_text(node) + color_code("reset")

    def _handle_text(self, node: JSONMessagePart):
        return node.get("text", "")

    def _handle_player_id(self, node: JSONMessagePart):
        player = int(node["text"])
        node["color"] = 'player1_color' if self.ctx.slot_concerns_self(player) else 'player2_color'
        node["text"] = self.ctx.player_names[player]
        return self._handle_color(node)

    # for other teams, spectators etc.? Only useful if player isn't in the clientside mapping
    def _handle_player_name(self, node: JSONMessagePart):
        node["color"] = 'player2_color'
        return self._handle_color(node)

    def _handle_item_name(self, node: JSONMessagePart):
        flags = node.get("flags", 0)
        if flags == 0:
            node["color"] = 'regular_item_color' # filler
        elif flags & 0b00010:  # useful
            node["color"] = 'useful_item_color'  # lime for useful items
        if flags & 0b00100:  # "useful trap" gets marked trap
            node["color"] = 'trap_item_color'  # salmon for traps
        elif flags & 0b00001:  # progression is the third flag checked, so it can overwrite. 
            # "useful progression" gets marked progression
            node["color"] = 'progression_item_color'  # "dulled" gold for regular progression
            if flags & 0b10000:  # deprioritized, but still progression (skulls etc)
                node["color"] = 'progression_deprioritized_item_color'  # Citron for progression deprioritized
            else:
                if flags & 0b01000:  # skip_balancing bit set
                    node["color"] = 'progression_goal_item_color'  # Gold for progression skip items/macguffins
        if not node["color"]:
            # if we can't find the flag set, use the command echo color to indicate it doesn't know what kind of item it is
            node["color"] = 'command_echo_color'
        node["text"] = escape_markup(node["text"])
        return self._handle_color(node)

    def _handle_item_id(self, node: JSONMessagePart):
        item_id = int(node["text"])
        node["text"] = self.ctx.item_names.lookup_in_slot(item_id, node["player"])
        return self._handle_item_name(node)

    def _handle_location_name(self, node: JSONMessagePart):
        node["color"] = 'location_color'
        node["text"] = escape_markup(node["text"])
        return self._handle_color(node)

    def _handle_location_id(self, node: JSONMessagePart):
        location_id = int(node["text"])
        node["text"] = self.ctx.location_names.lookup_in_slot(location_id, node["player"])
        return self._handle_location_name(node)

    def _handle_entrance_name(self, node: JSONMessagePart):
        node["color"] = 'entrance_color'
        node["text"] = escape_markup(node["text"])
        return self._handle_color(node)

    def _handle_hint_status(self, node: JSONMessagePart):
        node["color"] = status_colors.get(node["hint_status"], "red")
        node["text"] = escape_markup(node["text"])
        return self._handle_color(node)

    def _handle_plaintext(self, node: JSONMessagePart):
        if "[color=" in node["text"]:
            import re
            node_list = []
            text = node["text"]
            
            # Use regex to find all [color=value]text[/color] patterns
            pattern = r'\[color=([^\]]+)\](.*?)\[/color\]'
            # pattern = r'\[i\](.*?)\[/i\]'
            # pattern = r'\[b\](.*?)\[/b\]'
            last_end = 0
            
            for match in re.finditer(pattern, text):
                # Add text before the color tag
                if match.start() > last_end:
                    before_text = text[last_end:match.start()]
                    if before_text:
                        node_list.append({"color": 'default_color', "text": escape_markup(before_text)})
                
                # Add the color tag content
                color_value = match.group(1)
                tag_text = match.group(2)
                node_list.append({"color": color_value, "text": escape_markup(tag_text)})
                last_end = match.end()
            
            # Add any remaining text after the last color tag
            if last_end < len(text):
                remaining_text = text[last_end:]
                if remaining_text:
                    node_list.append({"color": 'default_color', "text": escape_markup(remaining_text)})
            
            return_text = ""
            for node_part in node_list:
                return_text += self._handle_color(node_part)
            return return_text
        
        node["text"] = escape_markup(node["text"])
        node["color"] = 'default_color'
        return self._handle_color(node)


class RawJSONtoTextParser(JSONtoTextParser):
    def _handle_color(self, node: JSONMessagePart):
        return self._handle_text(node)


class KivyMarkupJSONtoTextParser(JSONtoTextParser):
    """JSON parser that converts to Kivy markup format with hex colors"""
    def __init__(self, ctx):
        super().__init__(ctx)
        for key,value in TEXT_COLORS.items():
            self.color_codes[key] = value

    def _handle_color(self, node: JSONMessagePart):
        codes = node["color"].split(";")
        # Find the first valid color code
        color_hex = None
        for code in codes:
            if code in self.color_codes:
                color_hex = self.color_codes[code]
                break
        
        if color_hex:
            # Get the plain text without wrapping it in default color
            text = node.get("text", "")
            return f'[color={color_hex}]{text}[/color]'
        else:
            return self._handle_text(node)
    
    def _handle_text(self, node: JSONMessagePart):
        return node.get("text", "")

# setting ansi colors - Added many 8 bit to go with the 4 bit.
color_codes = {'reset': 0, 'bold': 1, 'underline': 4, 'black': 30, 'red': 31, 'green': 32, 'yellow': 33, 'blue': 34,
                'magenta': 35, 'cyan': 36, 'white': 37, 'black_bg': 40, 'red_bg': 41, 'green_bg': 42, 'yellow_bg': 43,
                'blue_bg': 44, 'magenta_bg': 45, 'cyan_bg': 46, 'white_bg': 47,
                'plum': 33, 'slateblue': 32, 'salmon': 31, 'limegreen': 32, 'lightgray': 37, 'gold': 33,
                'default_color': 37, #white
                'location_color': '38;5;34', #green
                'player1_color': '38;5;212', #atzpink
                'player2_color': '38;5;75', #ltblue
                'entrance_color': '38;5;27', #blue
                'trap_item_color': '38;5;167', #salmon
                'regular_item_color': '38;5;249', #gray
                'useful_item_color': '38;5;149', #lime
                'progression_skip_item_color': '38;5;220', #gold
                'progression_item_color': '38;5;220', #gold
                'command_echo_color': '38;5;208' #orange
}

def color_code(*args):
    return '\033[' + ';'.join([str(color_codes[arg]) for arg in args]) + 'm'


def color(text, *args):
    return color_code(*args) + text + '\33[0m'


def add_json_text(parts: list, text: typing.Any, **kwargs) -> None:
    parts.append({"text": str(text), **kwargs})


def add_json_item(parts: list, item_id: int, player: int = 0, item_flags: int = 0, **kwargs) -> None:
    parts.append({"text": str(item_id), "player": player, "flags": item_flags, "type": JSONTypes.item_id, **kwargs})


def add_json_location(parts: list, location_id: int, player: int = 0, **kwargs) -> None:
    parts.append({"text": str(location_id), "player": player, "type": JSONTypes.location_id, **kwargs})

## HintStatus map is the location identifier/colors
status_names: typing.Dict[HintStatus, str] = {
    HintStatus.HINT_FOUND: "(found)",
    HintStatus.HINT_UNSPECIFIED: "(unspecified)",
    HintStatus.HINT_NO_PRIORITY: "(no priority)",
    HintStatus.HINT_AVOID: "(avoid)",
    HintStatus.HINT_PRIORITY: "(priority)",
}
status_colors: typing.Dict[HintStatus, str] = {
    HintStatus.HINT_FOUND: "green",
    HintStatus.HINT_UNSPECIFIED: "white",
    HintStatus.HINT_NO_PRIORITY: "lightgray",
    HintStatus.HINT_AVOID: "salmon",
    HintStatus.HINT_PRIORITY: "gold",
}
## HintStatus map is the location identifier/colors
mwggui_status_names: typing.Dict[MWGGUIHintStatus, str] = {
    MWGGUIHintStatus.HINT_SHOP: "(shop)",
    MWGGUIHintStatus.HINT_GOAL: "(goal)",
    MWGGUIHintStatus.HINT_BK_MODE: "(bk_mode)",
}
mwggui_status_colors: typing.Dict[MWGGUIHintStatus, str] = {
    MWGGUIHintStatus.HINT_SHOP: "grey",
    MWGGUIHintStatus.HINT_GOAL: "gold",
    MWGGUIHintStatus.HINT_BK_MODE: "salmon",
}

def add_json_hint_status(parts: list, hint_status: HintStatus, text: typing.Optional[str] = None, **kwargs):
    parts.append({"text": text if text != None else status_names.get(hint_status, "(unknown)"),
                  "hint_status": hint_status, "type": JSONTypes.hint_status, **kwargs})


class Hint(typing.NamedTuple):
    receiving_player: int
    finding_player: int
    location: int
    item: int
    found: bool
    entrance: str = ""
    item_flags: int = 0
    status: HintStatus = HintStatus.HINT_UNSPECIFIED

    def re_check(self, ctx, team) -> Hint:
        if self.found and self.status == HintStatus.HINT_FOUND:
            return self
        found = self.location in ctx.location_checks[team, self.finding_player]
        if found:
            return self._replace(found=found, status=HintStatus.HINT_FOUND)
        return self
    
    def re_prioritize(self, ctx, status: HintStatus) -> Hint:
        if self.found and status != HintStatus.HINT_FOUND:
            status = HintStatus.HINT_FOUND
        if status != self.status:
            return self._replace(status=status)
        return self

    def __hash__(self):
        return hash((self.receiving_player, self.finding_player, self.location, self.item, self.entrance))

    def as_network_message(self) -> dict:
        parts = []
        add_json_text(parts, "[Hint]: ")
        add_json_text(parts, self.receiving_player, type="player_id")
        add_json_text(parts, "'s ")
        add_json_item(parts, self.item, self.receiving_player, self.item_flags)
        add_json_text(parts, " is at ")
        add_json_location(parts, self.location, self.finding_player)
        add_json_text(parts, " in ")
        add_json_text(parts, self.finding_player, type="player_id")
        if self.entrance:
            add_json_text(parts, "'s World at ")
            add_json_text(parts, self.entrance, type="entrance_name")
        else:
            add_json_text(parts, "'s World")
        add_json_text(parts, ". ")
        add_json_text(parts, status_names.get(self.status, "(unknown)"), type="color",
                      color=status_colors.get(self.status, "red"))

        return {"cmd": "PrintJSON", "data": parts, "type": "Hint",
                "receiving": self.receiving_player,
                "item": NetworkItem(self.item, self.location, self.finding_player, self.item_flags),
                "found": self.found}

    @property
    def local(self):
        return self.receiving_player == self.finding_player


class _LocationStore(dict, typing.MutableMapping[int, typing.Dict[int, typing.Tuple[int, int, int]]]):
    def __init__(self, values: typing.MutableMapping[int, typing.Dict[int, typing.Tuple[int, int, int]]]):
        super().__init__(values)

        if not self:
            raise ValueError(f"Rejecting game with 0 players")

        if len(self) != max(self):
            raise ValueError("Player IDs not continuous")

        if len(self.get(0, {})):
            raise ValueError("Invalid player id 0 for location")

    def find_item(self, slots: typing.Set[int], seeked_item_id: int
                  ) -> typing.Generator[typing.Tuple[int, int, int, int, int], None, None]:
        for finding_player, check_data in self.items():
            for location_id, (item_id, receiving_player, item_flags) in check_data.items():
                if receiving_player in slots and item_id == seeked_item_id:
                    yield finding_player, location_id, item_id, receiving_player, item_flags

    def get_for_player(self, slot: int) -> typing.Dict[int, typing.Set[int]]:
        import collections
        all_locations: typing.Dict[int, typing.Set[int]] = collections.defaultdict(set)
        for source_slot, location_data in self.items():
            for location_id, values in location_data.items():
                if values[1] == slot:
                    all_locations[source_slot].add(location_id)
        return all_locations

    def get_checked(self, state: typing.Dict[typing.Tuple[int, int], typing.Set[int]], team: int, slot: int
                    ) -> typing.List[int]:
        checked = state[team, slot]
        if not checked:
            # This optimizes the case where everyone connects to a fresh game at the same time.
            if slot not in self:
                raise KeyError(slot)
            return []
        return [location_id for
                location_id in self[slot] if
                location_id in checked]

    def get_missing(self, state: typing.Dict[typing.Tuple[int, int], typing.Set[int]], team: int, slot: int
                    ) -> typing.List[int]:
        checked = state[team, slot]
        if not checked:
            # This optimizes the case where everyone connects to a fresh game at the same time.
            return list(self[slot])
        return [location_id for
                location_id in self[slot] if
                location_id not in checked]

    def get_remaining(self, state: typing.Dict[typing.Tuple[int, int], typing.Set[int]], team: int, slot: int
                      ) -> typing.List[typing.Tuple[int, int]]:
        checked = state[team, slot]
        player_locations = self[slot]
        return sorted([(player_locations[location_id][1], player_locations[location_id][0]) for
                        location_id in player_locations if
                        location_id not in checked])


class MinimumVersions(typing.TypedDict):
    server: tuple[int, int, int]
    clients: dict[int, tuple[int, int, int]]


class GamesPackage(typing.TypedDict, total=False):
    item_name_groups: dict[str, list[str]]
    item_name_to_id: dict[str, int]
    location_name_groups: dict[str, list[str]]
    location_name_to_id: dict[str, int]
    checksum: str


class DataPackage(typing.TypedDict):
    games: dict[str, GamesPackage]


class MultiData(typing.TypedDict):
    slot_data: dict[int, Mapping[str, typing.Any]]
    slot_info: dict[int, NetworkSlot]
    connect_names: dict[str, tuple[int, int]]
    locations: dict[int, dict[int, tuple[int, int, int]]]
    checks_in_area: dict[int, dict[str, int | list[int]]]
    server_options: dict[str, object]
    er_hint_data: dict[int, dict[int, str]]
    precollected_items: dict[int, list[int]]
    precollected_hints: dict[int, set[Hint]]
    version: tuple[int, int, int]
    tags: list[str]
    minimum_versions: MinimumVersions
    seed_name: str
    spheres: list[dict[int, set[int]]]
    datapackage: dict[str, GamesPackage]
    race_mode: int


if typing.TYPE_CHECKING:  # type-check with pure python implementation until we have a typing stub
    LocationStore = _LocationStore
else:
    try:
        from _speedups import LocationStore
        import _speedups
        import os.path
        if os.path.isfile("_speedups.pyx") and os.path.getctime(_speedups.__file__) < os.path.getctime("_speedups.pyx"):
            warnings.warn(f"{_speedups.__file__} outdated! "
                          f"Please rebuild with `cythonize -b -i _speedups.pyx` or delete it!")
    except ImportError:
        try:
            import pyximport
            pyximport.install()
        except ImportError:
            pyximport = None
        try:
            from _speedups import LocationStore
        except ImportError:
            warnings.warn("_speedups not available. Falling back to pure python LocationStore. "
                          "Install a matching C++ compiler for your platform to compile _speedups.")
            LocationStore = _LocationStore