agents.py

import os
import json
import asyncio
import random

# --- OpenAI ---
from openai import AsyncOpenAI, APIError

# --- Google Gemini ---
from google import genai
from google.genai import types

# --- Mistral AI ---
from mistralai import Mistral

# --- Poke-Env ---
from poke_env.player import Player
from poke_env.environment.battle import Battle
from poke_env.environment.move import Move
from poke_env.environment.pokemon import Pokemon
from typing import Optional, Dict, Any, Union

# --- Helper Function & Base Class ---
def normalize_name(name: str) -> str:
    """Lowercase and remove non-alphanumeric characters."""
    return "".join(filter(str.isalnum, name)).lower()

STANDARD_TOOL_SCHEMA = {
    "choose_move": {
        "name": "choose_move",
        "description": "Selects and executes an available attacking or status move.",
        "parameters": {
            "type": "object",
            "properties": {
                "move_name": {
                    "type": "string",
                    "description": "The exact name or ID (e.g., 'thunderbolt', 'swordsdance') of the move to use. Must be one of the available moves.",
                },
            },
            "required": ["move_name"],
        },
    },
    "choose_switch": {
        "name": "choose_switch",
        "description": "Selects an available Pokémon from the bench to switch into.",
        "parameters": {
            "type": "object",
            "properties": {
                "pokemon_name": {
                    "type": "string",
                    "description": "The exact name of the Pokémon species to switch to (e.g., 'Pikachu', 'Charizard'). Must be one of the available switches.",
                },
            },
            "required": ["pokemon_name"],
        },
    },
}

# --- OpenAI Tools Schema (with 'type' field) ---
OPENAI_TOOL_SCHEMA = {
    "choose_move": {
        "type": "function",
        "function": {
            "name": "choose_move",
            "description": "Selects and executes an available attacking or status move.",
            "parameters": {
                "type": "object",
                "properties": {
                    "move_name": {
                        "type": "string",
                        "description": "The exact name or ID (e.g., 'thunderbolt', 'swordsdance') of the move to use. Must be one of the available moves.",
                    },
                },
                "required": ["move_name"],
            },
        }
    },
    "choose_switch": {
        "type": "function",
        "function": {
            "name": "choose_switch",
            "description": "Selects an available Pokémon from the bench to switch into.",
            "parameters": {
                "type": "object",
                "properties": {
                    "pokemon_name": {
                        "type": "string",
                        "description": "The exact name of the Pokémon species to switch to (e.g., 'Pikachu', 'Charizard'). Must be one of the available switches.",
                    },
                },
                "required": ["pokemon_name"],
            },
        }
    },
}


