auto_server.py

import asyncio
import sys
import websockets
import requests
import re
import json
import os
from dotenv import load_dotenv

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)
from src.monte_carlo_search import TURN, TEAM_PREVIEW, FAINT


load_dotenv()

# WebSocket URL for Pokémon Showdown
WS_URL = "wss://sim3.psim.us/showdown/websocket"
# Server URL for predictions
SERVER_URL = "http://localhost:5000/"

# Global state for the auto bot
bot_state = {
    "running": True,
    "room_id": os.getenv("SHOWDOWN_ROOM_ID"),
    "bot_name": os.getenv("SHOWDOWN_BOT_NAME"),
    "password": os.getenv("SHOWDOWN_PASSWORD"),
    "team_url": os.getenv("SHOWDOWN_TEAM_URL"),
    "team_data": None,
    "packed_team": os.getenv("SHOWDOWN_PACKED_TEAM"),
    "websocket": None,
    "bot_side": None,
    "accumulated_log": [],
    "my_full_team": [],
    "status": "Idle",
    "last_action": None,
}


def fetch_pokepaste(url):
    if not url.endswith("/raw"):
        if url.endswith("/"):
            url = url[:-1]
        url = url + "/raw"
    try:
        response = requests.get(url)
        response.raise_for_status()
        return response.text
    except Exception as e:
        print(f"Error fetching Pokepaste: {e}")
        return None


def parse_team(text: str):
    blocks = re.split(r"\n\s*\n", text.strip())
    team = {}
    for block in blocks:
        lines = [line.strip() for line in block.split("\n") if line.strip()]
        if not lines:
            continue
        mon = {
            "species": None,
            "item": None,
            "ability": None,
            "tera_type": None,
            "moves": [],
        }
        first_line = lines[0]
        if " @ " in first_line:
            parts = first_line.split(" @ ")
            name_part = parts[0].strip()
            mon["item"] = parts[1].strip()
        else:
            name_part = first_line.strip()
        species_match = re.search(r"\(([^)]+)\)", name_part)
        if species_match:
            mon["species"] = species_match.group(1).strip()
        else:
            mon["species"] = name_part.strip()
        mon["species"] = re.sub(r"\s*\([MF]\)$", "", mon["species"])
        for line in lines[1:]:
            if line.startswith("Ability:"):
                mon["ability"] = line.replace("Ability:", "").strip()
            elif line.startswith("Tera Type:"):
                mon["tera_type"] = line.replace("Tera Type:", "").strip()
            elif line.startswith("- "):
                mon["moves"].append(line.replace("- ", "").strip())
        if mon["species"]:
            team[mon["species"].lower()] = mon
    return team


def _normalize_move_name(name: str) -> str:
    """Convert MOVE:ICE_SPINNER or 'Ice Spinner' to showdown format 'icespinner'."""
    name = name.split(":", 1)[-1] if ":" in name else name
    return re.sub(r"[^a-z0-9]", "", name.lower())


def _normalize_mon_name(name: str) -> str:
    """Convert MON:CHIEN_PAO or 'Chien Pao' to normalized 'chienpao'."""
    name = name.split(":", 1)[-1] if ":" in name else name
    return re.sub(r"[^a-z0-9]", "", name.lower())


def _find_team_index(mon_token: str) -> int:
    """Find 1-based team index for a mon token like MON:INCINEROAR."""
    norm = _normalize_mon_name(mon_token)
    for i, (_, species) in enumerate(bot_state["my_full_team"]):
        if norm in _normalize_mon_name(species):
            return i + 1
    return -1


# Move target types from Showdown that do NOT require a target number.
# Only moves whose target field is "normal", "any", "adjacentAlly",
# "adjacentFoe", or "adjacentAllyOrSelf" need an explicit target.
NO_TARGET_KINDS = frozenset(
    {
        "self",  # Protect, Swords Dance, etc.
        "allAdjacentFoes",  # Dazzling Gleam, Heat Wave, etc.
        "allAdjacent",  # Earthquake, Surf, etc.
        "all",  # Perish Song, Haze, etc.
        "allyTeam",  # Heal Bell, Aromatherapy, etc.
        "allySide",  # Reflect, Light Screen, etc.
        "foeSide",  # Stealth Rock, Spikes, etc.
        "randomNormal",  # Metronome, etc.
        "scripted",  # Struggle, etc.
    }
)


