let hexSize = 40;
let hexGrid = [];
let buildings = [];
let train = { x: 0, y: 0, dir: 1 };
let monsterMode = false;
let currentMonster = 'Godzilla';
let monsterX, monsterY;
let angle = Math.PI / 6;
let debugText = "Loading...";
let playerId = null;
const sqrt3 = Math.sqrt(3);

// Streamlit connection
const { Streamlit } = window;
Streamlit.events.addEventListener("message", (event) => {
  const data = event.data;
  hexGrid = data.game_state.hex_grid;
  buildings = data.game_state.buildings;
  train = data.game_state.train;
  monsterMode = data.game_state.monster_mode;
  currentMonster = data.game_state.current_monster;
  monsterX = data.game_state.monster_x;
  monsterY = data.game_state.monster_y;
  playerId = data.player_id;
  selectedEmoji = data.selected_emoji;
  debugText = `Player ${playerId} joined. Click to build!`;
  // Set train track on row 5 if not already set
  for (let i = 0; i < hexGrid.length; i++) {
    if (hexGrid[i][5].type === 'empty') hexGrid[i][5].type = 'track';
  }
});

function setup() {
  createCanvas(800, 600);
  Streamlit.setFrameHeight(650);
}

function draw() {
  background(220);
  drawHexGrid();
  drawBuildings();
  drawTrain();
  if (monsterMode) drawMonster();
  fill(0);
  textSize(12);
  text(debugText, 10, 20);
}

function drawHexGrid() {
  for (let i = 0; i < hexGrid.length; i++) {
    for (let j = 0; j < hexGrid[i].length; j++) {
      let x = i * hexSize * 1.5;
      let y = j * hexSize * sqrt3 + (i % 2 === 1 ? hexSize * sqrt3 / 2 : 0);
      let z = j * hexSize * Math.sin(angle);
      fill(hexGrid[i][j].type === 'track' ? 100 : 150, 200, 150);
      stroke(0);
      beginShape();
      for (let k = 0; k < 6; k++) {
        let angleRad = Math.PI / 3 * k;
        vertex(x + hexSize * Math.cos(angleRad), y + hexSize * Math.sin(angleRad) - z);
      }
      endShape(CLOSE);
      if (hexGrid[i][j].emoji) {
        textSize(20);
        text(hexGrid[i][j].emoji, x - 10, y + 5 - z);
      }
    }
  }
}

function drawBuildings() {
  buildings.forEach(b => {
    let x = b.x * hexSize * 1.5;
    let y = b.y * hexSize * sqrt3 + (b.x % 2 === 1 ? hexSize * sqrt3 / 2 : 0);
    let z = b.y * hexSize * Math.sin(angle);
    fill(b.color);
    noStroke();
    beginShape();
    if (b.type === 'Residential') {
      vertex(x + 10, y + 30 - z); vertex(x + 20, y + 10 - z); vertex(x + 30, y + 30 - z);
      vertex(x + 25, y + 30 - z); vertex(x + 25, y + 35 - z); vertex(x + 15, y + 35 - z);
      vertex(x + 15, y + 30 - z); vertex(x + 17, y + 25 - z); vertex(x + 23, y + 25 - z);
    } else if (b.type === 'Commercial') {
      vertex(x + 10, y + 35 - z); vertex(x + 15, y + 15 - z); vertex(x + 25, y + 15 - z);
      vertex(x + 30, y + 35 - z); vertex(x + 27, y + 35 - z); vertex(x + 27, y + 20 - z);
      vertex(x + 13, y + 20 - z); vertex(x + 13, y + 35 - z);
    } else if (b.type === 'Industrial') {
      vertex(x + 5, y + 35 - z); vertex(x + 15, y + 20 - z); vertex(x + 25, y + 20 - z);
      vertex(x + 35, y + 35 - z); vertex(x + 30, y + 35 - z); vertex(x + 30, y + 15 - z);
      vertex(x + 33, y + 15 - z); vertex(x + 33, y + 35 - z);
    } else if (b.type === 'School') {
      vertex(x + 5, y + 35 - z); vertex(x + 10, y + 20 - z); vertex(x + 30, y + 20 - z);
      vertex(x + 35, y + 35 - z); vertex(x + 25, y + 35 - z); vertex(x + 25, y + 10 - z);
      vertex(x + 27, y + 10 - z); vertex(x + 27, y + 35 - z);
    } else if (b.type === 'PowerPlant') {
      vertex(x + 5, y + 35 - z); vertex(x + 15, y + 15 - z); vertex(x + 25, y + 15 - z);
      vertex(x + 35, y + 35 - z); vertex(x + 30, y + 35 - z); vertex(x + 30, y + 25 - z);
      vertex(x + 20, y + 25 - z); vertex(x + 20, y + 35 - z);
    }
    endShape(CLOSE);
  });
}

