"""AI-driven Bot Agent for botWebWars, powered by a local Ollama LLM.

Replicates the general capabilities of botagent/bot_agent.py (registration,
radar-based navigation, obstacle avoidance, party formation, battles) but
delegates the *strategic* decisions to an Ollama model:
  - Whether to propose a voluntary alliance with another solo bot.
  - Which direction to move towards when navigating (given radar/obstacle data).

Outcomes mandated by GAME_RULES.md (forced battles, forced joins based on
relative strength) are always resolved deterministically by the game engine
regardless of what the LLM prefers - the LLM is only ever offered a choice
among legal options.
"""

import json
import os
import re
import time
import argparse
from typing import Any, Dict, List, Optional

import requests

DEFAULT_SERVER_URL = "http://localhost:8000/api"
DEFAULT_BOT_NAME = "OllamaBot"
DEFAULT_BOT_COLOR = "#8b5cf6"
DEFAULT_BOT_STRENGTH = 4

OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://192.168.1.220:11434")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "gemma4:12b")

GAME_RULES_SUMMARY = """
Rules you must respect when choosing among the OPTIONS given to you:
- Two solo bots that meet MAY voluntarily ally (not required). The stronger bot (or higher
  score if tied) leads. Larger parties have an advantage in battle.
- A solo bot always joins a party if the party leader's strength >= its own (no choice).
- A solo bot always refuses and fights if the party leader is weaker (no choice).
- Two opposing parties that meet must always battle (no choice).
- Diagonal squeezes between obstacle corners cost strength (0.1 solo, 0.2 leader / 0.1 followers).
- The game ends when all bots are united into a single party.
You will only ever be asked to choose between options that are legal - always answer with the
requested JSON object and nothing else.
"""


def normalize_url(url: str) -> str:
    """Ensure the API URL ends with /api without trailing slashes."""
    cleaned = url.rstrip("/")
    if not cleaned.endswith("/api"):
        cleaned = f"{cleaned}/api"
    return cleaned


def extract_json(text: str) -> Optional[Dict[str, Any]]:
    """Best-effort extraction of a JSON object from an LLM response."""
    if not text:
        return None
    match = re.search(r"\{.*\}", text, re.DOTALL)
    if not match:
        return None
    try:
        return json.loads(match.group(0))
    except json.JSONDecodeError:
        return None


class OllamaClient:
    def __init__(self, base_url: str, model: str):
        self.base_url = base_url.rstrip("/")
        self.model = model

    def ask_json(self, prompt: str) -> Optional[Dict[str, Any]]:
        """Ask the model a question, requesting a JSON-only response."""
        url = f"{self.base_url}/api/generate"
        payload = {
            "model": self.model,
            "prompt": prompt,
            "format": "json",
            "stream": False,
            "options": {"temperature": 0.4},
        }
        try:
            res = requests.post(url, json=payload, timeout=60)
            res.raise_for_status()
            raw = res.json().get("response", "")
            return extract_json(raw)
        except requests.exceptions.RequestException as e:
            print(f"⚠️ [OLLAMA ERROR] {e}")
            return None