def _get_move_info(mon_moves: list, move_token: str) -> tuple:
    """Find 1-based move index and the move's target kind from request data.

    Returns (index, target_kind).  index=0 means not found."""
    norm_move = _normalize_move_name(move_token)
    for i, m in enumerate(mon_moves):
        if _normalize_move_name(m.get("move", m.get("id", ""))) == norm_move:
            return i + 1, m.get("target", "normal")
    return 0, "normal"


def _slot_action_to_showdown(
    slot_action: dict, slot_key: str, req_data: dict, bot_side: int
) -> str:
    """Convert a single slot action dict to a Showdown /choose fragment.

    slot_key: "[SIDE_A]" or "[SIDE_B]"
    Returns e.g. "move icespinner 1" or "switch 3"
    """
    if slot_action["type"] == "switch":
        idx = _find_team_index(slot_action["mon"])
        if idx > 0 and _is_valid_switch_target(req_data, idx):
            return f"switch {idx}"
        # Mon is fainted / active / not found – pick best valid switch
        fallback = _first_valid_switch(req_data)
        print(
            f"  ↳ Switch to '{slot_action.get('mon')}' invalid, falling back to: {fallback}"
        )
        return fallback

    if slot_action["type"] != "move":
        return "default"

    move_name = _normalize_move_name(slot_action.get("move", ""))
    if not move_name:
        return "default"

    # Look up the move in Showdown's request data to get both the index
    # and the target kind (tells us whether a target number is needed).
    move_str = f"move {move_name}"
    move_target_kind = "normal"  # assume needs target unless we learn otherwise

    active_list = req_data.get("active", []) if req_data else []
    slot_idx = 0 if slot_key == "[SIDE_A]" else 1
    if slot_idx < len(active_list):
        available_moves = active_list[slot_idx].get("moves", [])
        midx, move_target_kind = _get_move_info(
            available_moves, slot_action.get("move", "")
        )
        if midx > 0:
            move_str = f"move {midx}"
        else:
            # Model suggested a move this pokemon can't use right now
            # (e.g. locked out, not in moveset). Use first available instead.
            fallback = _first_available_move(active_list[slot_idx], slot_key, bot_side)
            print(f"  ↳ Move '{move_name}' not available, falling back to: {fallback}")
            return fallback

    # Only append a target number when the move actually requires one.
    needs_target = move_target_kind not in NO_TARGET_KINDS

    if needs_target:
        target_player = slot_action.get("target_player")
        target_side = slot_action.get("target_side")

        # Tier 2: search lower-ranked rollout groups for same move with a target
        if not target_player or not target_side:
            all_moves = getattr(_parse_turn_action, "_all_best_moves", [])
            target_player, target_side = _infer_target_from_alternatives(
                slot_key, slot_action.get("move", ""), all_moves
            )

        if target_player and target_side:
            our_player = f"[P{bot_side}]"
            is_opponent = target_player != our_player

            if is_opponent:
                target_num = 1 if target_side == "[SIDE_A]" else 2
            else:
                target_num = -1 if target_side == "[SIDE_A]" else -2
            move_str += f" {target_num}"
        else:
            # Tier 3: smart default based on move target type
            target_num = _default_target(slot_key, move_target_kind, bot_side)
            print(
                f"  ↳ No target for {slot_action.get('move', '?')}, defaulting to {target_num}"
            )
            move_str += f" {target_num}"

    # Handle tera
    if slot_action.get("tera"):
        move_str += " terastallize"

    return move_str


def _first_available_move(
    active_slot: dict, slot_key: str = "[SIDE_A]", bot_side: int = 1
) -> str:
    """Return the first non-disabled move fragment, *with a target* if needed."""
    for i, m in enumerate(active_slot.get("moves", []), 1):
        if not m.get("disabled", False):
            frag = f"move {i}"
            target_kind = m.get("target", "normal")
            if target_kind not in NO_TARGET_KINDS:
                target_num = _default_target(slot_key, target_kind, bot_side)
                frag += f" {target_num}"
            return frag
    return "move 1 1"  # absolute last resort (with a target)


