parser.py

import re
from typing import List
from .vocab_manager import VocabManager


class LogTokenizer:
    """
    Parses raw Showdown log lines into a sequence of tokens from our Vocab.
    """

    def __init__(self, vocab_manager: VocabManager):
        self.vm = vocab_manager
        self.vm.load_vocab()
        self.pokedex = self.vm.get_pokedex()

        # Map commands to their handler functions
        self.command_handlers = {
            "player": self._handle_player,
            "clearpoke": self._handle_clearpoke,
            "poke": self._handle_poke,
            "teampreview": self._handle_teampreview,
            "start": self._handle_start,
            "turn": self._handle_turn,
            "move": self._handle_move,
            "switch": self._handle_switch,
            "drag": self._handle_switch,
            "cant": self._handle_cant,
            "faint": self._handle_faint,
            "-damage": self._handle_generic_with_token(self.vm.DAMAGE),
            "-heal": self._handle_generic_with_token(self.vm.HEAL),
            "-status": self._handle_generic_with_token(self.vm.STATUS),
            "-curestatus": self._handle_generic_with_token(self.vm.CURE_STATUS),
            "-boost": self._handle_generic_with_token(self.vm.BOOST),
            "-unboost": self._handle_generic_with_token(self.vm.UNBOOST),
            "-weather": self._handle_generic_with_token(self.vm.WEATHER),
            "-fieldstart": self._handle_generic_with_token(self.vm.FIELD_START),
            "-fieldend": self._handle_generic_with_token(self.vm.FIELD_END),
            "-sidestart": self._handle_generic_with_token(self.vm.SIDE_START),
            "-sideend": self._handle_generic_with_token(self.vm.SIDE_END),
            "-crit": self._handle_simple_token(self.vm.CRIT),
            "-supereffective": self._handle_simple_token(self.vm.SUPEREFFECTIVE),
            "-resisted": self._handle_simple_token(self.vm.RESISTED),
            "-immune": self._handle_simple_token(self.vm.IMMUNE),
            "-miss": self._handle_simple_token(self.vm.MISS),
            "-fail": self._handle_simple_token(self.vm.FAIL),
            "-block": self._handle_simple_token(self.vm.BLOCK),
            "-activate": self._handle_generic_with_token(self.vm.ACTIVATE),
            "-item": self._handle_generic_with_token(self.vm.ITEM_CMD),
            "-enditem": self._handle_generic_with_token(self.vm.END_ITEM),
            "-ability": self._handle_generic_with_token(self.vm.ABILITY_CMD),
            "-endability": self._handle_generic_with_token(self.vm.END_ABILITY),
            "-transform": self._handle_generic_with_token(self.vm.TRANSFORM),
            "-formechange": self._handle_generic_with_token(self.vm.FORM_CHANGE),
            "-mega": self._handle_generic_with_token(self.vm.MEGA),
            "-primal": self._handle_generic_with_token(self.vm.PRIMAL),
            "-zpower": self._handle_generic_with_token(self.vm.Z_POWER),
            "-terastallize": self._handle_generic_with_token(self.vm.TERASTALLIZE),
            "terastallize": self._handle_generic_with_token(self.vm.TERASTALLIZE),
            "-dynamax": self._handle_generic_with_token(self.vm.DYNAMAX),
            "-singleturn": self._handle_generic_with_token(self.vm.SINGLE_TURN),
            "-singlemove": self._handle_generic_with_token(self.vm.SINGLE_MOVE),
            "win": self._handle_win,
        }

    def _handle_player(self, parts, ctx):
        if len(parts) > 3:
            pid = parts[2]
            name = parts[3]
            if pid == "p1":
                ctx["p1_name"] = name
            elif pid == "p2":
                ctx["p2_name"] = name

    def _handle_clearpoke(self, parts, ctx):
        ctx["tokens"].append(self.vm.tokenize(self.vm.CLEARED_POKE))

    def _handle_poke(self, parts, ctx):
        if len(parts) > 3:
            ctx["tokens"].append(self.vm.tokenize(self.vm.POKE))
            player = parts[2]
            species = parts[3].split(",")[0].strip()

            if player == "p1":
                ctx["tokens"].append(self.vm.tokenize(self.vm.PLAYER_1))
                ctx["p1_roster"].append(species)
            elif player == "p2":
                ctx["tokens"].append(self.vm.tokenize(self.vm.PLAYER_2))
                ctx["p2_roster"].append(species)

            mon_token = self.vm.formatted_token("MON", species)
            if mon_token in self.vm.vocab:
                ctx["tokens"].append(self.vm.vocab[mon_token])
            else:
                ctx["tokens"].append(self.vm.vocab[self.vm.UNK_TOKEN])

    def _handle_teampreview(self, parts, ctx):
        ctx["tokens"].append(self.vm.tokenize(self.vm.TEAMS_PREVIEW))

    def _handle_start(self, parts, ctx):
        ctx["tokens"].append(self.vm.tokenize(self.vm.START))

    def _handle_turn(self, parts, ctx):
        ctx["tokens"].append(self.vm.tokenize(self.vm.TURN))

    def _handle_move(self, parts, ctx):
        ctx["tokens"].append(self.vm.tokenize(self.vm.MOVE_CMD))
        self._parse_move(parts, ctx["tokens"])

    def _handle_switch(self, parts, ctx):
        ctx["tokens"].append(
            self.vm.tokenize(self.vm.SWITCH if parts[1] == "switch" else self.vm.DRAG)
        )
        self._parse_switch(parts, ctx["tokens"], ctx["p1_roster"], ctx["p2_roster"])

    def _handle_cant(self, parts, ctx):
        self._parse_cant(parts, ctx["tokens"])

    def _handle_faint(self, parts, ctx):
        self._parse_faint(parts, ctx["tokens"])

    def _handle_generic_with_token(self, token_name):
        def handler(parts, ctx):
            ctx["tokens"].append(self.vm.tokenize(token_name))
            self._parse_generic_args(parts, ctx["tokens"])

        return handler

    def _handle_simple_token(self, token_name):
        def handler(parts, ctx):
            ctx["tokens"].append(self.vm.tokenize(token_name))

        return handler

    def _handle_win(self, parts, ctx):
        ctx["tokens"].append(self.vm.tokenize(self.vm.WIN))
        if len(parts) > 2:
            winner = parts[2]
            if winner == ctx["p1_name"]:
                ctx["tokens"].append(self.vm.tokenize(self.vm.PLAYER_1))
            elif winner == ctx["p2_name"]:
                ctx["tokens"].append(self.vm.tokenize(self.vm.PLAYER_2))

    def tokenize_replay(self, log_path: str) -> List[int]:
        """
        Reads a log file and returns a list of Token IDs.
        """
        tokens = [self.vm.vocab[self.vm.SOS_TOKEN]]

        # Context object to pass state to handlers
        ctx = {
            "tokens": tokens,
            "p1_roster": [],
            "p2_roster": [],
            "p1_name": None,
            "p2_name": None,
        }

        # Track nicknames to species mapping for the current match
        # Key: "p1a: Nickname", Value: "SpeciesName"
        self.nickname_map = {}

        with open(log_path, "r") as f:
            for line in f:
                line = line.strip()
                if not line.startswith("|"):
                    continue

                parts = line.split("|")
                if len(parts) < 2:
                    continue

                cmd = parts[1]

                # Filter out chat and cosmetic lines
                if cmd in [
                    "chat",
                    "j",
                    "l",
                    "search",
                    "join",
                    "leave",
                    "html",
                    "raw",
                    "dealloc",
                    "seed",
                ]:
                    continue

                # Dispatch to handler
                handler = self.command_handlers.get(cmd)
                if handler:
                    handler(parts, ctx)
                else:
                    # For generic lines (damage, boost, etc), we can try to tokenize the arguments
                    # This is a "blind" tokenization of the state updates
                    # |cmd|arg1|arg2...
                    self._parse_generic_args(parts, tokens)

        tokens.append(self.vm.vocab[self.vm.EOS_TOKEN])
        return tokens

    def _parse_generic_args(self, parts, tokens):
        for part in parts[2:]:
            self._blind_tokenize_part(part, tokens)

    def extract_metadata(self, log_path: str) -> dict:
        """
        Extracts metadata: winner, ratings.
        """
        meta = {"winner": None, "p1_rating": 1000, "p2_rating": 1000}

        p1_name = None
        p2_name = None

        with open(log_path, "r") as f:
            for line in f:
                parts = line.strip().split("|")
                if len(parts) < 2:
                    continue
                cmd = parts[1]

                if cmd == "player":
                    # |player|p1|Username|Avatar|Rating
                    if len(parts) > 2:
                        pid = parts[2]
                        if len(parts) > 3:
                            name = parts[3]
                            if pid == "p1":
                                p1_name = name
                            elif pid == "p2":
                                p2_name = name

                        if len(parts) > 5:
                            try:
                                rating = float(parts[5])
                                if pid == "p1":
                                    meta["p1_rating"] = rating
                                elif pid == "p2":
                                    meta["p2_rating"] = rating
                            except:
                                pass

                elif cmd == "win":
                    # |win|Username
                    if len(parts) > 2:
                        winner_name = parts[2]
                        if winner_name == p1_name:
                            meta["winner"] = "p1"
                        elif winner_name == p2_name:
                            meta["winner"] = "p2"
                        break
        return meta

    def _blind_tokenize_part(self, text: str, tokens: List[int]):
        """Attempts to find known keywords in simple text fields."""

        # Handle [spread]
        if text.startswith("[spread]"):
            tokens.append(self.vm.tokenize(self.vm.formatted_token("KW", "spread")))
            remaining = text.replace("[spread]", "").strip()
            if remaining:
                refs = remaining.split(",")
                for r in refs:
                    self._blind_tokenize_part(r.strip(), tokens)
            return

        # Handle [notarget]
        if "[notarget]" in text:
            tokens.append(self.vm.tokenize(self.vm.formatted_token("KW", "notarget")))
            return

        # Handle [of] clauses
        if text.startswith("[of]"):
            source = text.replace("[of]", "").strip()
            self._blind_tokenize_part(source, tokens)
            return

        # Handle direct prefixes (item: X, move: Y)
        if text.startswith("item:"):
            item_name = text.split(":")[1].strip()
            tokens.append(self.vm.tokenize(self.vm.formatted_token("ITEM", item_name)))
            return
        elif text.startswith("ability:"):
            abil_name = text.split(":")[1].strip()
            tokens.append(self.vm.tokenize(self.vm.formatted_token("ABIL", abil_name)))
            return
        elif text.startswith("move:"):
            move_name = text.split(":")[1].strip()
            tokens.append(self.vm.tokenize(self.vm.formatted_token("MOVE", move_name)))
            return

        # Handle [from] clauses
        if text.startswith("[from]"):
            source = text.replace("[from]", "").strip()

            # Specific prefix handling
            if source.startswith("item:"):
                item_name = source.split(":")[1].strip()
                tokens.append(
                    self.vm.tokenize(self.vm.formatted_token("ITEM", item_name))
                )
                return
            elif source.startswith("ability:"):
                abil_name = source.split(":")[1].strip()
                tokens.append(
                    self.vm.tokenize(self.vm.formatted_token("ABIL", abil_name))
                )
                return
            elif source.startswith("move:"):
                move_name = source.split(":")[1].strip()
                tokens.append(
                    self.vm.tokenize(self.vm.formatted_token("MOVE", move_name))
                )
                return

            # Fallback: try key vocab categories
            for category in ["MOVE", "ABIL", "ITEM", "STATUS", "MECH"]:
                candidate = self.vm.formatted_token(category, source)
                if candidate in self.vm.vocab:
                    tokens.append(self.vm.vocab[candidate])
                    return

        # Check if it's a player ref (p1a: Mon)
        if re.match(r"^p[12][abcd]:", text):
            self._parse_player_ref(text, tokens)
            return

        # Try prefixes for plain text (e.g. weather "none")
        if text == "none":
            if "WEATHER:NONE" in self.vm.vocab:
                tokens.append(self.vm.vocab["WEATHER:NONE"])
                return

        # Try prefixes
        for category in [
            "MON",
            "MOVE",
            "ABIL",
            "ITEM",
            "TYPE",
            "STATUS",
            "MECH",
            "STAT",
        ]:
            candidate = self.vm.formatted_token(category, text)
            if candidate in self.vm.vocab:
                tokens.append(self.vm.vocab[candidate])
                return

        # Check for numbers (HP tokens: HP:100)
        # HP format often: 100/100 or 45/100
        if "/" in text:
            try:
                current, total = text.split("/")
                # Clean up status if attached (e.g. 90/100 par)
                current = current.split(" ")[0]
                total_part = total.strip()

                status_suffix = None
                if " " in total_part:
                    # "100 brn"
                    total_num_str, status_suffix = total_part.split(" ", 1)
                else:
                    total_num_str = total_part

                if total_num_str.isdigit():
                    # Calculate percentage
                    num_curr = float(current)
                    num_tot = float(total_num_str)
                    pct = int((num_curr / num_tot) * 100)
                    tokens.append(self.vm.tokenize(f"HP:{pct}"))

                    if status_suffix:
                        self._blind_tokenize_part(status_suffix, tokens)
                    return
            except:
                # Ignore parsing errors
                pass

    def _parse_player_ref(self, text: str, tokens: List[int]):
        # "p1a: Incineroar"
        # Extract Player (P1/P2) and Slot (A/B)
        player_prefix = text[:3]  # p1a or p1

        if "p1" in player_prefix:
            tokens.append(self.vm.tokenize(self.vm.PLAYER_1))
        elif "p2" in player_prefix:
            tokens.append(self.vm.tokenize(self.vm.PLAYER_2))

        # Check for Side/Slot (a/b/c)
        if "a" in player_prefix or "a:" in text[:4]:
            tokens.append(self.vm.tokenize(self.vm.SIDE_A))
        elif "b" in player_prefix or "b:" in text[:4]:
            tokens.append(self.vm.tokenize(self.vm.SIDE_B))

        remaining = text.split(":", 1)[1].strip() if ":" in text else text

        # Resolve nickname if possible
        species_name = remaining
        if text in self.nickname_map:
            species_name = self.nickname_map[text]

        # Try to find this name as a Pokemon
        mon_token = self.vm.formatted_token("MON", species_name)
        if mon_token in self.vm.vocab:
            tokens.append(self.vm.vocab[mon_token])

    def _parse_move(self, parts, tokens):
        # |move|Source|MoveName|Target...
        if len(parts) > 2:
            self._parse_player_ref(parts[2], tokens)
        if len(parts) > 3:
            move_name = parts[3]
            tokens.append(self.vm.tokenize(self.vm.formatted_token("MOVE", move_name)))

        if len(parts) > 4:
            self._blind_tokenize_part(parts[4], tokens)

    def _parse_switch(self, parts, tokens, p1_roster, p2_roster):
        # |switch|p1a: Nickname|Species|HP
        player_ref = ""
        species = ""

        if len(parts) > 3:
            # Species (e.g. "Landorus-Therian, M")
            species_raw = parts[3]
            species = species_raw.split(",")[0].strip()  # Remove ", L50" etc

            # Update nickname map
            if len(parts) > 2:
                player_ref = parts[2]
                self.nickname_map[player_ref] = species

        if len(parts) > 2:
            player_ref = parts[2]
            self._parse_player_ref(player_ref, tokens)  # Source

        if species:
            # Inject Type Matrix Knowledge
            dex_key = re.sub(r"[^a-z0-9]", "", species.lower())
            types = self._get_types(dex_key)
            for t in types:
                tokens.append(self.vm.tokenize(self.vm.formatted_token("TYPE", t)))

            # Emit Swap Action Token (if we can find it in roster)
            # Determine which roster to check
            roster = p1_roster if "p1" in player_ref else p2_roster

            # Find index (0-based)
            if species in roster:
                idx = roster.index(species)
                # Emit SWAP:{idx+1} (1-based slot)
                tokens.append(self.vm.tokenize(f"SWAP:{idx+1}"))

        if len(parts) > 4:
            # HP
            self._blind_tokenize_part(parts[4], tokens)

    def _get_types(self, dex_key: str) -> List[str]:
        if dex_key in self.pokedex:
            return self.pokedex[dex_key].get("types", [])
        # Fallback: maybe strip cosmetic formes?
        return []

    def _parse_cant(self, parts, tokens):
        if len(parts) > 2:
            self._parse_player_ref(parts[2], tokens)

    def _parse_faint(self, parts, tokens):
        if len(parts) > 2:
            tokens.append(self.vm.tokenize(self.vm.FAINT))
            self._parse_player_ref(parts[2], tokens)