class LLMAgentBase(Player):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.standard_tools = STANDARD_TOOL_SCHEMA
        self.battle_history = []

    def _format_battle_state(self, battle: Battle) -> str:
        active_pkmn = battle.active_pokemon
        active_pkmn_info = f"Your active Pokemon: {active_pkmn.species} " \
                           f"(Type: {'/'.join(map(str, active_pkmn.types))}) " \
                           f"HP: {active_pkmn.current_hp_fraction * 100:.1f}% " \
                           f"Status: {active_pkmn.status.name if active_pkmn.status else 'None'} " \
                           f"Boosts: {active_pkmn.boosts}"

        opponent_pkmn = battle.opponent_active_pokemon
        opp_info_str = "Unknown"
        if opponent_pkmn:
            opp_info_str = f"{opponent_pkmn.species} " \
                           f"(Type: {'/'.join(map(str, opponent_pkmn.types))}) " \
                           f"HP: {opponent_pkmn.current_hp_fraction * 100:.1f}% " \
                           f"Status: {opponent_pkmn.status.name if opponent_pkmn.status else 'None'} " \
                           f"Boosts: {opponent_pkmn.boosts}"
        opponent_pkmn_info = f"Opponent's active Pokemon: {opp_info_str}"

        available_moves_info = "Available moves:\n"
        if battle.available_moves:
            available_moves_info += "\n".join(
                [f"- {move.id} (Type: {move.type}, BP: {move.base_power}, Acc: {move.accuracy}, PP: {move.current_pp}/{move.max_pp}, Cat: {move.category.name})"
                 for move in battle.available_moves]
            )
        else:
             available_moves_info += "- None (Must switch or Struggle)"

        available_switches_info = "Available switches:\n"
        if battle.available_switches:
              available_switches_info += "\n".join(
                  [f"- {pkmn.species} (HP: {pkmn.current_hp_fraction * 100:.1f}%, Status: {pkmn.status.name if pkmn.status else 'None'})"
                   for pkmn in battle.available_switches]
              )
        else:
            available_switches_info += "- None"

        state_str = f"{active_pkmn_info}\n" \
                    f"{opponent_pkmn_info}\n\n" \
                    f"{available_moves_info}\n\n" \
                    f"{available_switches_info}\n\n" \
                    f"Weather: {battle.weather}\n" \
                    f"Terrains: {battle.fields}\n" \
                    f"Your Side Conditions: {battle.side_conditions}\n" \
                    f"Opponent Side Conditions: {battle.opponent_side_conditions}"
        return state_str.strip()

    def _find_move_by_name(self, battle: Battle, move_name: str) -> Optional[Move]:
        normalized_name = normalize_name(move_name)
        # Prioritize exact ID match
        for move in battle.available_moves:
            if move.id == normalized_name:
                return move
        # Fallback: Check display name (less reliable)
        for move in battle.available_moves:
            if move.name.lower() == move_name.lower():
                print(f"Warning: Matched move by display name '{move.name}' instead of ID '{move.id}'. Input was '{move_name}'.")
                return move
        return None

    def _find_pokemon_by_name(self, battle: Battle, pokemon_name: str) -> Optional[Pokemon]:
        normalized_name = normalize_name(pokemon_name)
        for pkmn in battle.available_switches:
            # Normalize the species name for comparison
            if normalize_name(pkmn.species) == normalized_name:
                return pkmn
        return None

    async def choose_move(self, battle: Battle) -> str:
        battle_state_str = self._format_battle_state(battle)
        decision_result = await self._get_llm_decision(battle_state_str)
        print(decision_result)
        decision = decision_result.get("decision")
        error_message = decision_result.get("error")
        action_taken = False
        fallback_reason = ""

        if decision:
            function_name = decision.get("name")
            args = decision.get("arguments", {})
            if function_name == "choose_move":
                move_name = args.get("move_name")
                if move_name:
                    chosen_move = self._find_move_by_name(battle, move_name)
                    if chosen_move and chosen_move in battle.available_moves:
                        action_taken = True
                        chat_msg = f"AI Decision: Using move '{chosen_move.id}'."
                        print(chat_msg)
                        return self.create_order(chosen_move)
                    else:
                        fallback_reason = f"LLM chose unavailable/invalid move '{move_name}'."
                else:
                     fallback_reason = "LLM 'choose_move' called without 'move_name'."
            elif function_name == "choose_switch":
                pokemon_name = args.get("pokemon_name")
                if pokemon_name:
                    chosen_switch = self._find_pokemon_by_name(battle, pokemon_name)
                    if chosen_switch and chosen_switch in battle.available_switches:
                        action_taken = True
                        chat_msg = f"AI Decision: Switching to '{chosen_switch.species}'."
                        print(chat_msg)
                        return self.create_order(chosen_switch)
                    else:
                        fallback_reason = f"LLM chose unavailable/invalid switch '{pokemon_name}'."
                else:
                    fallback_reason = "LLM 'choose_switch' called without 'pokemon_name'."
            else:
                fallback_reason = f"LLM called unknown function '{function_name}'."

        if not action_taken:
            if not fallback_reason:
                 if error_message:
                     fallback_reason = f"API Error: {error_message}"
                 elif decision is None:
                      fallback_reason = "LLM did not provide a valid function call."
                 else:
                      fallback_reason = "Unknown error processing LLM decision."

            print(f"Warning: {fallback_reason} Choosing random action.")

            if battle.available_moves or battle.available_switches:
                 return self.choose_random_move(battle)
            else:
                 print("AI Fallback: No moves or switches available. Using Struggle/Default.")
                 return self.choose_default_move(battle)

    async def _get_llm_decision(self, battle_state: str) -> Dict[str, Any]:
        raise NotImplementedError("Subclasses must implement _get_llm_decision")