def _is_valid_switch_target(req_data, team_index_1based: int) -> bool:
    """Check whether the Pokémon at this 1-based team slot can be switched to."""
    if not req_data:
        return True  # can't validate, assume ok
    pokemon_list = req_data.get("side", {}).get("pokemon", [])
    idx = team_index_1based - 1
    if idx < 0 or idx >= len(pokemon_list):
        return False
    mon = pokemon_list[idx]
    if "fnt" in mon.get("condition", ""):
        return False
    if mon.get("active", False):
        return False
    return True


def _first_valid_switch(req_data) -> str:
    """Return the first team slot that is alive and not currently active."""
    if req_data:
        for i, mon in enumerate(req_data.get("side", {}).get("pokemon", []), 1):
            if "fnt" not in mon.get("condition", "") and not mon.get("active", False):
                return f"switch {i}"
    return "default"


def _infer_target_from_alternatives(
    slot_key: str, move_token: str, all_best_moves: list
) -> tuple:
    """Search lower-ranked rollout groups for a target for the same move.

    Returns (target_player, target_side) or (None, None).
    """
    norm = _normalize_move_name(move_token)
    for candidate in all_best_moves:
        sa = candidate.get("slot_actions", {}).get(slot_key)
        if sa and sa.get("type") == "move":
            if _normalize_move_name(sa.get("move", "")) == norm:
                tp = sa.get("target_player")
                ts = sa.get("target_side")
                if tp and ts:
                    return tp, ts
    return None, None


def _default_target(slot_key: str, target_kind: str, bot_side: int) -> int:
    """Pick a sensible default target when the model didn't specify one.

    Offensive moves target the opponent slot across from the attacker.
    Ally-targeting moves target our own other slot.
    """
    if target_kind == "adjacentAlly":
        return -2 if slot_key == "[SIDE_A]" else -1
    if target_kind == "adjacentAllyOrSelf":
        return -1 if slot_key == "[SIDE_A]" else -2
    # "normal", "adjacentFoe", "any" → target the opponent
    return 1 if slot_key == "[SIDE_A]" else 2


def _mine_slot_from_alternatives(
    missing_slot: str, primary_actions: dict, all_best_moves: list
) -> dict | None:
    """Search lower-ranked rollout groups for an action for *missing_slot*.

    Prefers a candidate whose OTHER slot matches the primary actions
    (i.e. same [SIDE_A] move → look for its companion [SIDE_B]).
    Falls back to any candidate that has the missing slot.
    """
    other_slot = "[SIDE_B]" if missing_slot == "[SIDE_A]" else "[SIDE_A]"
    primary_other = primary_actions.get(other_slot)
    primary_sig = None
    if primary_other:
        primary_sig = (
            primary_other.get("type", ""),
            _normalize_move_name(primary_other.get("move", "")),
            primary_other.get("mon", ""),
        )

    # First pass: find a candidate whose companion slot matches our primary action
    best_fallback = None
    for candidate in all_best_moves:
        sa = candidate.get("slot_actions", {})
        if missing_slot not in sa:
            continue
        # This candidate has the slot we need
        if best_fallback is None:
            best_fallback = sa[missing_slot]
        # Check if the companion slot matches our primary
        companion = sa.get(other_slot)
        if companion and primary_sig:
            comp_sig = (
                companion.get("type", ""),
                _normalize_move_name(companion.get("move", "")),
                companion.get("mon", ""),
            )
            if comp_sig == primary_sig:
                return sa[missing_slot]

    return best_fallback


