import streamlit as st
from streamlit.components.v1 import html
import random
import json
import math

# Initialize game state
if 'game_state' not in st.session_state:
    st.session_state.game_state = {
        'hex_grid': [[{'type': 'empty', 'emoji': ''} for _ in range(12)] for _ in range(16)],  # Smaller 16x12 grid for better visibility
        'buildings': [],
        'players': {},
        'train': {'x': 0, 'y': 5, 'dir': 1},  # Using grid coordinates
        'monster_mode': False,
        'current_monster': 'Godzilla',
        'monster_x': 8,
        'monster_y': 6,
        'last_update': None
    }

# Constants and assets
HEX_SIZE = 40
random_names = ["SkyWalker", "ForestRanger", "CityBuilder", "MonsterTamer", "RailMaster"]
plants = ["🌱", "🌲", "🌳", "🌴", "🌵"]
buildings = ["🏠", "🏡", "🏢", "🏥", "🏦"]
creatures = ["🐾", "🐱", "🐶", "🐭", "🐰"]

# Updated p5.js code
p5js_code = """
const HEX_SIZE = 40;
const SQRT_3 = Math.sqrt(3);
let hexGrid = [];
let buildings = [];
let train = {};
let monsterMode = false;
let currentMonster = '';
let monsterX, monsterY;
let playerId;

function setup() {
  createCanvas(640, 480);  // Adjusted canvas size for 16x12 grid
  updateFromState();
}

function updateFromState() {
  hexGrid = JSON.parse(document.getElementById('game_state').innerHTML);
  buildings = JSON.parse(document.getElementById('buildings').innerHTML);
  train = JSON.parse(document.getElementById('train').innerHTML);
  monsterMode = JSON.parse(document.getElementById('monster_mode').innerHTML);
  currentMonster = document.getElementById('current_monster').innerHTML;
  monsterX = parseInt(document.getElementById('monster_x').innerHTML);
  monsterY = parseInt(document.getElementById('monster_y').innerHTML);
  playerId = document.getElementById('player_id').innerHTML;
}

function draw() {
  background(220);
  drawHexGrid();
  drawBuildings();
  drawTrain();
  if (monsterMode) drawMonster();
  
  // Train movement
  train.x += train.dir * 0.05;
  if (train.x >= hexGrid.length || train.x < 0) train.dir *= -1;
}

function drawHexGrid() {
  for (let i = 0; i < hexGrid.length; i++) {
    for (let j = 0; j < hexGrid[i].length; j++) {
      let x = i * HEX_SIZE * 1.5;
      let y = j * HEX_SIZE * SQRT_3 + (i % 2 ? HEX_SIZE * SQRT_3 / 2 : 0);
      fill(hexGrid[i][j].type === 'track' ? '#808080' : '#90EE90');
      stroke(0);
      drawHex(x, y);
      if (hexGrid[i][j].emoji) {
        textSize(20);
        text(hexGrid[i][j].emoji, x - 10, y + 5);
      }
    }
  }
}

function drawHex(x, y) {
  beginShape();
  for (let a = 0; a < 6; a++) {
    let angle = TWO_PI / 6 * a;
    vertex(x + HEX_SIZE * cos(angle), y + HEX_SIZE * sin(angle));
  }
  endShape(CLOSE);
}

function drawBuildings() {
  buildings.forEach(b => {
    let x = b.x * HEX_SIZE * 1.5;
    let y = b.y * HEX_SIZE * SQRT_3 + (b.x % 2 ? HEX_SIZE * SQRT_3 / 2 : 0);
    fill(b.color[0], b.color[1], b.color[2]);
    noStroke();
    ellipse(x, y, HEX_SIZE * 0.8);
    textSize(20);
    text(b.emoji, x - 10, y + 5);
  });
}

function drawTrain() {
  let x = train.x * HEX_SIZE * 1.5;
  let y = train.y * HEX_SIZE * SQRT_3;
  fill(150, 50, 50);
  rect(x - 15, y - 10, 30, 20);
  text('🚂', x - 10, y + 5);
}

function drawMonster() {
  let x = monsterX * HEX_SIZE * 1.5;
  let y = monsterY * HEX_SIZE * SQRT_3 + (monsterX % 2 ? HEX_SIZE * SQRT_3 / 2 : 0);
  fill(255, 0, 0);
  ellipse(x, y, HEX_SIZE * 1.2);
  text(currentMonster === 'Godzilla' ? '🦖' : '🤖', x - 10, y + 5);
}

function mousePressed() {
  let i = Math.floor(mouseX / (HEX_SIZE * 1.5));
  let j = Math.floor((mouseY - (i % 2 ? HEX_SIZE * SQRT_3 / 2 : 0)) / (HEX_SIZE * SQRT_3));
  if (i >= 0 && i < hexGrid.length && j >= 0 && j < hexGrid[0].length && hexGrid[i][j].type === 'empty') {
    let emoji = document.getElementById('selected_emoji').innerHTML;
    if (!monsterMode) {
      if (emoji.startsWith('🏠') || emoji.startsWith('🏡') || emoji.startsWith('🏢') || 
          emoji.startsWith('🏥') || emoji.startsWith('🏦')) {
        buildings.push({
          x: i, 
          y: j, 
          emoji: emoji,
          color: [random(100, 255), random(100, 255), random(100, 255)],
          player: playerId
        });
        hexGrid[i][j].type = 'building';
      } else {
        hexGrid[i][j].type = 'placed';
        hexGrid[i][j].emoji = emoji;
      }
      updateState();
    }
  }
}

function keyPressed() {
  if (key === 'm' || key === 'M') monsterMode = !monsterMode;
  if (key === 'g' || key === 'G') currentMonster = 'Godzilla';
  if (key === 'r' || key === 'R') currentMonster = 'GiantRobot';
  if (key === 't' || key === 'T') train.dir *= -1;
  updateState();
}

function updateState() {
  fetch('/update_state', {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({
      hex_grid: hexGrid,
      buildings: buildings,
      train: train,
      monster_mode: monsterMode,
      current_monster: currentMonster,
      monster_x: monsterX,
      monster_y: monsterY,
      player_id: playerId
    })
  });
}
"""

