import streamlit as st
import numpy as np
import pandas as pd
import time
import random
from PIL import Image, ImageDraw
import math
import uuid
import base64
from io import BytesIO

# Set page config
st.set_page_config(
    page_title="Battle Simulator",
    page_icon="🎮",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Constants
MAP_WIDTH = 800
MAP_HEIGHT = 600
TEAM_RED = "red"
TEAM_BLUE = "blue"
UNIT_TYPES = ["infantry", "tank", "artillery"]
TEAM_COLORS = {
    TEAM_RED: "#e74c3c",
    TEAM_BLUE: "#3498db"
}
UNIT_STATS = {
    "infantry": {
        "max_health": 80,
        "attack": 8,
        "defense": 3,
        "range": 40,
        "speed": 1.5,
        "size": 10,
        "shape": "circle",
        "respawn_time": 120
    },
    "tank": {
        "max_health": 150, 
        "attack": 15,
        "defense": 10,
        "range": 60,
        "speed": 0.8,
        "size": 12,
        "shape": "square",
        "respawn_time": 240
    },
    "artillery": {
        "max_health": 60,
        "attack": 25,
        "defense": 2,
        "range": 120,
        "speed": 0.5,
        "size": 13,
        "shape": "triangle",
        "respawn_time": 180
    }
}

# Utility functions
def calculate_distance(x1, y1, x2, y2):
    """Calculate Euclidean distance between two points"""
    return math.sqrt((x2 - x1)**2 + (y2 - y1)**2)

def angle_between(x1, y1, x2, y2):
    """Calculate angle between two points in radians"""
    return math.atan2(y2 - y1, x2 - x1)

class Unit:
    def __init__(self, id, team, unit_type, x, y):
        self.id = id
        self.team = team
        self.type = unit_type
        self.x = x
        self.y = y
        self.target_x = x
        self.target_y = y
        
        # Get stats from the unit type
        stats = UNIT_STATS[unit_type]
        self.max_health = stats["max_health"]
        self.health = self.max_health
        self.attack = stats["attack"]
        self.defense = stats["defense"]
        self.range = stats["range"]
        self.speed = stats["speed"]
        self.size = stats["size"]
        self.shape = stats["shape"]
        self.respawn_time = stats["respawn_time"]
        
        # Battle state
        self.state = "idle"  # idle, moving, attacking, dead
        self.target = None
        self.dead_timer = 0
        self.attack_cooldown = 0
        self.attack_cooldown_max = 60
        self.attack_line = None

    def find_target(self, units):
        """Find nearest enemy unit"""
        if self.state == "dead":
            return

        nearest_dist = float("inf")
        nearest_enemy = None

        for unit in units:
            if unit.team != self.team and unit.state != "dead":
                dist = calculate_distance(self.x, self.y, unit.x, unit.y)
                if dist < nearest_dist:
                    nearest_dist = dist
                    nearest_enemy = unit

        self.target = nearest_enemy

    def update(self, units):
        """Update unit state and position"""
        if self.state == "dead":
            self.dead_timer += 1
            if self.dead_timer >= self.respawn_time:
                self.respawn()
            return

        # Find target if no current target or target is dead
        if not self.target or self.target.state == "dead":
            self.find_target(units)
            self.state = "idle"

        # If we have a target, pursue it
        if self.target:
            dist_to_target = calculate_distance(self.x, self.y, self.target.x, self.target.y)
            
            # If in attack range, attack
            if dist_to_target <= self.range:
                self.attack_target()
                self.state = "attacking"
            else:
                # Move toward target
                self.target_x = self.target.x
                self.target_y = self.target.y
                self.state = "moving"
        else:
            # Wander if no target
            if random.random() < 0.02:
                team_zone = 300 if self.team == TEAM_RED else 500
                center_x = 200 if self.team == TEAM_RED else 600
                
                self.target_x = max(50, min(MAP_WIDTH - 50, 
                                     center_x + (random.random() - 0.5) * team_zone))
                self.target_y = max(50, min(MAP_HEIGHT - 50, 
                                     300 + (random.random() - 0.5) * 400))

        # Move toward target position
        self.move()
        
        # Update attack cooldown
        if self.attack_cooldown > 0:
            self.attack_cooldown -= 1
            
        # Update attack line
        if self.attack_line:
            self.attack_line["duration"] -= 1
            if self.attack_line["duration"] <= 0:
                self.attack_line = None

    def move(self):
        """Move unit toward target position"""
        if self.state == "dead":
            return

        dx = self.target_x - self.x
        dy = self.target_y - self.y
        dist = math.sqrt(dx*dx + dy*dy)
        
        if dist > 5:
            self.x += (dx / dist) * self.speed
            self.y += (dy / dist) * self.speed
            
        # Keep units within map bounds
        self.x = max(self.size, min(MAP_WIDTH - self.size, self.x))
        self.y = max(self.size, min(MAP_HEIGHT - self.size, self.y))

    def attack_target(self):
        """Attack the current target"""
        if self.attack_cooldown == 0 and self.target and self.target.state != "dead":
            # Apply random variance to attack damage
            damage_multiplier = 0.8 + random.random() * 0.4
            base_damage = self.attack * damage_multiplier
            final_damage = max(1, base_damage - self.target.defense / 2)
            
            # Apply damage to target
            self.target.take_damage(final_damage)
            
            # Create attack line for visualization
            self.attack_line = {
                "from_x": self.x,
                "from_y": self.y,
                "to_x": self.target.x,
                "to_y": self.target.y,
                "duration": 10
            }
            
            # Reset cooldown
            self.attack_cooldown = self.attack_cooldown_max

    def take_damage(self, amount):
        """Take damage from an attack"""
        self.health -= amount
        if self.health <= 0:
            self.die()

    def die(self):
        """Unit death"""
        self.state = "dead"
        self.health = 0
        self.dead_timer = 0
        self.attack_line = None

    def respawn(self):
        """Respawn the unit"""
        margin = 50
        if self.team == TEAM_RED:
            self.x = margin + random.random() * 100
        else:
            self.x = MAP_WIDTH - margin - random.random() * 100
        
        self.y = 100 + random.random() * 400
        self.health = self.max_health
        self.state = "idle"
        self.target = None
        self.attack_line = None
        self.dead_timer = 0
        self.attack_cooldown = 0

    def is_point_inside(self, point_x, point_y):
        """Check if a point is inside the unit"""
        if self.state == "dead":
            return False
            
        return calculate_distance(self.x, self.y, point_x, point_y) <= self.size

class BattleSimulation:
    def __init__(self):
        self.units = []
        self.selected_unit = None
        self.frame_count = 0
        
        # Create initial units
        self.create_units()
        
    def create_units(self):
        """Create initial units for both teams"""
        # Create red team (left side)
        for i in range(5):
            self.units.append(Unit(
                id=f"red-infantry-{i}",
                team=TEAM_RED,
                unit_type="infantry",
                x=50 + random.random() * 100,
                y=100 + i * 100
            ))
            
        for i in range(3):
            self.units.append(Unit(
                id=f"red-tank-{i}",
                team=TEAM_RED,
                unit_type="tank",
                x=100 + random.random() * 100,
                y=150 + i * 150
            ))
            
        for i in range(2):
            self.units.append(Unit(
                id=f"red-artillery-{i}",
                team=TEAM_RED,
                unit_type="artillery",
                x=50 + random.random() * 80,
                y=200 + i * 200
            ))
        
        # Create blue team (right side)
        for i in range(5):
            self.units.append(Unit(
                id=f"blue-infantry-{i}",
                team=TEAM_BLUE,
                unit_type="infantry",
                x=650 + random.random() * 100,
                y=100 + i * 100
            ))
            
        for i in range(3):
            self.units.append(Unit(
                id=f"blue-tank-{i}",
                team=TEAM_BLUE,
                unit_type="tank",
                x=600 + random.random() * 100,
                y=150 + i * 150
            ))
            
        for i in range(2):
            self.units.append(Unit(
                id=f"blue-artillery-{i}",
                team=TEAM_BLUE,
                unit_type="artillery",
                x=670 + random.random() * 80,
                y=200 + i * 200
            ))
    
    def update(self):
        """Update all units in the simulation"""
        for unit in self.units:
            unit.update(self.units)
        
        self.frame_count += 1
        
    def render(self):
        """Render the current state of the battle"""
        # Create a new image
        img = Image.new('RGBA', (MAP_WIDTH, MAP_HEIGHT), (232, 244, 229, 255))
        draw = ImageDraw.Draw(img)
        
        # Draw territory gradients
        # (This is approximate as PIL doesn't have easy gradients)
        # Red territory
        for x in range(250):
            alpha = int(50 * (1 - x / 250))
            draw.line([(x, 0), (x, MAP_HEIGHT)], fill=(231, 76, 60, alpha), width=1)
            
        # Blue territory
        for x in range(MAP_WIDTH - 250, MAP_WIDTH):
            alpha = int(50 * (x - (MAP_WIDTH - 250)) / 250)
            draw.line([(x, 0), (x, MAP_HEIGHT)], fill=(52, 152, 219, alpha), width=1)
        
        # Draw attack lines first (so they appear behind units)
        for unit in self.units:
            if unit.attack_line:
                line_color = TEAM_COLORS[unit.team]
                draw.line(
                    [(unit.attack_line["from_x"], unit.attack_line["from_y"]), 
                     (unit.attack_line["to_x"], unit.attack_line["to_y"])],
                    fill=line_color,
                    width=2
                )
        
        # Draw units
        for unit in self.units:
            if unit.state == "dead":
                continue
                
            # Draw the unit shape
            fill_color = TEAM_COLORS[unit.team]
            
            if unit.shape == "circle":  # Infantry
                draw.ellipse(
                    [(unit.x - unit.size, unit.y - unit.size), 
                     (unit.x + unit.size, unit.y + unit.size)],
                    fill=fill_color
                )
            elif unit.shape == "square":  # Tank
                draw.rectangle(
                    [(unit.x - unit.size, unit.y - unit.size), 
                     (unit.x + unit.size, unit.y + unit.size)],
                    fill=fill_color
                )
            elif unit.shape == "triangle":  # Artillery
                draw.polygon(
                    [(unit.x, unit.y - unit.size),
                     (unit.x - unit.size, unit.y + unit.size),
                     (unit.x + unit.size, unit.y + unit.size)],
                    fill=fill_color
                )
            
            # Draw health bar
            health_ratio = unit.health / unit.max_health
            health_bar_width = 20
            health_bar_height = 3
            
            # Health bar background (red)
            draw.rectangle(
                [(unit.x - health_bar_width/2, unit.y - 20),
                 (unit.x + health_bar_width/2, unit.y - 20 + health_bar_height)],
                fill="#e74c3c"
            )
            
            # Health bar foreground (green)
            draw.rectangle(
                [(unit.x - health_bar_width/2, unit.y - 20),
                 (unit.x - health_bar_width/2 + health_bar_width * health_ratio, unit.y - 20 + health_bar_height)],
                fill="#2ecc71"
            )
            
            # Draw state indicator
            state_colors = {
                "idle": "#95a5a6",     # Gray
                "moving": "#3498db",    # Blue
                "attacking": "#e74c3c"  # Red
            }
            
            draw.ellipse(
                [(unit.x - 3, unit.y - 25 - 3),
                 (unit.x + 3, unit.y - 25 + 3)],
                fill=state_colors.get(unit.state, "#95a5a6")
            )
            
            # Highlight selected unit
            if self.selected_unit and unit.id == self.selected_unit.id:
                draw.ellipse(
                    [(unit.x - unit.size - 5, unit.y - unit.size - 5), 
                     (unit.x + unit.size + 5, unit.y + unit.size + 5)],
                    outline="#f1c40f",
                    width=2
                )
        
        # Convert the image to base64 for displaying in Streamlit
        buffered = BytesIO()
        img.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        
        return img_str
        
    def handle_click(self, x, y):
        """Handle click on the battle map"""
        self.selected_unit = None
        for unit in self.units:
            if unit.is_point_inside(x, y) and unit.state != "dead":
                self.selected_unit = unit
                break
        
        return self.selected_unit

def get_unit_info_html(unit):
    """Generate HTML for unit info panel"""
    if not unit:
        return ""
        
    type_icon = {
        "infantry": "⭕",
        "tank": "⬛",
        "artillery": "△"
    }.get(unit.type, "")
    
    team_color = TEAM_COLORS[unit.team]
    
    html = f"""
    <div style="background-color: rgba(255, 255, 255, 0.95); padding: 20px; 
                border-radius: 8px; border: 2px solid {team_color}; box-shadow: 0 4px 15px rgba(0, 0, 0, 0.3);">
        <h3 style="margin-bottom: 15px; color: #333; border-bottom: 1px solid #ddd; padding-bottom: 10px;">
            {unit.team.capitalize()} {unit.type.capitalize()} {type_icon}
        </h3>
        <p style="margin-bottom: 10px; display: flex; justify-content: space-between;">
            <span>Health:</span> <span>{int(unit.health)}/{unit.max_health}</span>
        </p>
        <p style="margin-bottom: 10px; display: flex; justify-content: space-between;">
            <span>Attack:</span> <span>{unit.attack}</span>
        </p>
        <p style="margin-bottom: 10px; display: flex; justify-content: space-between;">
            <span>Defense:</span> <span>{unit.defense}</span>
        </p>
        <p style="margin-bottom: 10px; display: flex; justify-content: space-between;">
            <span>Range:</span> <span>{unit.range}</span>
        </p>
        <p style="margin-bottom: 10px; display: flex; justify-content: space-between;">
            <span>Speed:</span> <span>{unit.speed}</span>
        </p>
        <p style="margin-bottom: 10px; display: flex; justify-content: space-between;">
            <span>Status:</span> <span>{unit.state.capitalize()}</span>
        </p>
    </div>
    """
    return html

def main():
    # Initialize session state
    if 'battle_sim' not in st.session_state:
        st.session_state.battle_sim = BattleSimulation()
        st.session_state.last_click = None
    
    # Title and description
    st.title("Battle Simulator")
    st.markdown("An interactive battle simulation where red and blue forces fight on a dynamic battlefield.")
    
    # Create a layout with columns
    col1, col2 = st.columns([3, 1])
    
    with col1:
        # Display battle map with click handler
        battle_sim = st.session_state.battle_sim
        battle_sim.update()  # Update simulation
        img_str = battle_sim.render()  # Render the current state
        
        # Display the battle map with click handling
        st.markdown(f"""
        <div style="position: relative; width: {MAP_WIDTH}px; margin: 0 auto;">
            <img src="data:image/png;base64,{img_str}" width="{MAP_WIDTH}" id="battlemap" 
                 style="border: 1px solid #ccc; cursor: pointer;" onclick="handleMapClick(event)">
        </div>
        <script>
            function handleMapClick(e) {
                const rect = e.target.getBoundingClientRect();
                const x = e.clientX - rect.left;
                const y = e.clientY - rect.top;
                // Pass click coordinates to Streamlit
                window.parent.postMessage({{
                    type: "streamlit:setComponentValue",
                    value: [x, y]
                }}, "*");
            }
        </script>
        """, unsafe_allow_html=True)
        
        # Use a streamlit empty element to trigger rerun on map click
        map_click = st.empty()
        
    with col2:
        # Show information about selected unit
        st.markdown("### Unit Information")
        
        unit_info = st.empty()
        if battle_sim.selected_unit:
            unit_info.markdown(get_unit_info_html(battle_sim.selected_unit), unsafe_allow_html=True)
        else:
            unit_info.markdown("Click on a unit to see its details")
            
        # Unit legend
        st.markdown("### Unit Types")
        st.markdown("""
        - ⭕ **Infantry**: Fast but weaker units
        - ⬛ **Tank**: Slow but powerful with strong defense
        - △ **Artillery**: Long-range attacks but vulnerable
        """)
        
        # Battle statistics
        st.markdown("### Battle Statistics")
        
        # Count units by team and status
        red_active = sum(1 for unit in battle_sim.units if unit.team == TEAM_RED and unit.state != "dead")
        blue_active = sum(1 for unit in battle_sim.units if unit.team == TEAM_BLUE and unit.state != "dead")
        
        # Display counters
        st.markdown(f"""
        <div style="display: flex; justify-content: space-between; margin-bottom: 10px;">
            <div style="text-align: center; padding: 10px; background-color: rgba(231, 76, 60, 0.2); border-radius: 5px; width: 45%;">
                <div style="font-size: 24px; font-weight: bold;">{red_active}</div>
                <div>Red Forces</div>
            </div>
            <div style="text-align: center; padding: 10px; background-color: rgba(52, 152, 219, 0.2); border-radius: 5px; width: 45%;">
                <div style="font-size: 24px; font-weight: bold;">{blue_active}</div>
                <div>Blue Forces</div>
            </div>
        </div>
        """, unsafe_allow_html=True)
    
    # Handle map clicks
    clicked = st.button('Refresh Battle State', key='refresh_battle')
    
    # Create a container for receiving click data and auto-refreshing
    click_container = st.container()
    with click_container:
        click_data = st.text_input("Click coordinates", 
                                  value="", 
                                  label_visibility="collapsed", 
                                  key="click_data")
        
        if click_data:
            try:
                x, y = map(float, click_data.strip('[]').split(','))
                selected = battle_sim.handle_click(x, y)
                st.session_state.last_click = (x, y)
                st.experimental_rerun()
            except:
                pass
    
    # Auto-refresh the page
    st.markdown("""
    <script>
        // Set up to receive messages from the iframe
        window.addEventListener('message', function(e) {
            if (e.data.type === 'streamlit:setComponentValue') {
                const coordinates = e.data.value;
                document.querySelector('input[key="click_data"]').value = coordinates;
                document.querySelector('input[key="click_data"]').dispatchEvent(new Event('input'));
            }
        });
        
        // Auto-refresh the simulation every 1 second
        const refreshInterval = setInterval(function() {
            const refreshButton = document.querySelector('button[key="refresh_battle"]');
            if (refreshButton) {
                refreshButton.click();
            }
        }, 1000);
    </script>
    """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()