def _parse_turn_action(req_data: dict) -> str:
    """Parse structured turn data into a Showdown /choose command.

    For doubles: /choose move X target, move Y target
    For singles: /choose move X target
    """
    # The best_move dict is stored on the module level by the caller
    slot_actions = getattr(_parse_turn_action, "_slot_actions", None)
    if not slot_actions:
        print(f"No structured slot_actions, falling back to default.")
        return "/choose default"

    bot_side = bot_state["bot_side"] or 1
    active_list = req_data.get("active", []) if req_data else []

    all_best_moves = getattr(_parse_turn_action, "_all_best_moves", [])

    parts = []
    for slot_idx, slot_key in enumerate((("[SIDE_A]", "[SIDE_B]"))):
        # Only generate a choice for slots Showdown is asking us to fill
        if slot_idx >= len(active_list):
            break

        if slot_key in slot_actions:
            fragment = _slot_action_to_showdown(
                slot_actions[slot_key], slot_key, req_data, bot_side
            )
        else:
            # Model didn't generate an action for this slot in the top
            # candidate.  Mine lower-ranked rollouts for a usable action.
            mined = _mine_slot_from_alternatives(slot_key, slot_actions, all_best_moves)
            if mined:
                print(
                    f"  ↳ Mined {slot_key} action from alternative rollout: {mined.get('move', mined.get('mon', '?'))}"
                )
                fragment = _slot_action_to_showdown(mined, slot_key, req_data, bot_side)
            else:
                # Absolute fallback: first available move with target
                fragment = _first_available_move(
                    active_list[slot_idx], slot_key, bot_side
                )
                print(
                    f"  ↳ No model action for {slot_key} in any rollout, falling back to: {fragment}"
                )

        parts.append(fragment)

    if not parts:
        return "/choose default"

    return "/choose " + ", ".join(parts)


async def send_action_to_showdown(
    action_str, objective, room, req_data=None, move_data=None, all_best_moves=None
):
    """
    Sends a command to the Showdown WebSocket server.

    move_data: the full best_move dict from the server (includes slot_actions)
    all_best_moves: all ranked candidates so fallbacks can try alternatives
    """
    if not bot_state["websocket"]:
        return
    cmd = "/choose default"  # Fallback to keep game moving

    # Stash slot_actions so _parse_turn_action can access them
    if move_data and move_data.get("slot_actions"):
        _parse_turn_action._slot_actions = move_data["slot_actions"]
    else:
        _parse_turn_action._slot_actions = None
    _parse_turn_action._all_best_moves = all_best_moves or []

    def normalize_name(name):
        return re.sub(r"[^a-z0-9]", "", name.lower())

    if objective == TEAM_PREVIEW:
        try:
            parts = action_str.replace("Leads:", "").replace("Back:", "").split("|")
            leads = [p.strip() for p in parts[0].split(",")]
            back = [p.strip() for p in parts[1].split(",")]
            all_chosen = leads + back
            indices = []
            for chosen_mon in all_chosen:
                norm_chosen = normalize_name(chosen_mon)
                for i, (_, species) in enumerate(bot_state["my_full_team"]):
                    if norm_chosen in normalize_name(species):
                        indices.append(str(i + 1))
                        break
            if len(indices) == 4:
                cmd = f"/team {''.join(indices)}"
            else:
                print(f"Warning: Could not match all 4 mons. Found indices: {indices}")
        except Exception as e:
            print(f"Failed to parse team preview: {e}")

    elif objective == FAINT:
        try:
            # Try each candidate in LCB order until we find a living, non-active mon
            candidates = all_best_moves or [{"action": action_str}]
            for candidate in candidates:
                cand_action = re.sub(r"<[^>]+>", "", candidate.get("action", ""))
                target_mon = cand_action.replace("Switch to", "").strip()
                norm_target = normalize_name(target_mon)
                found = False
                for i, (_, species) in enumerate(bot_state["my_full_team"]):
                    if norm_target in normalize_name(species):
                        if _is_valid_switch_target(req_data, i + 1):
                            cmd = f"/choose switch {i + 1}"
                            found = True
                        else:
                            print(
                                f"  ↳ {target_mon} is fainted/active, trying next candidate"
                            )
                        break
                if found:
                    break
            if cmd == "/choose default":
                fallback = _first_valid_switch(req_data)
                cmd = f"/choose {fallback}"
                print(f"  ↳ All FAINT candidates invalid, falling back to: {cmd}")
        except Exception as e:
            print(f"Failed to parse faint: {e}")

    elif objective == TURN:
        try:
            cmd = _parse_turn_action(req_data)
        except Exception as e:
            print(f"Failed to parse turn action: {e}")
            cmd = "/choose default"

    full_cmd = f"{room}|{cmd}"
    print("\n" + "=" * 50)
    print(f"🎯 ACTION DECIDED: {action_str}")
    print(f"🚀 SENDING TO SHOWDOWN: {full_cmd}")
    print("=" * 50 + "\n")
    bot_state["last_action"] = full_cmd
    await bot_state["websocket"].send(full_cmd)


