vocab_manager.py

import json
import os
import requests
import re
from typing import Dict


class VocabManager:
    """
    Manages the vocabulary of tokens for the Pokemon Showdown AI.
    Features:
    - Fetches official data from Pokemon Showdown (Pokemon, Moves, Abilities, Items, Types)
    - Builds a deterministic mapping from string -> ID
    - Handles special tokens (PAD, SOS, EOS, UNK, etc.)
    """

    DATA_URLS = {
        "pokedex": "https://play.pokemonshowdown.com/data/pokedex.json",
        "moves": "https://play.pokemonshowdown.com/data/moves.json",
        "abilities": "https://play.pokemonshowdown.com/data/abilities.js",
        "items": "https://play.pokemonshowdown.com/data/items.js",
        "types": "https://play.pokemonshowdown.com/data/typechart.js",
    }

    # Special Tokens
    PAD_TOKEN = "[PAD]"
    UNK_TOKEN = "[UNK]"
    SOS_TOKEN = "[SOS]"  # Start of Sequence (Start of Battle or Turn)
    EOS_TOKEN = "[EOS]"  # End of Sequence

    # Game Logic Tokens
    PLAYER_1 = "[P1]"
    PLAYER_2 = "[P2]"
    SIDE_A = "[SIDE_A]"
    SIDE_B = "[SIDE_B]"
    WIN = "[WIN]"
    LOSE = "[LOSE]"
    FAINT = "[FAINT]"

    # Mechanics and Log Command Tokens
    TURN = "[TURN]"
    START = "[START]"
    MOVE_CMD = "[CMD_MOVE]"  # To avoid conflict with MOVE category
    SWITCH = "[SWITCH]"
    DRAG = "[DRAG]"
    DETAILS_CHANGE = "[DETAILS_CHANGE]"

    DAMAGE = "[DAMAGE]"
    HEAL = "[HEAL]"
    STATUS = "[STATUS]"  # Generic or prefix? Usually |-status|...
    CURE_STATUS = "[CURE_STATUS]"
    CURE_TEAM = "[CURE_TEAM]"
    BOOST = "[BOOST]"
    UNBOOST = "[UNBOOST]"
    SET_BOOST = "[SET_BOOST]"
    CLEAR_BOOST = "[CLEAR_BOOST]"
    CLEAR_POSITIVE_BOOST = "[CLEAR_POSITIVE_BOOST]"
    CLEAR_NEGATIVE_BOOST = "[CLEAR_NEGATIVE_BOOST]"
    COPY_BOOST = "[COPY_BOOST]"
    SWAP_BOOST = "[SWAP_BOOST]"
    INVERT_BOOST = "[INVERT_BOOST]"

    WEATHER = "[WEATHER]"
    FIELD_START = "[FIELD_START]"
    FIELD_END = "[FIELD_END]"
    SIDE_START = "[SIDE_START]"
    SIDE_END = "[SIDE_END]"

    CRIT = "[CRIT]"
    SUPEREFFECTIVE = "[SUPEREFFECTIVE]"
    RESISTED = "[RESISTED]"
    IMMUNE = "[IMMUNE]"
    MISS = "[MISS]"
    FAIL = "[FAIL]"
    BLOCK = "[BLOCK]"
    ACTIVATE = "[ACTIVATE]"

    ITEM_CMD = "[CMD_ITEM]"  # |-item|
    END_ITEM = "[END_ITEM]"
    ABILITY_CMD = "[CMD_ABILITY]"  # |-ability|
    END_ABILITY = "[END_ABILITY]"

    TRANSFORM = "[TRANSFORM]"
    FORM_CHANGE = "[FORM_CHANGE]"
    MEGA = "[MEGA]"
    PRIMAL = "[PRIMAL]"
    BURST = "[BURST]"
    Z_POWER = "[Z_POWER]"
    TERASTALLIZE = "[TERASTALLIZE]"
    DYNAMAX = "[DYNAMAX]"
    MAX_GUARD = "[MAX_GUARD]"  # Not a command but often special?
    SINGLE_TURN = "[SINGLE_TURN]"
    SINGLE_MOVE = "[SINGLE_MOVE]"

    CLEARED_POKE = "[CLEARED_POKE]"
    POKE = "[POKE]"
    TEAMS_PREVIEW = "[TEAMS_PREVIEW]"

    def __init__(self, cache_dir: str = "data/cache"):
        self.cache_dir = cache_dir
        self.vocab: Dict[str, int] = {}
        self.reverse_vocab: Dict[int, str] = {}
        self.special_tokens = [
            self.PAD_TOKEN,
            self.UNK_TOKEN,
            self.SOS_TOKEN,
            self.EOS_TOKEN,
            self.PLAYER_1,
            self.PLAYER_2,
            self.SIDE_A,
            self.SIDE_B,
            self.WIN,
            self.LOSE,
            self.FAINT,
            self.TURN,
            self.START,
            self.MOVE_CMD,
            self.SWITCH,
            self.DRAG,
            self.DETAILS_CHANGE,
            self.DAMAGE,
            self.HEAL,
            self.STATUS,
            self.CURE_STATUS,
            self.CURE_TEAM,
            self.BOOST,
            self.UNBOOST,
            self.SET_BOOST,
            self.CLEAR_BOOST,
            self.CLEAR_POSITIVE_BOOST,
            self.CLEAR_NEGATIVE_BOOST,
            self.COPY_BOOST,
            self.SWAP_BOOST,
            self.INVERT_BOOST,
            self.WEATHER,
            self.FIELD_START,
            self.FIELD_END,
            self.SIDE_START,
            self.SIDE_END,
            self.CRIT,
            self.SUPEREFFECTIVE,
            self.RESISTED,
            self.IMMUNE,
            self.MISS,
            self.FAIL,
            self.BLOCK,
            self.ACTIVATE,
            self.ITEM_CMD,
            self.END_ITEM,
            self.ABILITY_CMD,
            self.END_ABILITY,
            self.TRANSFORM,
            self.FORM_CHANGE,
            self.MEGA,
            self.PRIMAL,
            self.BURST,
            self.Z_POWER,
            self.TERASTALLIZE,
            self.DYNAMAX,
            self.SINGLE_TURN,
            self.SINGLE_MOVE,
            self.CLEARED_POKE,
            self.POKE,
            self.TEAMS_PREVIEW,
        ]

    def formatted_token(self, category: str, name: str) -> str:
        """Formats a token with its category prefix."""
        return f"{category}:{name.upper().replace(' ', '_').replace('-', '_')}"

    def build_vocab(self):
        """Fetches data and builds the vocabulary dictionary."""
        print("Building vocabulary...")
        os.makedirs(self.cache_dir, exist_ok=True)

        tokens = set(self.special_tokens)

        # 1. Fetch and process Pokedex
        pokedex = self._fetch_json("pokedex")
        for key, mon in pokedex.items():
            tokens.add(self.formatted_token("MON", mon["name"]))  # e.g. MON:CHARIZARD

        # 2. Fetch and process Moves
        moves = self._fetch_json("moves")
        for key, move in moves.items():
            tokens.add(self.formatted_token("MOVE", move["name"]))

        # 3. Fetch and process Abilities
        abilities = self._fetch_json("abilities")
        for key, ability in abilities.items():
            tokens.add(self.formatted_token("ABIL", ability["name"]))

        # 4. Fetch and process Items
        items = self._fetch_json("items")
        for key, item in items.items():
            tokens.add(self.formatted_token("ITEM", item["name"]))

        # 5. Fetch and process Types
        types = self._fetch_json("types")
        for key in types.keys():
            tokens.add(self.formatted_token("TYPE", key))

        # 6. Manual Additions (Status, Mechanics, Keywords, Stats)
        for s in ["brn", "par", "slp", "frz", "psn", "tox"]:
            tokens.add(self.formatted_token("STATUS", s))

        for m in ["drain", "recoil", "confusion", "trap", "lockedmove", "struggle"]:
            tokens.add(self.formatted_token("MECH", m))

        for stat in ["atk", "def", "spa", "spd", "spe", "accuracy", "evasion"]:
            tokens.add(self.formatted_token("STAT", stat))

        tokens.add(self.formatted_token("WEATHER", "none"))
        tokens.add(self.formatted_token("KW", "spread"))
        tokens.add(self.formatted_token("KW", "notarget"))

        # 7. Add HP buckets (0-100)
        for i in range(101):
            tokens.add(self.formatted_token("HP", str(i)))

        # 8. Add Swap Tokens (Slots 1-6)
        for i in range(1, 7):
            tokens.add(self.formatted_token("SWAP", str(i)))

        # Sort tokens for deterministic ordering
        sorted_tokens = sorted(list(tokens))

        # Build dicts
        for idx, token in enumerate(sorted_tokens):
            self.vocab[token] = idx
            self.reverse_vocab[idx] = token

        print(f"Vocabulary built with {len(self.vocab)} tokens.")
        self.save_vocab()

    def _parse_js_content(self, content: str) -> dict:
        # Strip exports assignment
        content = content.strip()
        # Handle "exports.Foo = {...}" or "var Foo = {...}"
        if "=" in content:
            content = content.split("=", 1)[1].strip()
        if content.endswith(";"):
            content = content[:-1].strip()

        # Regex to capture strings (double and single quoted) OR keys
        # Group 1: Double Quoted String
        # Group 2: Single Quoted String
        # Group 3: Key (identifier followed by colon)
        pattern = r'("[^"\\]*(?:\\.[^"\\]*)*"|\'[^\'\\]*(?:\\.[^\'\\]*)*\')|([a-zA-Z0-9_]+)\s*:'

        def repl(match):
            if match.group(1):
                return match.group(1)  # Return the string as is
            else:
                return f'"{match.group(2)}":'  # Quote the key

        json_str = re.sub(pattern, repl, content)

        # Javascript might have trailing commas which JSON does not allow.
        # regex to remove trailing commas: ,} -> } and ,] -> ]
        json_str = re.sub(r",\s*}", "}", json_str)
        json_str = re.sub(r",\s*]", "]", json_str)

        return json.loads(json_str)

    def _fetch_json(self, name: str) -> dict:
        path = os.path.join(self.cache_dir, f"{name}.json")
        if os.path.exists(path):
            with open(path, "r") as f:
                return json.load(f)

        url = self.DATA_URLS[name]
        print(f"Downloading {name} from {url}...")

        if url.endswith(".json"):
            resp = requests.get(url)
            if resp.status_code != 200:
                print(f"Failed to fetch {url}: {resp.status_code}")
                return {}
            try:
                data = resp.json()
            except Exception as e:
                print(f"Error parsing JSON for {name}: {e}")
                return {}
        else:
            # Handle .js files via internal Python parsing
            # 1. Download .js file
            js_path = os.path.join(self.cache_dir, f"{name}.js")
            resp = requests.get(url)
            if resp.status_code != 200:
                print(f"Failed to fetch {url}: {resp.status_code}")
                return {}

            with open(js_path, "w") as f:
                f.write(resp.text)

            # 2. Parse JS content using regex
            try:
                data = self._parse_js_content(resp.text)
            except Exception as e:
                print(f"Error parsing {name}.js: {e}")
                return {}

        with open(path, "w") as f:
            json.dump(data, f)
        return data

    def get_pokedex(self) -> dict:
        """Returns the loaded Pokedex dictionary."""
        return self._fetch_json("pokedex")

    def save_vocab(self):
        path = os.path.join(self.cache_dir, "vocab.json")
        with open(path, "w") as f:
            json.dump({"vocab": self.vocab}, f, indent=2)

    def load_vocab(self):
        path = os.path.join(self.cache_dir, "vocab.json")
        if not os.path.exists(path):
            self.build_vocab()
        else:
            with open(path, "r") as f:
                data = json.load(f)
                self.vocab = data["vocab"]
                self.reverse_vocab = {v: k for k, v in self.vocab.items()}

    def tokenize(self, text: str) -> int:
        """Simple lookup"""
        return self.vocab.get(text, self.vocab[self.UNK_TOKEN])

    def decode(self, token_id: int) -> str:
        return self.reverse_vocab.get(token_id, self.UNK_TOKEN)

    def __len__(self):
        return len(self.vocab)