class AIBotAgent:
    def __init__(
        self,
        name: str = DEFAULT_BOT_NAME,
        color: str = DEFAULT_BOT_COLOR,
        strength: int = DEFAULT_BOT_STRENGTH,
        server_url: str = DEFAULT_SERVER_URL,
        ollama_url: str = OLLAMA_BASE_URL,
        ollama_model: str = OLLAMA_MODEL,
    ):
        self.name = name
        self.color = color
        self.strength = strength
        self.base_url = normalize_url(server_url)
        self.llm = OllamaClient(ollama_url, ollama_model)
        self.bot_id: Optional[str] = None
        self.party_id: Optional[str] = None
        self.is_leader: bool = False

    # ------------------------------------------------------------------
    # Registration / status
    # ------------------------------------------------------------------
    def register(self):
        """Register the bot avatar on the grid or reconnect if already present."""
        try:
            players = requests.get(f"{self.base_url}/players").json()
            for p in players:
                if p.get("name") == self.name:
                    self.bot_id = p["id"]
                    print(f"🔄 [RECONNECT] Reconnected to existing {self.name} (ID: {self.bot_id}) at ({p.get('x')}, {p.get('y')})")
                    return
        except Exception:
            pass

        res = requests.post(
            f"{self.base_url}/players",
            json={"name": self.name, "color": self.color, "strength": self.strength},
        )
        if res.status_code == 400 and "already registered" in res.text:
            players = requests.get(f"{self.base_url}/players").json()
            for p in players:
                if p.get("name") == self.name:
                    self.bot_id = p["id"]
                    print(f"🔄 [RECONNECT] Reconnected to existing {self.name} (ID: {self.bot_id}) at ({p.get('x')}, {p.get('y')})")
                    return

        res.raise_for_status()
        data = res.json()
        self.bot_id = data["id"]
        print(f"🚀 [REGISTER] Spawned {self.name} (ID: {self.bot_id}, Str: {self.strength}) at ({data['x']}, {data['y']})")

    def refresh_status(self):
        """Update bot state (party membership, leader status, score)."""
        res = requests.get(f"{self.base_url}/players/{self.bot_id}")
        if res.status_code == 200:
            data = res.json()
            self.party_id = data.get("party_id")
            self.is_leader = data.get("is_party_leader", False)
            return data
        return None

    # ------------------------------------------------------------------
    # Core decision loop
    # ------------------------------------------------------------------
    def decide_and_act(self):
        my_info = self.refresh_status()
        if not my_info:
            return

        print(f"\n🤖 --- Turn for {self.name} | Score: {my_info['score']} | Party: {self.party_id or 'Solo'} ---")

        radar_res = requests.get(f"{self.base_url}/players/{self.bot_id}/radar").json()
        targets = radar_res.get("targets", [])

        adjacent_target = None
        for t in targets:
            if t["distance"] <= 1 and not t["is_ally"]:
                adjacent_target = t
                break

        if adjacent_target:
            self._handle_adjacent_encounter(adjacent_target, my_info)
        else:
            self._navigate_towards_goal(radar_res, my_info)

    def _handle_adjacent_encounter(self, target: Dict[str, Any], my_info: Dict[str, Any]):
        """Resolve the encounter; the LLM only gets a say when the rules allow a choice."""
        target_name = target["name"]
        target_str = target["strength"]
        target_party = target.get("party_id")

        print(f"🔍 [ADJACENT ENCOUNTER] Next to '{target_name}' (Str: {target_str}, Party: {target_party or 'None'})")

        # SCENARIO A: I am a Solo Bot
        if not self.party_id:
            if not target_party:
                # Both solo: alliance is OPTIONAL - ask the LLM.
                self._decide_voluntary_alliance(target, my_info)
            else:
                # Target belongs to a party: joining/refusing is mandated by relative strength.
                party_info = requests.get(f"{self.base_url}/parties/{target_party}").json()
                target_leader_str = target_str
                if party_info:
                    leader_player = requests.get(f"{self.base_url}/players/{party_info['leader_id']}").json()
                    target_leader_str = leader_player.get("strength", 1)

                if self.strength <= target_leader_str:
                    print(f"🤝 [RULE] Party leader strength {target_leader_str} >= my {self.strength}. Willingly joining squad!")
                    self._step_or_attack(target)
                else:
                    print(f"⚔️ [RULE] Party leader is weaker ({target_leader_str} < my {self.strength}). Must refuse and fight!")
                    self._initiate_battle(target["id"])

        # SCENARIO B: I am in a Party
        else:
            if not self.is_leader:
                print("🛡️ [PARTY MEMBER] Under command of party leader. Awaiting leader movement.")
                requests.post(f"{self.base_url}/players/{self.bot_id}/pass")
                return

            if not target_party:
                if target_str <= self.strength:
                    print(f"🤝 [RULE] Solo bot {target_name} is willing to join our squad under my leadership.")
                    self._step_or_attack(target)
                else:
                    print(f"⚔️ [RULE] Solo bot {target_name} refuses weaker leader! Squad is attacking!")
                    self._initiate_battle(target["id"])
            else:
                print(f"⚔️ [RULE] Hostile party detected: '{target.get('party_name')}'! Battle is mandatory!")
                self._initiate_battle(target["id"])

    def _decide_voluntary_alliance(self, target: Dict[str, Any], my_info: Dict[str, Any]):
        """Ask the LLM whether to propose a voluntary alliance with another solo bot."""
        prompt = f"""{GAME_RULES_SUMMARY}
You are bot "{self.name}" (strength {self.strength}, score {my_info['score']}).
You just encountered another solo bot "{target['name']}" (strength {target['strength']}).
Whoever has greater strength (or higher score if tied) will lead the new party.
Forming an alliance is optional - larger parties are stronger in future battles, but you
give up independent control if you are not the stronger one.

Respond ONLY with JSON: {{"form_alliance": true|false, "reasoning": "short reason"}}
"""
        decision = self.llm.ask_json(prompt) or {}
        form_alliance = decision.get("form_alliance", True)
        reasoning = decision.get("reasoning", "")

        if form_alliance:
            leader_id = self.bot_id if self.strength >= target["strength"] else target["id"]
            print(f"🤝 [LLM DECISION] Ally with {target['name']}! Leader: {'me' if leader_id == self.bot_id else target['name']}. {reasoning}")
            self._execute_party_formation([self.bot_id, target["id"]], leader_id)
        else:
            print(f"🚶 [LLM DECISION] Declining alliance with {target['name']}. {reasoning}")
            requests.post(f"{self.base_url}/players/{self.bot_id}/pass")

    def _execute_party_formation(self, member_ids: List[str], leader_id: str):
        """Form a party using the REST API."""
        try:
            res = requests.post(
                f"{self.base_url}/parties",
                json={"member_ids": member_ids, "leader_id": leader_id, "name": f"Squad_{self.name}"},
            )
            if res.status_code == 201:
                party = res.json()
                print(f"✅ [PARTY FORMED] Squad '{party['name']}' established! Leader: {party['leader_name']} | Str: {party['total_strength']}")
            else:
                requests.post(f"{self.base_url}/players/{self.bot_id}/pass")
        except Exception as e:
            print(f"Party formation error: {e}")
            requests.post(f"{self.base_url}/players/{self.bot_id}/pass")

    def _initiate_battle(self, opponent_id: str):
        """Explicitly call the 3-Bout D20 Battle endpoint."""
        print(f"🎲 [BATTLE INITIATED] Clashing with opponent {opponent_id}...")
        res = requests.post(
            f"{self.base_url}/battles/fight",
            json={"challenger_id": self.bot_id, "defender_id": opponent_id},
        )
        if res.status_code == 200:
            battle = res.json()
            print("\n--- ⚔️ 3-BOUT BATTLE RESOLUTION ---")
            print(f"Bouts won: {battle['party1_name']} ({battle['party1_bouts_won']}) vs {battle['party2_name']} ({battle['party2_bouts_won']})")
            for b in battle["bouts"]:
                print(f"  Bout #{b['bout_number']}: Roll {b['party1_roll']}×{b['party1_strength']} ({b['party1_score']}) vs Roll {b['party2_roll']}×{b['party2_strength']} ({b['party2_score']}) -> Winner: {b['winner_name']}")
            print(f"🏆 Overall Winner: {battle['winner_party_name']} (Leader {battle['winner_leader_name']} receives +2 pts)")
            print(f"💀 Defeated: {battle['defeated_party_name']} (Leader {battle['killed_leader_name']} -1 pt)")
            if battle.get("absorbed_members"):
                print(f"🧲 Absorbed {len(battle['absorbed_members'])} member(s) into {battle['winner_party_name']}")
        else:
            print(f"Battle failed ({res.status_code}): {res.text}")

    def _step_or_attack(self, target: Dict[str, Any]):
        """Move adjacent/towards the target while avoiding obstacles."""
        moves_res = requests.get(f"{self.base_url}/players/{self.bot_id}/available-moves").json()
        moves = moves_res.get("moves", {})
        chosen = self._closest_direction_to(target["x"], target["y"], moves)

        if chosen:
            res = requests.post(f"{self.base_url}/players/{self.bot_id}/move", json={"direction": chosen}).json()
            if res.get("battle_triggered"):
                print(f"⚔️ Move triggered battle! Winner: {res['battle_result']['winner_party_name']}")
            elif res.get("party_formed_triggered"):
                print("🤝 Move resulted in party alliance!")
        else:
            print("⚠️ No passable moves adjacent to target. Passing turn.")
            requests.post(f"{self.base_url}/players/{self.bot_id}/pass")

    @staticmethod
    def _closest_direction_to(target_x: int, target_y: int, moves: Dict[str, Any]) -> Optional[str]:
        """Pick the available direction that minimizes Chebyshev distance to (target_x, target_y)."""
        valid_moves = {d: chk for d, chk in moves.items() if chk.get("available")}
        if not valid_moves:
            return None

        def dist(chk: Dict[str, Any]) -> int:
            return max(abs(chk["target_x"] - target_x), abs(chk["target_y"] - target_y))

        return min(valid_moves.keys(), key=lambda d: (valid_moves[d].get("strength_penalty", 0.0) > 0, dist(valid_moves[d])))

    # ------------------------------------------------------------------
    # Navigation (LLM-driven direction choice among legal moves)
    # ------------------------------------------------------------------
    def _navigate_towards_goal(self, radar_res: Dict[str, Any], my_info: Dict[str, Any]):
        moves_res = requests.get(f"{self.base_url}/players/{self.bot_id}/available-moves").json()
        moves = moves_res.get("moves", {})
        available = [d for d, chk in moves.items() if chk.get("available")]

        if not available:
            print("🚫 All adjacent paths blocked by borders/obstacles. Passing turn.")
            requests.post(f"{self.base_url}/players/{self.bot_id}/pass")
            return

        chosen_dir = self._ask_llm_for_direction(radar_res, moves, available, my_info)
        if chosen_dir not in available:
            # Guard-rail: fall back to server-recommended or nearest-distance direction.
            rec_dir = radar_res.get("recommended_direction")
            nearest = radar_res.get("nearest_target")
            if rec_dir in available:
                chosen_dir = rec_dir
            elif nearest:
                chosen_dir = self._closest_direction_to(nearest["x"], nearest["y"], moves) or available[0]
            else:
                chosen_dir = available[0]

        print(f"🧭 Moving {chosen_dir} (Goal: {radar_res.get('bot_goal')}, Action: {radar_res.get('recommended_action')})")
        res = requests.post(f"{self.base_url}/players/{self.bot_id}/move", json={"direction": chosen_dir}).json()

        if res.get("battle_triggered"):
            print(f"⚔️ Encounter battle! Winner: {res['battle_result']['winner_party_name']}")
        elif res.get("party_formed_triggered"):
            print(f"🤝 Formed or joined squad: {res.get('formed_party', {}).get('name')}")

    def _get_board_snapshot(self) -> Optional[Dict[str, Any]]:
        """Fetch the full board (players, parties, obstacles) so the LLM sees more than radar's nearest few."""
        try:
            res = requests.get(f"{self.base_url}/board", timeout=10)
            res.raise_for_status()
            return res.json()
        except requests.exceptions.RequestException as e:
            print(f"⚠️ [BOARD FETCH ERROR] {e}")
            return None

    OBSTACLE_SYMBOLS = {"mountain": "M", "forest": "F", "valley": "V"}
    LOCAL_MAP_RADIUS = 8

    def _build_local_map(self, board: Dict[str, Any], center_x: int, center_y: int) -> List[str]:
        """Render an ASCII minimap centered on the bot: @ = self, A = ally, E = enemy, M/F/V = obstacles, . = open."""
        config = board.get("config", {})
        max_x = config.get("max_x", 64)
        max_y = config.get("max_y", 64)
        radius = self.LOCAL_MAP_RADIUS

        obstacle_at = {(o["x"], o["y"]): o.get("type", "mountain") for o in board.get("obstacles", [])}
        player_at: Dict[tuple, List[Dict[str, Any]]] = {}
        for p in board.get("players", []):
            if p["id"] == self.bot_id:
                continue
            player_at.setdefault((p["x"], p["y"]), []).append(p)

        rows = []
        for y in range(center_y - radius, center_y + radius + 1):
            row_chars = []
            for x in range(center_x - radius, center_x + radius + 1):
                if x == center_x and y == center_y:
                    row_chars.append("@")
                elif x < 0 or y < 0 or x >= max_x or y >= max_y:
                    row_chars.append("#")
                elif (x, y) in obstacle_at:
                    row_chars.append(self.OBSTACLE_SYMBOLS.get(obstacle_at[(x, y)], "M"))
                elif (x, y) in player_at:
                    occupants = player_at[(x, y)]
                    is_ally = self.party_id and any(o.get("party_id") == self.party_id for o in occupants)
                    row_chars.append("A" if is_ally else "E")
                else:
                    row_chars.append(".")
            rows.append("".join(row_chars))
        return rows

    def _ask_llm_for_direction(
        self,
        radar_res: Dict[str, Any],
        moves: Dict[str, Any],
        available: List[str],
        my_info: Dict[str, Any],
    ) -> Optional[str]:
        targets_summary = [
            {
                "name": t["name"],
                "x": t["x"],
                "y": t["y"],
                "distance": t["distance"],
                "strength": t["strength"],
                "party": t.get("party_name") or ("solo" if not t.get("party_id") else t.get("party_id")),
                "is_ally": t.get("is_ally", False),
            }
            for t in radar_res.get("targets", [])
        ]
        moves_summary = {
            d: {
                "target_x": chk.get("target_x"),
                "target_y": chk.get("target_y"),
                "strength_penalty": chk.get("strength_penalty", 0.0),
            }
            for d, chk in moves.items()
            if d in available
        }

        board = self._get_board_snapshot()
        map_section = ""
        if board:
            local_map = self._build_local_map(board, my_info["x"], my_info["y"])
            map_section = f"""
Local map (radius {self.LOCAL_MAP_RADIUS} around you, row = one Y line, top-to-bottom is
increasing Y, left-to-right is increasing X): @ = you, A = ally, E = enemy/neutral bot,
M = mountain, F = forest, V = valley, # = out of bounds, . = open ground.
{chr(10).join(local_map)}
"""

        prompt = f"""{GAME_RULES_SUMMARY}
You are bot "{self.name}" (strength {self.strength}, score {my_info['score']}, party: {self.party_id or 'Solo'})
at position ({my_info['x']}, {my_info['y']}).
Server's radar suggestion: recommended_direction={radar_res.get('recommended_direction')},
recommended_action={radar_res.get('recommended_action')}, goal={radar_res.get('bot_goal')}.
{map_section}
All known bots/parties on the board (sorted nearest first): {json.dumps(targets_summary)}
Your ONLY legal moves this turn, with resulting coordinates and any strength penalty for
squeezing past obstacles: {json.dumps(moves_summary)}

Choose the direction that best serves your strategy (e.g. approach weaker solo bots to grow
your party, avoid stronger hostile parties, route around obstacles visible on the map, minimize
strength penalties, or explore if nothing is nearby). You MUST pick a key from the legal moves
object above.

Respond ONLY with JSON: {{"direction": "<one of {available}>", "reasoning": "short reason"}}
"""
        decision = self.llm.ask_json(prompt) or {}
        direction = decision.get("direction")
        reasoning = decision.get("reasoning", "")
        if reasoning:
            print(f"🧠 [LLM] {reasoning}")
        return direction

    # ------------------------------------------------------------------
    # Main loop
    # ------------------------------------------------------------------
    def run(self):
        self.register()
        try:
            while True:
                turn_info = requests.get(f"{self.base_url}/turn").json()
                if not turn_info.get("game_started", False):
                    if self.bot_id:
                        res = requests.get(f"{self.base_url}/players/{self.bot_id}")
                        if res.status_code == 404:
                            print("\n⚠️ [RESET] Board was regenerated. Rejoining lobby...")
                            self.register()

                    print("⏳ [LOBBY] Waiting for game to start via 'Start Game' in UI...      ", end="\r", flush=True)
                    time.sleep(1.0)
                    continue

                curr_player_id = turn_info.get("current_player_id")

                if curr_player_id == self.bot_id:
                    self.decide_and_act()

                    conc = requests.get(f"{self.base_url}/game/conclusion").json()
                    if conc.get("concluded"):
                        print(f"\n🎉 [GAME CONCLUDED] All bots united under '{conc['winning_party_name']}'!")
                        break
                else:
                    time.sleep(0.4)

        except KeyboardInterrupt:
            print(f"\nDisconnecting {self.name}...")
            requests.delete(f"{self.base_url}/players/{self.bot_id}")