async def get_prediction_from_server(objective, room, req_data):
    try:
        full_log = "\n".join(bot_state["accumulated_log"])
        payload = {
            "log": full_log,
            "player": bot_state["bot_side"],
            "num_pokemons": 2,
            "objective": objective,
        }
        if bot_state["team_data"]:
            payload["team_data"] = json.dumps(bot_state["team_data"])

        loop = asyncio.get_event_loop()
        response = await loop.run_in_executor(
            None,
            lambda: requests.post(
                SERVER_URL,
                data=payload,
                headers={"Accept": "application/json"},
                timeout=120,
            ),
        )

        if response.status_code == 200:
            data = response.json()
            if data.get("status") == "success" and data.get("best_moves"):
                all_best_moves = data["best_moves"]
                # For TURN: pick the highest-ranked candidate that has
                # actual slot_actions (skip "WIN" / empty predictions).
                best_move = all_best_moves[0]
                if objective == TURN:
                    n_active = len(req_data.get("active", [])) if req_data else 1
                    for candidate in all_best_moves:
                        sa = candidate.get("slot_actions", {})
                        if not sa:
                            continue
                        # Prefer candidates covering both slots in doubles
                        if n_active >= 2 and "[SIDE_A]" in sa and "[SIDE_B]" in sa:
                            best_move = candidate
                            break
                        # Accept any candidate that has at least one action
                        if n_active < 2 and sa:
                            best_move = candidate
                            break
                    else:
                        # No candidate covers both; pick first with any actions
                        for candidate in all_best_moves:
                            if candidate.get("slot_actions"):
                                best_move = candidate
                                break

                best_action = best_move["action"]
                print(f"Got prediction: {best_action}")
                # Pass all ranked candidates so fallbacks can pick targets / valid switches
                await send_action_to_showdown(
                    best_action, objective, room, req_data, best_move, all_best_moves
                )
        else:
            print(f"Server error: {response.status_code}")
    except Exception as e:
        print(f"Failed to contact server: {e}")


