# LLMArena_no_outlines.py
import io
import os
import re
import time
import math
import random
from collections import defaultdict
from datetime import datetime

import numpy as np
import pandas as pd
import requests
from tqdm.auto import tqdm
from PIL import Image, ImageDraw, ImageFont
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import gradio as gr
import gistyc

# Try optional transformers
try:
    from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
    TRANSFORMERS_AVAILABLE = True
except Exception:
    TRANSFORMERS_AVAILABLE = False

# Constants
ARENA_SIZE = 10
ROBOT_SIZE = 1.5
MAX_HEALTH = 100
MAX_ENERGY = 50
MAX_ROUNDS = 20
TURN_TIME_LIMIT = 30

# Environment variables
GITHUB_TOKEN = os.environ.get("GITHUB_TOKEN", "")
RESULTS_GIST_ID = os.environ.get("RESULTS_GIST_ID", "battle_results_gist")
LEADERBOARD_GIST_ID = os.environ.get("LEADERBOARD_GIST_ID", "battle_leaderboard_gist")

# --- Robot and Arena classes (identical logic, trimmed) ---
class Robot:
    def __init__(self, name, model_id, gen_mode="heuristic"):
        self.name = name
        self.model_id = model_id
        self.health = MAX_HEALTH
        self.energy = MAX_ENERGY
        self.position = [0, 0]
        self.facing = 0  # degrees
        self.actions = []
        self.gen_mode = gen_mode
        self._generator = None

    def load_model(self):
        # For transformer mode, create a generation pipeline cached on robot
        if self.gen_mode == "transformers":
            if not TRANSFORMERS_AVAILABLE:
                raise RuntimeError("transformers not available in environment")
            if self._generator is None:
                # Attempt to build a text-generation pipeline from the model_id
                try:
                    # Use model_id as HF model identifier
                    tok = AutoTokenizer.from_pretrained(self.model_id)
                    model = AutoModelForCausalLM.from_pretrained(self.model_id)
                    self._generator = pipeline("text-generation", model=model, tokenizer=tok, device=-1)
                except Exception as e:
                    raise RuntimeError(f"Failed to load HF model {self.model_id}: {e}")
        return self._generator