# --- Google Gemini Agent ---
class GeminiAgent(LLMAgentBase):
    """Uses Google Gemini API for decisions."""
    def __init__(self, api_key: str = None, model: str = "gemini-2.5-pro-preview-03-25", avatar: str = "steven", *args, **kwargs):
        # Set avatar before calling parent constructor
        kwargs['avatar'] = avatar
        kwargs['start_timer_on_battle_start'] = True
        super().__init__(*args, **kwargs)
        self.model_name = model
        used_api_key = api_key or os.environ.get("GOOGLE_API_KEY")
        if not used_api_key:
            raise ValueError("Google API key not provided or found in GOOGLE_API_KEY env var.")
        
        # Initialize Gemini client using the correct API
        self.genai_client = genai.Client(api_key=used_api_key)
        
        # Configure the tools for function calling
        self.function_declarations = list(self.standard_tools.values())

    async def _get_llm_decision(self, battle_state: str) -> Dict[str, Any]:
        """Sends state to the Gemini API and gets back the function call decision."""
        prompt = (
            "Based on the current battle state, decide the best action: either use an available move or switch to an available Pokémon. "
            "Consider type matchups, HP, status conditions, field effects, entry hazards, and potential opponent actions. "
            "Only choose actions listed as available using their exact ID (for moves) or species name (for switches). "
            "Use the provided functions to indicate your choice.\n\n"
            f"Current Battle State:\n{battle_state}\n\n"
            "Choose the best action by calling the appropriate function ('choose_move' or 'choose_switch')."
        )

        try:
            # Configure tools using the Gemini API format
            tools = genai.types.Tool(function_declarations=self.function_declarations)
            config = genai.types.GenerateContentConfig(tools=[tools],automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=True))
            
            # Send request to the model
            response = self.genai_client.models.generate_content(
                model=self.model_name,
                contents=prompt,
                config=config
            )
            try:
                function_calls = response.function_calls
                if function_calls:
                    function_name = function_calls[0].name
                    arguments = function_calls[0].args
                    return {"decision": {"name": function_name, "arguments": arguments}}
                else:
                    return {"error": "Gemini did not return a function call."}
            except Exception as e:
                    return {"error": f"Model called unknown @function '{function_name}'."}
            
            # No function call found
            return {"error": "Gemini did not return a function call."}
            
        except Exception as e:
            print(f"Unexpected error during Gemini processing: {e}")
            import traceback
            traceback.print_exc()
            return {"error": f"Unexpected error: {str(e)}"}


# --- OpenAI Agent ---
class OpenAIAgent(LLMAgentBase):
    """Uses OpenAI API for decisions."""
    def __init__(self, api_key: str = None, model: str = "gpt-4.1", avatar: str = "giovanni", *args, **kwargs):
        # Set avatar before calling parent constructor
        kwargs['avatar'] = avatar
        kwargs['start_timer_on_battle_start'] = True
        super().__init__(*args, **kwargs)
        self.model = model
        used_api_key = api_key or os.environ.get("OPENAI_API_KEY")
        if not used_api_key:
            raise ValueError("OpenAI API key not provided or found in OPENAI_API_KEY env var.")
        self.openai_client = AsyncOpenAI(api_key=used_api_key)

        # Use the OpenAI-specific schema with type field
        self.openai_tools = list(OPENAI_TOOL_SCHEMA.values())

    async def _get_llm_decision(self, battle_state: str) -> Dict[str, Any]:
        system_prompt = (
            "You are a skilled Pokemon battle AI. Your goal is to win the battle. "
            "Based on the current battle state, decide the best action: either use an available move or switch to an available Pokémon. "
            "Consider type matchups, HP, status conditions, field effects, entry hazards, and potential opponent actions. "
            "Only choose actions listed as available using their exact ID (for moves) or species name (for switches). "
            "Use the provided functions to indicate your choice."
        )
        user_prompt = f"Current Battle State:\n{battle_state}\n\nChoose the best action by calling the appropriate function ('choose_move' or 'choose_switch')."

        try:
            response = await self.openai_client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                tools=self.openai_tools,
                tool_choice="auto",  # Let the model choose
                temperature=0.5,
            )
            message = response.choices[0].message
            print("OPENAI RESPONSE : ",response)
            # Check for tool calls in the response
            if message.tool_calls:
                tool_call = message.tool_calls[0]  # Get the first tool call
                function_name = tool_call.function.name
                try:
                    arguments = json.loads(tool_call.function.arguments or '{}')
                    if function_name in self.standard_tools:
                        return {"decision": {"name": function_name, "arguments": arguments}}
                    else:
                        return {"error": f"Model called unknown function '{function_name}'."}
                except json.JSONDecodeError:
                    return {"error": f"Error decoding function arguments: {tool_call.function.arguments}"}
            else:
                # Model decided not to call a function
                return {"error": f"OpenAI did not return a function call. Response: {message.content}"}

        except APIError as e:
            print(f"Error during OpenAI API call: {e}")
            return {"error": f"OpenAI API Error: {e.status_code} - {e.message}"}
        except Exception as e:
            print(f"Unexpected error during OpenAI API call: {e}")
            return {"error": f"Unexpected error: {e}"}