async def listen_to_battle():
    bot_state["status"] = "Connecting..."
    if bot_state["team_url"]:
        team_text = fetch_pokepaste(bot_state["team_url"])
        if team_text:
            bot_state["team_data"] = parse_team(team_text)

    try:
        async with websockets.connect(WS_URL) as websocket:
            bot_state["websocket"] = websocket
            bot_state["status"] = "Connected & Waiting for Challenges"

            async for message in websocket:
                if not bot_state["running"]:
                    break
                if message == "o" or message.startswith("h"):
                    continue

                if "|challstr|" in message:
                    parts = message.split("|")
                    challstr = f"{parts[2]}|{parts[3]}"
                    if bot_state["password"]:
                        res = requests.post(
                            "https://play.pokemonshowdown.com/action.php",
                            data={
                                "act": "login",
                                "name": bot_state["bot_name"],
                                "pass": bot_state["password"],
                                "challstr": challstr,
                            },
                        )
                        assertion = json.loads(res.text[1:])["assertion"]
                    else:
                        res = requests.post(
                            "https://play.pokemonshowdown.com/action.php",
                            data={
                                "act": "getassertion",
                                "userid": bot_state["bot_name"],
                                "challstr": challstr,
                            },
                        )
                        assertion = res.text
                    await websocket.send(f"|/trn {bot_state['bot_name']},0,{assertion}")
                    await websocket.send("|/avatar 167")
                    if bot_state["room_id"]:
                        await websocket.send(f"|/join {bot_state['room_id']}")

                if message.startswith("|updatechallenges|"):
                    try:
                        # The message format is usually: |updatechallenges|{"challengesFrom":{"username":"format"},"challengeTo":null}
                        parts = message.split("|")
                        if len(parts) > 2:
                            data = json.loads(parts[2])
                            challenges = data.get("challengesFrom", {})
                            for challenger, format_id in challenges.items():
                                print(
                                    f"Accepting challenge from {challenger} for format {format_id}"
                                )
                                # Showdown requires the team to be sent when accepting if it's a format that requires a team
                                if bot_state["packed_team"]:
                                    await websocket.send(
                                        f"|/utm {bot_state['packed_team']}"
                                    )
                                else:
                                    await websocket.send(f"|/utm null")
                                await websocket.send(f"|/accept {challenger}")
                    except Exception as e:
                        print(f"Error parsing challenges: {e}")

                # Handle PMs (Private Messages) which are sometimes used for challenges
                if message.startswith("|pm|"):
                    parts = message.split("|")
                    if len(parts) >= 5:
                        sender = parts[2].strip()
                        receiver = parts[3].strip()
                        pm_msg = parts[4].strip()

                        # Ignore our own messages
                        if sender.lower() == bot_state["bot_name"].lower():
                            continue

                        if pm_msg.startswith("/challenge"):
                            print(f"Received challenge via PM from {sender}")
                            if bot_state["packed_team"]:
                                await websocket.send(
                                    f"|/utm {bot_state['packed_team']}"
                                )
                            else:
                                await websocket.send(f"|/utm null")
                            await websocket.send(f"|/accept {sender}")

                if message.startswith(">"):
                    lines = message.split("\n")
                    room = lines[0][1:]

                    if room.startswith("battle-"):
                        bot_state["room_id"] = room

                        for line in lines[1:]:
                            line = line.strip()
                            if not line:
                                continue
                            bot_state["accumulated_log"].append(line)

                            if line.startswith("|player|"):
                                parts = line.split("|")
                                if (
                                    len(parts) >= 4
                                    and parts[3].lower()
                                    == bot_state["bot_name"].lower()
                                ):
                                    bot_state["bot_side"] = 1 if parts[2] == "p1" else 2
                                    bot_state["status"] = (
                                        f"Playing as P{bot_state['bot_side']} in {room}"
                                    )

                            if bot_state["bot_side"] is not None and line.startswith(
                                f"|poke|p{bot_state['bot_side']}|"
                            ):
                                species = line.split("|")[3].split(",")[0].strip()
                                bot_state["my_full_team"].append((line, species))

                            if line.startswith("|error|"):
                                print(f"Showdown Error: {line}")
                                if "[Invalid choice]" in line:
                                    print(
                                        "Invalid choice detected, falling back to default..."
                                    )
                                    asyncio.create_task(
                                        bot_state["websocket"].send(
                                            f"{room}|/choose default"
                                        )
                                    )

                            if line.startswith("|request|"):
                                try:
                                    req_json = line.split("|", 2)[2]
                                    if not req_json:
                                        continue
                                    req_data = json.loads(req_json)
                                    if req_data.get("wait"):
                                        continue

                                    objective = None
                                    if req_data.get("teamPreview"):
                                        objective = TEAM_PREVIEW
                                    elif req_data.get("forceSwitch"):
                                        objective = FAINT
                                    elif req_data.get("active"):
                                        objective = TURN

                                    if objective:
                                        if bot_state["bot_side"] is None:
                                            bot_state["bot_side"] = 2
                                        bot_state["status"] = (
                                            f"Thinking ({objective})..."
                                        )
                                        asyncio.create_task(
                                            get_prediction_from_server(
                                                objective, room, req_data
                                            )
                                        )
                                except Exception as e:
                                    print(f"Error parsing request: {e}")

    except Exception as e:
        bot_state["status"] = f"Error: {e}"
        print(f"WebSocket Error: {e}")
    finally:
        bot_state["running"] = False
        bot_state["websocket"] = None


if __name__ == "__main__":
    print(f"Starting Auto-Bot for {bot_state['bot_name']}...")
    if bot_state["room_id"]:
        print(f"Will automatically join room: {bot_state['room_id']}")
    else:
        print("Waiting for challenges...")

    try:
        asyncio.run(listen_to_battle())
    except KeyboardInterrupt:
        print("\nBot stopped by user.")