class BattleArena:
    def __init__(self):
        self.robot1 = None
        self.robot2 = None
        self.turn = 1
        self.round = 1
        self.game_over = False
        self.winner = None
        self.history = []

    def initialize_battle(self, model_id1, model_id2, mode1="heuristic", mode2="heuristic"):
        self.robot1 = Robot("Robot A", model_id1, gen_mode=mode1)
        self.robot2 = Robot("Robot B", model_id2, gen_mode=mode2)
        self.robot1.position = [-3, 0]
        self.robot2.position = [3, 0]
        self.robot1.facing = 0
        self.robot2.facing = 180
        self.turn = 1
        self.round = 1
        self.game_over = False
        self.winner = None
        self.history = []

    def get_legal_actions(self, robot):
        actions = [
            "MOVE_FORWARD", "MOVE_BACKWARD", "TURN_LEFT", "TURN_RIGHT",
            "PUNCH", "KICK", "ENERGY_BLAST", "BLOCK", "CHARGE"
        ]
        legal = []
        for action in actions:
            if action in ["PUNCH", "KICK"] and robot.energy < 10:
                continue
            if action == "ENERGY_BLAST" and robot.energy < 25:
                continue
            if action == "CHARGE" and robot.energy >= MAX_ENERGY:
                continue
            if "MOVE" in action:
                new_pos = self.calculate_new_position(robot, action)
                if not (-ARENA_SIZE/2 <= new_pos[0] <= ARENA_SIZE/2 and -ARENA_SIZE/2 <= new_pos[1] <= ARENA_SIZE/2):
                    continue
            legal.append(action)
        return legal

    def calculate_new_position(self, robot, action):
        x, y = robot.position
        angle = math.radians(robot.facing)
        if action == "MOVE_FORWARD":
            x += math.cos(angle) * 1.5
            y += math.sin(angle) * 1.5
        elif action == "MOVE_BACKWARD":
            x -= math.cos(angle) * 1.0
            y -= math.sin(angle) * 1.0
        return [x, y]

    def calculate_distance(self):
        x1, y1 = self.robot1.position
        x2, y2 = self.robot2.position
        return math.sqrt((x2-x1)**2 + (y2-y1)**2)

    def is_facing_opponent(self, attacker, target):
        dx = target.position[0] - attacker.position[0]
        dy = target.position[1] - attacker.position[1]
        target_angle = math.degrees(math.atan2(dy, dx)) % 360
        angle_diff = abs((attacker.facing - target_angle + 180) % 360 - 180)
        return angle_diff <= 45

    def execute_action(self, attacker, target, action):
        damage = 0
        energy_cost = 0
        message = ""
        distance = self.calculate_distance()

        if action == "MOVE_FORWARD":
            attacker.position = self.calculate_new_position(attacker, action)
            energy_cost = 5
            message = f"{attacker.name} moves forward"
        elif action == "MOVE_BACKWARD":
            attacker.position = self.calculate_new_position(attacker, action)
            energy_cost = 3
            message = f"{attacker.name} moves backward"
        elif action == "TURN_LEFT":
            attacker.facing = (attacker.facing - 45) % 360
            energy_cost = 2
            message = f"{attacker.name} turns left"
        elif action == "TURN_RIGHT":
            attacker.facing = (attacker.facing + 45) % 360
            energy_cost = 2
            message = f"{attacker.name} turns right"
        elif action == "PUNCH":
            if distance <= 2.5 and self.is_facing_opponent(attacker, target):
                damage = 15
                message = f"{attacker.name} lands a solid punch on {target.name}!"
            else:
                message = f"{attacker.name} tries to punch but misses!"
            energy_cost = 10
        elif action == "KICK":
            if distance <= 3.0 and self.is_facing_opponent(attacker, target):
                damage = 25
                message = f"{attacker.name} delivers a powerful kick to {target.name}!"
            else:
                message = f"{attacker.name}'s kick misses the target!"
            energy_cost = 15
        elif action == "ENERGY_BLAST":
            if distance <= 6.0 and self.is_facing_opponent(attacker, target):
                damage = 35
                message = f"{attacker.name} hits {target.name} with an energy blast! 💥"
            else:
                message = f"{attacker.name}'s energy blast misses!"
            energy_cost = 25
        elif action == "BLOCK":
            energy_cost = 5
            message = f"{attacker.name} assumes defensive stance"
        elif action == "CHARGE":
            energy_gain = 15
            attacker.energy = min(MAX_ENERGY, attacker.energy + energy_gain)
            message = f"{attacker.name} charges up energy +{energy_gain}"

        if damage > 0:
            target.health = max(0, target.health - damage)
        attacker.energy = max(0, attacker.energy - energy_cost)
        attacker.actions.append(action)

        return message, damage

    def check_game_over(self):
        if self.robot1.health <= 0:
            self.game_over = True
            self.winner = self.robot2
            return True
        elif self.robot2.health <= 0:
            self.game_over = True
            self.winner = self.robot1
            return True
        return False

# --- Prompt generator (for transformer mode) ---
def generate_action_prompt(arena, current_robot, opponent):
    legal_actions = arena.get_legal_actions(current_robot)
    actions_str = "|".join(legal_actions)
    prompt = (
        f"BATTLE ARENA - ROUND {arena.round}\n\n"
        f"You are {current_robot.name} controlling a battle robot.\n"
        f"Your opponent is {opponent.name}.\n\n"
        "CURRENT STATUS:\n"
        f"- Your Health: {current_robot.health}/{MAX_HEALTH}\n"
        f"- Your Energy: {current_robot.energy}/{MAX_ENERGY}\n"
        f"- Your Position: ({current_robot.position[0]:.1f}, {current_robot.position[1]:.1f})\n"
        f"- Your Facing: {current_robot.facing}°\n"
        f"- Opponent Health: {opponent.health}/{MAX_HEALTH}\n"
        f"- Opponent Position: ({opponent.position[0]:.1f}, {opponent.position[1]:.1f})\n"
        f"- Distance to opponent: {arena.calculate_distance():.1f}\n\n"
        f"LEGAL ACTIONS: {actions_str}\n\n"
        "Choose ONE action from the legal moves above. Respond with only the action name.\n\n"
        "ACTION:"
    )
    return prompt, legal_actions