# --- Mistral Agent ---
class MistralAgent(LLMAgentBase):
    """Uses Mistral AI API for decisions."""
    def __init__(self, api_key: str = None, model: str = "mistral-large-latest", avatar: str = "alder", *args, **kwargs):
        # Set avatar before calling parent constructor
        kwargs['avatar'] = avatar
        kwargs['start_timer_on_battle_start'] = True
        super().__init__(*args, **kwargs)
        self.model = model
        used_api_key = api_key or os.environ.get("MISTRAL_API_KEY")
        if not used_api_key:
            raise ValueError("Mistral API key not provided or found in MISTRAL_API_KEY env var.")
        self.mistral_client = Mistral(api_key=used_api_key)

        # Convert standard schema to Mistral's tool format with "function" wrapper
        self.mistral_tools = []
        for tool_name, tool_schema in self.standard_tools.items():
            self.mistral_tools.append({
                "type": "function",
                "function": {
                    "name": tool_schema["name"],
                    "description": tool_schema["description"],
                    "parameters": tool_schema["parameters"]
                }
            })

    async def _get_llm_decision(self, battle_state: str) -> Dict[str, Any]:
        system_prompt = (
            "You are a skilled Pokemon battle AI. Your goal is to win the battle. "
            "Based on the current battle state, decide the best action: either use an available move or switch to an available Pokémon. "
            "Consider type matchups, HP, status conditions, field effects, entry hazards, and potential opponent actions. "
            "Only choose actions listed as available using their exact ID (for moves) or species name (for switches). "
            "Use the provided tools to indicate your choice."
        )
        user_prompt = f"Current Battle State:\n{battle_state}\n\nChoose the best action by calling the appropriate function ('choose_move' or 'choose_switch')."

        try:
            # Create the messages array
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ]
            
            # Call the Mistral API with tool_choice set to "any" to force tool usage
            response = self.mistral_client.chat.complete(
                model=self.model,
                messages=messages,
                tools=self.mistral_tools,
                tool_choice="any",  # Force the model to use a tool
                temperature=0.3,
            )
            print("Mistral RESPONSE : ", response)
            
            # Check for tool calls in the response
            message = response.choices[0].message
            if hasattr(message, 'tool_calls') and message.tool_calls:
                tool_call = message.tool_calls[0]  # Get the first tool call
                function_name = tool_call.function.name
                try:
                    # Parse the function arguments from JSON string
                    arguments = json.loads(tool_call.function.arguments or '{}')
                    if function_name in self.standard_tools:
                        return {"decision": {"name": function_name, "arguments": arguments}}
                    else:
                        return {"error": f"Model called unknown function '{function_name}'."}
                except json.JSONDecodeError:
                    return {"error": f"Error decoding function arguments: {tool_call.function.arguments}"}
            else:
                # Model did not return a tool call
                return {"error": f"Mistral did not return a tool call. Response: {message.content}"}

        except Exception as e:
            print(f"Error during Mistral API call: {e}")
            import traceback
            traceback.print_exc()
            return {"error": f"Unexpected error: {str(e)}"}