from __future__ import annotations
import time
from openai import AsyncOpenAI
from .llm_client import get_llm_client
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from ..environment.config import GameConfig

ToolCall = dict


class Agent:
    def __init__(
        self,
        agent_id: str,
        model: str = "google/gemma-4-26B-A4B-it",
        provider: str = "modal",
        system_prompt: str | None = None,
        config: GameConfig | None = None,
    ):
        self.id = agent_id
        self.model = model
        self.provider = provider
        self.client: AsyncOpenAI = get_llm_client(provider)
        
        if system_prompt:
            self.system_prompt = system_prompt
        else:
            self.system_prompt = self.get_default_system_prompt(config)

    @staticmethod
    def get_default_system_prompt(config: GameConfig | None = None, grid_size: int | None = None) -> str:
        from backend.environment.config import GameConfig
        cfg = config or GameConfig()
        gs = grid_size if grid_size is not None else cfg.grid_size
        kill_sc = cfg.kill_score
        srv_sc = cfg.survival_score_per_turn
        att_dmg = cfg.attack_damage
        att_rng = cfg.attack_range
        
        return (
            "You are a competitive AI agent in a grid-based battle royale game on a "
            f"{gs}x{gs} grid. "
            "Your goal is to maximize your score and be the last agent standing by eliminating others, collecting loot, and surviving.\n"
            "CRITICAL STRATEGIC RULES:\n"
            "1. ALWAYS call observe() at the start of your turn to get details on your HP, abilities, cooldowns, nearby tiles with loot, and visible agents.\n"
            f"2. SCORING: Eliminating an agent yields +{kill_sc} points. Surviving only gives +{srv_sc} point(s) per turn. Passive play will lose. Actively hunt other agents, especially those with low HP.\n"
            "3. ABILITIES: Walk over tiles marked with '[has loot]' to automatically collect them. Once acquired, activate them using activate_ability(ability, args):\n"
            f"   - 'attack': args={{'x': target_x, 'y': target_y}}. Attack an agent within {att_rng} tiles (Manhattan distance <= {att_rng}) for {att_dmg} damage (3 uses). If they are shielded, their shield breaks instead. Focus fire to eliminate them!\n"
            "   - 'dash': args={{'dx': dx, 'dy': dy}}. Teleport up to 3 tiles (dx/dy between -3 and 3) to chase agents, escape, or grab loot (3 uses).\n"
            "   - 'shield': args={}. Activate a shield to completely block the next incoming attack (2 uses).\n"
            "   - 'heal': args={}. Restore 40 HP when your health is low (2 uses).\n"
            "4. MOVE: Use move(dx, dy) to move 1 step to an adjacent tile (dx/dy in -1, 0, 1) and automatically pick up any loot on that tile."
        )

    async def decide(self, messages: list, tools: list[dict]) -> list[ToolCall]:
        import json
        system_msg = {"role": "system", "content": self.system_prompt}
        full_messages = [system_msg] + messages

        t0 = time.time()
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=full_messages,
            tools=tools,
            tool_choice="auto",
        )
        elapsed = round((time.time() - t0) * 1000)

        choice = response.choices[0]
        if not choice.message.tool_calls:
            return {"calls": [], "time_ms": elapsed, "raw": response.usage.model_dump() if response.usage else None}

        results = []
        for tc in choice.message.tool_calls:
            try:
                args = json.loads(tc.function.arguments)
            except json.JSONDecodeError:
                args = {}
            results.append({"name": tc.function.name, "args": args})

        return {
            "calls": results,
            "time_ms": elapsed,
            "raw": response.usage.model_dump() if response.usage else None,
        }