# UI Components
if 'player_id' not in st.session_state:
    st.session_state.player_id = random.choice(random_names)

player_id = st.sidebar.selectbox("Choose Player", random_names, index=random_names.index(st.session_state.player_id))
st.session_state.player_id = player_id

# Emoji palette
st.sidebar.subheader("🎨 Palette")
cols = st.sidebar.columns(5)
selected_emoji = st.session_state.get('selected_emoji', plants[0])
for i, emoji in enumerate(plants + buildings + creatures):
    if cols[i % 5].button(emoji, key=f"emoji_{i}"):
        st.session_state.selected_emoji = emoji
        selected_emoji = emoji

# Scoreboard
if player_id not in st.session_state.game_state['players']:
    st.session_state.game_state['players'][player_id] = 0
st.sidebar.subheader("🏆 Scores")
for p, s in st.session_state.game_state['players'].items():
    st.sidebar.write(f"{p}: {s}")

# Game HTML
game_html = f"""
<script src="https://cdnjs.cloudflare.com/ajax/libs/p5.js/1.4.2/p5.min.js"></script>
<div id="game_state" style="display:none">{json.dumps(st.session_state.game_state['hex_grid'])}</div>
<div id="buildings" style="display:none">{json.dumps(st.session_state.game_state['buildings'])}</div>
<div id="train" style="display:none">{json.dumps(st.session_state.game_state['train'])}</div>
<div id="monster_mode" style="display:none">{json.dumps(st.session_state.game_state['monster_mode'])}</div>
<div id="current_monster" style="display:none">{st.session_state.game_state['current_monster']}</div>
<div id="monster_x" style="display:none">{st.session_state.game_state['monster_x']}</div>
<div id="monster_y" style="display:none">{st.session_state.game_state['monster_y']}</div>
<div id="player_id" style="display:none">{player_id}</div>
<div id="selected_emoji" style="display:none">{selected_emoji}</div>
<script>{p5js_code}</script>
"""

# Main layout
st.title("HexCity Adventure")
st.write(f"Player: {player_id}. Click to place {selected_emoji}. Keys: M (monster), G/R (monster type), T (train)")
html(game_html, height=500)

# State update handler
def update_state(data):
    st.session_state.game_state.update({
        'hex_grid': data['hex_grid'],
        'buildings': data['buildings'],
        'train': data['train'],
        'monster_mode': data['monster_mode'],
        'current_monster': data['current_monster'],
        'monster_x': data['monster_x'],
        'monster_y': data['monster_y'],
        'last_update': data
    })
    if data['player_id'] not in st.session_state.game_state['players']:
        st.session_state.game_state['players'][data['player_id']] = 0
    st.session_state.game_state['players'][data['player_id']] += 1
    st.rerun()

# Initialize tracks
for i in range(len(st.session_state.game_state['hex_grid'])):
    st.session_state.game_state['hex_grid'][i][5]['type'] = 'track'