# --- Validation ---
def validate_and_sanitize_action(generated_text, legal_actions):
    cleaned = (generated_text or "").strip().upper()
    # exact match
    if cleaned in legal_actions:
        return cleaned
    # partial match heuristics
    for action in legal_actions:
        if action in cleaned:
            return action
    # match keywords
    keywords = {
        "PUNCH": ["PUNCH", "HIT", "STRIKE"],
        "KICK": ["KICK", "KICKED"],
        "ENERGY_BLAST": ["ENERGY", "BLAST", "BEAM"],
        "BLOCK": ["BLOCK", "DEFEND", "GUARD"],
        "CHARGE": ["CHARGE", "RECHARGE", "ENERGIZE"],
        "MOVE_FORWARD": ["FORWARD", "ADVANCE", "MOVE_FORWARD"],
        "MOVE_BACKWARD": ["BACK", "RETREAT", "MOVE_BACKWARD"],
        "TURN_LEFT": ["LEFT", "TURN_LEFT"],
        "TURN_RIGHT": ["RIGHT", "TURN_RIGHT"]
    }
    for act, kws in keywords.items():
        if act in legal_actions:
            for kw in kws:
                if kw in cleaned:
                    return act
    # fallback: prefer CHARGE if available, else random legal
    if "CHARGE" in legal_actions:
        return "CHARGE"
    if legal_actions:
        return random.choice(legal_actions)
    return "CHARGE"

# --- Heuristic policy (fallback) ---
def heuristic_policy(arena, robot, opponent):
    legal = arena.get_legal_actions(robot)
    dist = arena.calculate_distance()

    # If low on energy, charge
    if robot.energy <= 10 and "CHARGE" in legal:
        return "CHARGE"
    # If close and facing, prefer PUNCH or KICK
    if dist <= 2.5:
        if "PUNCH" in legal:
            return "PUNCH"
        if "KICK" in legal:
            return "KICK"
    # If mid-range and have energy, ENERGY_BLAST
    if 2.5 < dist <= 6.0 and robot.energy >= 25 and "ENERGY_BLAST" in legal:
        return "ENERGY_BLAST"
    # If opponent behind, turn
    if not arena.is_facing_opponent(robot, opponent):
        # choose turn that minimizes angle difference
        # simple random left/right
        return random.choice([a for a in legal if a.startswith("TURN")] or legal)
    # Otherwise move towards opponent if not too close
    if dist > 3.0 and "MOVE_FORWARD" in legal:
        return "MOVE_FORWARD"
    # Default: random legal
    return random.choice(legal) if legal else "CHARGE"

# --- Transformers policy ---
def transformers_policy(arena, robot, opponent, prompt, legal_actions):
    """
    Use HF pipeline to generate text. We then extract action name by regex/validation.
    """
    try:
        gen = robot.load_model()
    except Exception as e:
        # If model fails to load, fallback to heuristic
        print(f"Transformer load error: {e}")
        return heuristic_policy(arena, robot, opponent)

    # generate short completion
    try:
        # some models respond better with small max_new_tokens
        out = gen(prompt, max_new_tokens=16, do_sample=True, temperature=0.7, top_k=50, num_return_sequences=1)
        text = out[0]["generated_text"] if isinstance(out, list) else str(out)
        # keep only the appended text after the prompt (some pipelines return full text)
        if text.startswith(prompt):
            text = text[len(prompt):]
        # extract first token-like word
        candidate = text.strip().splitlines()[0].strip()
        # sanitize to action
        return validate_and_sanitize_action(candidate, legal_actions)
    except Exception as e:
        print(f"Transformers generation failed: {e}")
        return heuristic_policy(arena, robot, opponent)

# --- ELO and persistence (same logic) ---
def calculate_elo(rank1, rank2, result):
    K = 32
    expected_score1 = 1 / (1 + 10 ** ((rank2 - rank1) / 400))
    new_rank1 = rank1 + K * (result - expected_score1)
    return round(new_rank1)

def update_elo_ratings(battle_data):
    elo_ratings = defaultdict(lambda: 1000)
    for index, row in battle_data.iterrows():
        if row["Result"] == "DRAW":
            continue
        model1 = row["Model1"]
        model2 = row["Model2"]
        result = row["Result"]
        model1_elo = elo_ratings[model1]
        model2_elo = elo_ratings[model2]
        if result == "WIN1":
            elo_ratings[model1] = calculate_elo(model1_elo, model2_elo, 1)
            elo_ratings[model2] = calculate_elo(model2_elo, model1_elo, 0)
        elif result == "WIN2":
            elo_ratings[model1] = calculate_elo(model1_elo, model2_elo, 0)
            elo_ratings[model2] = calculate_elo(model2_elo, model1_elo, 1)
    return elo_ratings