function drawTrain() {
  let tx = train.x;
  let ty = train.y;
  let tz = train.y * Math.sin(angle);
  fill(150, 50, 50);
  noStroke();
  beginShape();
  vertex(tx + 10, ty + 10 - tz); vertex(tx + 30, ty + 10 - tz); vertex(tx + 35, ty + 20 - tz);
  vertex(tx + 30, ty + 30 - tz); vertex(tx + 10, ty + 30 - tz); vertex(tx + 5, ty + 20 - tz);
  vertex(tx + 15, ty + 20 - tz); vertex(tx + 15, ty + 15 - tz); vertex(tx + 25, ty + 15 - tz);
  vertex(tx + 25, ty + 20 - tz);
  endShape(CLOSE);
  train.x += train.dir * 2;
  if (train.x > width || train.x < 0) train.dir *= -1;
  saveState();
}

function drawMonster() {
  let mz = monsterY * Math.sin(angle);
  fill(255, 0, 0);
  noStroke();
  beginShape();
  if (currentMonster === 'Godzilla') {
    vertex(monsterX, monsterY + 40 - mz); vertex(monsterX + 10, monsterY + 20 - mz);
    vertex(monsterX + 20, monsterY - mz); vertex(monsterX + 30, monsterY + 20 - mz);
    vertex(monsterX + 40, monsterY + 40 - mz); vertex(monsterX + 35, monsterY + 50 - mz);
    vertex(monsterX + 25, monsterY + 60 - mz); vertex(monsterX + 15, monsterY + 50 - mz);
    vertex(monsterX + 20, monsterY + 40 - mz); vertex(monsterX + 25, monsterY + 30 - mz);
  } else if (currentMonster === 'GiantRobot') {
    vertex(monsterX, monsterY + 40 - mz); vertex(monsterX + 10, monsterY + 20 - mz);
    vertex(monsterX + 15, monsterY - mz); vertex(monsterX + 25, monsterY - mz);
    vertex(monsterX + 30, monsterY + 20 - mz); vertex(monsterX + 40, monsterY + 40 - mz);
    vertex(monsterX + 35, monsterY + 50 - mz); vertex(monsterX + 20, monsterY + 60 - mz);
    vertex(monsterX + 5, monsterY + 50 - mz); vertex(monsterX + 15, monsterY + 30 - mz);
  }
  endShape(CLOSE);
  monsterX += random(-5, 5);
  monsterY += random(-5, 5);
  monsterX = constrain(monsterX, 0, width);
  monsterY = constrain(monsterY, 0, height);
  saveState();
}

function mousePressed() {
  let i = Math.floor(mouseX / (hexSize * 1.5));
  let j = Math.floor((mouseY - (i % 2 === 1 ? hexSize * sqrt3 / 2 : 0)) / (hexSize * sqrt3));
  if (i >= 0 && i < hexGrid.length && j >= 0 && j < hexGrid[0].length) {
    if (hexGrid[i][j].type === 'empty' && !monsterMode) {
      let selectedEmoji = window.Streamlit.initialArgs.selected_emoji;
      if (['🏠', '🏡', '🏢', '🏣', '🏤', '🏥', '🏦', '🏨', '🏩', '🏪'].includes(selectedEmoji)) {
        let types = ['Residential', 'Commercial', 'Industrial', 'School', 'PowerPlant'];
        let colors = [[0, 200, 0], [0, 0, 200], [200, 200, 0], [200, 0, 200], [100, 100, 100]];
        let idx = Math.floor(Math.random() * 5);
        buildings.push({ x: i, y: j, type: types[idx], color: colors[idx], player: playerId });
        hexGrid[i][j].type = 'building';
      } else {
        hexGrid[i][j].type = 'placed';
        hexGrid[i][j].emoji = selectedEmoji;
      }
      debugText = `Player ${playerId} placed ${selectedEmoji} at (${i}, ${j})`;
      saveState();
    }
  }
}

function keyPressed() {
  if (key === 'm' || key === 'M') {
    monsterMode = !monsterMode;
    debugText = `Monster Mode: ${monsterMode}`;
    saveState();
  }
  if (key === 'g' || key === 'G') {
    currentMonster = 'Godzilla';
    debugText = "Monster set to Godzilla";
    saveState();
  }
  if (key === 'r' || key === 'R') {
    currentMonster = 'GiantRobot';
    debugText = "Monster set to Giant Robot";
    saveState();
  }
  if (key === 't' || key === 'T') {
    train.dir *= -1;
    debugText = "Train direction reversed";
    saveState();
  }
}

function saveState() {
  const state = {
    hex_grid: hexGrid,
    buildings: buildings,
    train: train,
    monster_mode: monsterMode,
    current_monster: currentMonster,
    monster_x: monsterX,
    monster_y: monsterY
  };
  Streamlit.setComponentValue(state);
}

function toggleMonster() {
  monsterMode = !monsterMode;
  debugText = `Monster Mode: ${monsterMode}`;
  saveState();
}

function setMonster(monster) {
  currentMonster = monster;
  debugText = `Monster set to ${monster}`;
  saveState();
}

function reverseTrain() {
  train.dir *= -1;
  debugText = "Train direction reversed";
  saveState();
}