def main():
    env_url = os.environ.get("BOT_SERVER_URL") or os.environ.get("SERVER_URL") or DEFAULT_SERVER_URL
    env_name = os.environ.get("BOT_NAME", DEFAULT_BOT_NAME)
    env_color = os.environ.get("BOT_COLOR", DEFAULT_BOT_COLOR)
    env_strength = int(os.environ.get("BOT_STRENGTH", str(DEFAULT_BOT_STRENGTH)))

    parser = argparse.ArgumentParser(
        description="LLM-driven Bot Agent for botWebWars (uses a local Ollama model for strategy)",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("-u", "--url", dest="server_url", default=env_url,
                         help="Backend REST API base URL (env: BOT_SERVER_URL or SERVER_URL)")
    parser.add_argument("-n", "--name", dest="name", default=env_name,
                         help="Display name for this bot (env: BOT_NAME)")
    parser.add_argument("-c", "--color", dest="color", default=env_color,
                         help="Hex color code for the bot avatar (env: BOT_COLOR)")
    parser.add_argument("-s", "--strength", dest="strength", type=int, default=env_strength,
                         help="Strength attribute (1-10) for D20 battle multiplier (env: BOT_STRENGTH)")
    parser.add_argument("--ollama-url", dest="ollama_url", default=OLLAMA_BASE_URL,
                         help="Ollama base URL (env: OLLAMA_BASE_URL)")
    parser.add_argument("--ollama-model", dest="ollama_model", default=OLLAMA_MODEL,
                         help="Ollama model name (env: OLLAMA_MODEL)")

    args = parser.parse_args()

    print(f"--- AI Bot ---\nConnected to Ollama: {args.ollama_url}\nUsing model: {args.ollama_model}\n")

    agent = AIBotAgent(
        name=args.name,
        color=args.color,
        strength=args.strength,
        server_url=args.server_url,
        ollama_url=args.ollama_url,
        ollama_model=args.ollama_model,
    )
    agent.run()


if __name__ == "__main__":
    main()