def save_battle_result(model_id1, model_id2, winner, termination, rounds):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    if winner == "Robot A":
        result = "WIN1"
    elif winner == "Robot B":
        result = "WIN2"
    else:
        result = "DRAW"
    data_str = f"{timestamp},{model_id1},{model_id2},{result},{termination},{rounds}\n"
    with open("battle_results.csv", "a") as file:
        file.write(data_str)
    if GITHUB_TOKEN:
        try:
            gist_api = gistyc.GISTyc(auth_token=GITHUB_TOKEN)
            response = gist_api.update_gist(file_name="battle_results.csv", gist_id=RESULTS_GIST_ID)
            print("Results Gist updated")
        except Exception as e:
            print(f"Failed to update results Gist: {e}")

def update_leaderboard():
    try:
        battle_data = pd.read_csv("battle_results.csv")
        model_stats = defaultdict(lambda: {"Wins": 0, "Losses": 0, "Draws": 0})
        for _, row in battle_data.iterrows():
            if row["Result"] == "WIN1":
                model_stats[row["Model1"]]["Wins"] += 1
                model_stats[row["Model2"]]["Losses"] += 1
            elif row["Result"] == "WIN2":
                model_stats[row["Model2"]]["Wins"] += 1
                model_stats[row["Model1"]]["Losses"] += 1
            else:
                model_stats[row["Model1"]]["Draws"] += 1
                model_stats[row["Model2"]]["Draws"] += 1
        elo_ratings = update_elo_ratings(battle_data)
        leaderboard_data = []
        for model, stats in model_stats.items():
            total_games = stats["Wins"] + stats["Losses"] + stats["Draws"]
            win_rate = (stats["Wins"] / total_games * 100) if total_games > 0 else 0
            leaderboard_data.append({
                "Model": model,
                "Wins": stats["Wins"],
                "Losses": stats["Losses"],
                "Draws": stats["Draws"],
                "Win Rate": f"{win_rate:.1f}%",
                "ELO Rating": elo_ratings[model]
            })
        # If no models have played yet, return an empty dataframe with the correct columns
        if not leaderboard_data:
            cols = ["Model", "Wins", "Losses", "Draws", "Win Rate", "ELO Rating"]
            leaderboard_df = pd.DataFrame(columns=cols)
        else:
            leaderboard_df = pd.DataFrame(leaderboard_data)
        leaderboard_df.sort_values("ELO Rating", ascending=False, inplace=True)
        leaderboard_df.reset_index(drop=True, inplace=True)
        leaderboard_df.to_csv("battle_leaderboard.csv", index=False)
        if GITHUB_TOKEN:
            gist_api = gistyc.GISTyc(auth_token=GITHUB_TOKEN)
            response = gist_api.update_gist(file_name="battle_leaderboard.csv", gist_id=LEADERBOARD_GIST_ID)
            print("Leaderboard Gist updated")
        return leaderboard_df
    except Exception as e:
        print(f"Leaderboard update failed: {e}")
        # On error return an empty leaderboard (no example/demo rows)
        cols = ["Model", "Wins", "Losses", "Draws", "Win Rate", "ELO Rating"]
        return pd.DataFrame(columns=cols)

def get_leaderboard():
    try:
        return pd.read_csv("battle_leaderboard.csv")
    except:
        # If leaderboard file doesn't exist yet, return an empty dataframe with the expected columns
        cols = ["Model", "Wins", "Losses", "Draws", "Win Rate", "ELO Rating"]
        return pd.DataFrame(columns=cols)

# --- Rendering ---
def create_robot_mesh(x, y, color, name, facing):
    angle = math.radians(facing)
    base_vertices = [
        [x-0.4, y-0.3, 0], [x+0.4, y-0.3, 0],
        [x+0.4, y+0.3, 0], [x-0.4, y+0.3, 0],
        [x-0.4, y-0.3, 1.2], [x+0.4, y-0.3, 1.2],
        [x+0.4, y+0.3, 1.2], [x-0.4, y+0.3, 1.2],
    ]
    head_vertices = [
        [x-0.2, y-0.2, 1.2], [x+0.2, y-0.2, 1.2],
        [x+0.2, y+0.2, 1.2], [x-0.2, y+0.2, 1.2],
        [x-0.2, y-0.2, 1.8], [x+0.2, y-0.2, 1.8],
        [x+0.2, y+0.2, 1.8], [x-0.2, y+0.2, 1.8],
    ]
    vertices = base_vertices + head_vertices
    faces = [
        [0,1,2], [0,2,3], [4,5,6], [4,6,7],
        [0,4,7], [0,7,3], [1,5,6], [1,6,2],
        [8,9,10], [8,10,11], [12,13,14], [12,14,15]
    ]
    return go.Mesh3d(
        x=[v[0] for v in vertices],
        y=[v[1] for v in vertices],
        z=[v[2] for v in vertices],
        i=[f[0] for f in faces],
        j=[f[1] for f in faces],
        k=[f[2] for f in faces],
        opacity=0.9,
        name=name
    )

def render_arena_3d(arena):
    fig = make_subplots(rows=1, cols=1, specs=[[{'type':'scatter3d'}]])
    x_floor, y_floor = np.meshgrid(np.linspace(-5,5,10), np.linspace(-5,5,10))
    z_floor = np.zeros(x_floor.shape)
    fig.add_trace(go.Surface(x=x_floor, y=y_floor, z=z_floor, colorscale='Greys', opacity=0.7, showscale=False, name='Arena'))
    boundary_x = [-5,-5,5,5,-5]
    boundary_y = [-5,5,5,-5,-5]
    boundary_z = [0,0,0,0,0]
    fig.add_trace(go.Scatter3d(x=boundary_x, y=boundary_y, z=boundary_z, mode='lines', line=dict(color='red', width=4), name='Boundary'))
    fig.add_trace(create_robot_mesh(arena.robot1.position[0], arena.robot1.position[1], 'red', arena.robot1.name, arena.robot1.facing))
    fig.add_trace(create_robot_mesh(arena.robot2.position[0], arena.robot2.position[1], 'blue', arena.robot2.name, arena.robot2.facing))
    r1_angle = math.radians(arena.robot1.facing)
    r2_angle = math.radians(arena.robot2.facing)
    fig.add_trace(go.Scatter3d(x=[arena.robot1.position[0], arena.robot1.position[0] + math.cos(r1_angle)],
                               y=[arena.robot1.position[1], arena.robot1.position[1] + math.sin(r1_angle)],
                               z=[1.0,1.0], mode='lines', line=dict(color='darkred', width=6), showlegend=False))
    fig.add_trace(go.Scatter3d(x=[arena.robot2.position[0], arena.robot2.position[0] + math.cos(r2_angle)],
                               y=[arena.robot2.position[1], arena.robot2.position[1] + math.sin(r2_angle)],
                               z=[1.0,1.0], mode='lines', line=dict(color='darkblue', width=6), showlegend=False))
    fig.update_layout(scene=dict(xaxis=dict(range=[-6,6], showbackground=False, showticklabels=False, title=''),
                                 yaxis=dict(range=[-6,6], showbackground=False, showticklabels=False, title=''),
                                 zaxis=dict(range=[0,6], showbackground=False, showticklabels=False, title=''),
                                 aspectmode='cube', camera=dict(eye=dict(x=1.5, y=1.5, z=1.2))),
                      margin=dict(l=0,r=0,t=0,b=0), height=500, showlegend=True)
    return fig

def create_status_image(arena):
    img = Image.new('RGB', (800, 300), color='black')
    draw = ImageDraw.Draw(img)
    try:
        font_large = ImageFont.truetype("arial.ttf", 20)
        font_medium = ImageFont.truetype("arial.ttf", 16)
        font_small = ImageFont.truetype("arial.ttf", 14)
    except:
        font_large = font_medium = font_small = ImageFont.load_default()
    draw.rectangle([5,5,395,145], outline='red', width=2)
    draw.text((15,15), f"🤖 {arena.robot1.name}", fill='red', font=font_large)
    hp1 = arena.robot1.health / MAX_HEALTH
    bw1 = int(350 * hp1)
    draw.rectangle([15,45,15 + bw1,65], fill='red')
    draw.rectangle([15,45,365,65], outline='white', width=1)
    draw.text((15,70), f"Health: {arena.robot1.health}/{MAX_HEALTH}", fill='white', font=font_medium)
    e1 = arena.robot1.energy / MAX_ENERGY
    be1 = int(350 * e1)
    draw.rectangle([15,85,15 + be1,105], fill='yellow')
    draw.rectangle([15,85,365,105], outline='white', width=1)
    draw.text((15,110), f"Energy: {arena.robot1.energy}/{MAX_ENERGY}", fill='white', font=font_medium)
    draw.text((15,130), f"Position: ({arena.robot1.position[0]:.1f}, {arena.robot1.position[1]:.1f})", fill='white', font=font_small)
    draw.text((200,130), f"Facing: {arena.robot1.facing}°", fill='white', font=font_small)
    draw.rectangle([405,5,795,145], outline='blue', width=2)
    draw.text((415,15), f"🤖 {arena.robot2.name}", fill='blue', font=font_large)
    hp2 = arena.robot2.health / MAX_HEALTH
    bw2 = int(350 * hp2)
    draw.rectangle([415,45,415 + bw2,65], fill='blue')
    draw.rectangle([415,45,765,65], outline='white', width=1)
    draw.text((415,70), f"Health: {arena.robot2.health}/{MAX_HEALTH}", fill='white', font=font_medium)
    e2 = arena.robot2.energy / MAX_ENERGY
    be2 = int(350 * e2)
    draw.rectangle([415,85,415 + be2,105], fill='yellow')
    draw.rectangle([415,85,765,105], outline='white', width=1)
    draw.text((415,110), f"Energy: {arena.robot2.energy}/{MAX_ENERGY}", fill='white', font=font_medium)
    draw.text((415,130), f"Position: ({arena.robot2.position[0]:.1f}, {arena.robot2.position[1]:.1f})", fill='white', font=font_small)
    draw.text((600,130), f"Facing: {arena.robot2.facing}°", fill='white', font=font_small)
    draw.rectangle([5,155,795,295], outline='green', width=2)
    draw.text((15,165), f"⚔️ BATTLE ARENA - ROUND {arena.round}", fill='yellow', font=font_large)
    draw.text((15,195), f"Distance between robots: {arena.calculate_distance():.1f}", fill='white', font=font_medium)
    draw.text((15,220), f"Current turn: {arena.robot1.name if arena.turn == 1 else arena.robot2.name}", fill='white', font=font_medium)
    if arena.game_over and arena.winner:
        draw.text((15,245), f"🏆 WINNER: {arena.winner.name}! 🏆", fill='green', font=font_large)
    elif arena.game_over:
        draw.text((15,245), "DRAW - Maximum rounds reached", fill='yellow', font=font_large)
    return img

# --- Battle sequence generator (used by Gradio) ---
def battle_sequence(model_id1, model_id2, mode_choice="auto"):
    """
    mode_choice: "auto", "heuristic", "transformers"
    - "auto": if model_id looks like HF then try transformers; else heuristic
    """
    # decide modes
    def decide_mode(mid):
        if mode_choice == "heuristic":
            return "heuristic"
        if mode_choice == "transformers":
            return "transformers"
        # auto
        if TRANSFORMERS_AVAILABLE and "/" in mid:  # naive HF id detection
            return "transformers"
        return "heuristic"

    m1_mode = decide_mode(model_id1)
    m2_mode = decide_mode(model_id2)

    arena = BattleArena()
    arena.initialize_battle(model_id1, model_id2, mode1=m1_mode, mode2=m2_mode)

    battle_log = []
    fig = render_arena_3d(arena)
    status_img = create_status_image(arena)
    yield fig, status_img, "🤖 Battle starting! Initializing..."

    # Load transformer models if needed
    try:
        if arena.robot1.gen_mode == "transformers":
            arena.robot1.load_model()
        if arena.robot2.gen_mode == "transformers":
            arena.robot2.load_model()
        battle_log.append("✅ Models initialized (or heuristics ready).")
    except Exception as e:
        yield fig, status_img, f"❌ Model init error: {e}"
        return

    battle_log.append("⚔️ Let the battle begin!")

    while not arena.game_over and arena.round <= MAX_ROUNDS:
        current_robot = arena.robot1 if arena.turn == 1 else arena.robot2
        opponent = arena.robot2 if arena.turn == 1 else arena.robot1
        prompt, legal_actions = generate_action_prompt(arena, current_robot, opponent)
        # legal actions as list
        legal_list = legal_actions

        # choose action based on mode
        if current_robot.gen_mode == "transformers":
            action = transformers_policy(arena, current_robot, opponent, prompt, legal_list)
        else:
            action = heuristic_policy(arena, current_robot, opponent)

        # final sanity
        action = validate_and_sanitize_action(action, legal_list)

        message, damage = arena.execute_action(current_robot, opponent, action)
        log_entry = f"Round {arena.round} - {current_robot.name}: {action}"
        if message:
            log_entry += f" - {message}"
        if damage > 0:
            log_entry += f" 💥(-{damage} HP)"
        battle_log.append(log_entry)

        if arena.check_game_over():
            break

        if arena.turn == 2:
            arena.round += 1
        arena.turn = 3 - arena.turn

        fig = render_arena_3d(arena)
        status_img = create_status_image(arena)
        # yield last part of log
        yield fig, status_img, "\n".join(battle_log[-8:])
        time.sleep(0.6)

    if arena.game_over:
        conclusion = f"🏆 BATTLE OVER! {arena.winner.name} WINS! 🏆"
        battle_log.append(conclusion)
        save_battle_result(model_id1, model_id2, arena.winner.name, "KO", arena.round)
        update_leaderboard()
    else:
        conclusion = "🤝 Battle ended in draw - maximum rounds reached"
        battle_log.append(conclusion)
        save_battle_result(model_id1, model_id2, "DRAW", "Max Rounds", arena.round)
        update_leaderboard()

    fig = render_arena_3d(arena)
    status_img = create_status_image(arena)
    yield fig, status_img, "\n".join(battle_log[-12:])

# --- File initialization ---
def initialize_files():
    if not os.path.exists("battle_results.csv"):
        with open("battle_results.csv", "w") as f:
            f.write("Timestamp,Model1,Model2,Result,Termination,Rounds\n")
    if not os.path.exists("battle_leaderboard.csv"):
        df = get_leaderboard()
        df.to_csv("battle_leaderboard.csv", index=False)

initialize_files()

# --- Gradio UI ---
title = """
<div align="center">
  <p style="font-size: 32px;">🤖 LLM Battle Arena — No outlines</p>
  <p style="font-size: 16px;">Scegli due "modelli" o usa la modalità heuristica. Se vuoi usare un modello HuggingFace, inserisci l'ID (es. 'gpt2' o 'facebook/opt-125m') e assicurati di avere 'transformers' installato.</p>
</div>
"""

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(title)
    with gr.Row():
        with gr.Column():
            model_id1 = gr.Textbox(label="🤖 Robot A Model ID", value="heuristic", placeholder="Inserisci HF model id o 'heuristic'")
            model_id2 = gr.Textbox(label="🤖 Robot B Model ID", value="heuristic", placeholder="Inserisci HF model id o 'heuristic'")
            mode_choice = gr.Radio(choices=["auto", "heuristic", "transformers"], value="auto", label="Mode (auto sceglie in base all'ID e disponibilità transformers)")
            battle_btn = gr.Button("🎯 Start Battle!", variant="primary", size="lg")
        with gr.Column():
            arena_plot = gr.Plot(label="🎪 3D Battle Arena")
            status_display = gr.Image(label="📊 Battle Status", height=300)
    battle_log = gr.Textbox(label="📝 Battle Log", lines=6, max_lines=10)
    with gr.Row():
        gr.Markdown("### 🏆 Leaderboard")
    leaderboard = gr.Dataframe(value=get_leaderboard, every=60, label="Model Rankings")
    footer = """
    <div align="center">
    <p><em>Azioni: MOVE_FORWARD, MOVE_BACKWARD, TURN_LEFT, TURN_RIGHT, PUNCH, KICK, ENERGY_BLAST, BLOCK, CHARGE</em></p>
    </div>
    """
    gr.Markdown(footer)

    battle_btn.click(fn=battle_sequence, inputs=[model_id1, model_id2, mode_choice], outputs=[arena_plot, status_display, battle_log])

if __name__ == "__main__":
    demo.launch(share=True)