import streamlit as st
import numpy as np
import time
from dataclasses import dataclass
from typing import Dict, List, Optional

@dataclass
class Building:
    name: str
    cost: Dict[str, int]
    produces: str
    color: str
    icon: str
    description: str

class FactoryGame:
    def __init__(self):
        self.GRID_SIZE = 8
        self.buildings = {
            "solar_panel": Building(
                name="Solar Panel",
                cost={"minerals": 10},
                produces="energy",
                color="#FEF08A",
                icon="☀️",
                description="Produces energy from sunlight"
            ),
            "mineral_extractor": Building(
                name="Mineral Extractor",
                cost={"energy": 10},
                produces="minerals",
                color="#D1D5DB",
                icon="⛏️",
                description="Extracts minerals from the ground"
            ),
            "circuit_factory": Building(
                name="Circuit Factory",
                cost={"energy": 15, "minerals": 15},
                produces="circuits",
                color="#BBF7D0",
                icon="🔧",
                description="Produces electronic circuits"
            )
        }

        # Initialize session state if needed
        if "resources" not in st.session_state:
            st.session_state.resources = {
                "energy": 50,
                "minerals": 50,
                "circuits": 0
            }
        if "grid" not in st.session_state:
            st.session_state.grid = [[None for _ in range(self.GRID_SIZE)] 
                                   for _ in range(self.GRID_SIZE)]
        if "last_update" not in st.session_state:
            st.session_state.last_update = time.time()
        if "selected_building" not in st.session_state:
            st.session_state.selected_building = None

    def update_resources(self):
        current_time = time.time()
        elapsed = current_time - st.session_state.last_update
        
        if elapsed >= 1.0:  # Update every second
            for row in range(self.GRID_SIZE):
                for col in range(self.GRID_SIZE):
                    building_id = st.session_state.grid[row][col]
                    if building_id:
                        building = self.buildings[building_id]
                        st.session_state.resources[building.produces] += 1
            st.session_state.last_update = current_time

    def can_afford(self, building_id: str) -> bool:
        building = self.buildings[building_id]
        for resource, cost in building.cost.items():
            if st.session_state.resources.get(resource, 0) < cost:
                return False
        return True

    def place_building(self, row: int, col: int, building_id: str) -> bool:
        if not self.can_afford(building_id):
            st.error(f"Cannot afford {self.buildings[building_id].name}!")
            return False

        if st.session_state.grid[row][col] is not None:
            st.error("Space already occupied!")
            return False

        # Deduct costs
        building = self.buildings[building_id]
        for resource, cost in building.cost.items():
            st.session_state.resources[resource] -= cost

        # Place building
        st.session_state.grid[row][col] = building_id
        st.success(f"Placed {building.name}")
        return True

def select_building(building_id: str):
    st.session_state.selected_building = building_id

def handle_cell_click(row: int, col: int, game: FactoryGame):
    if st.session_state.selected_building:
        game.place_building(row, col, st.session_state.selected_building)

def main():
    st.set_page_config(page_title="Factory Game", layout="wide")
    
    game = FactoryGame()
    
    # Title and description
    st.title("🏭 Factory Game")
    st.markdown("""
    Build your factory empire! Place buildings, manage resources, and optimize production.
    """)
    
    # Layout columns
    col1, col2 = st.columns([3, 1])
    
    with col2:
        # Resources display
        st.subheader("Resources")
        st.metric("⚡ Energy", f"{st.session_state.resources['energy']}")
        st.metric("🪨 Minerals", f"{st.session_state.resources['minerals']}")
        st.metric("🔌 Circuits", f"{st.session_state.resources['circuits']}")
        
        # Building selection
        st.subheader("Buildings")
        for building_id, building in game.buildings.items():
            if st.button(
                f"{building.icon} {building.name}\n"
                f"Cost: {', '.join(f'{cost} {res}' for res, cost in building.cost.items())}",
                help=building.description,
                key=f"building_{building_id}",
                type="secondary" if st.session_state.selected_building != building_id else "primary"
            ):
                select_building(building_id)

        # Currently selected building
        if st.session_state.selected_building:
            st.caption(f"Selected: {game.buildings[st.session_state.selected_building].name}")
        else:
            st.caption("Select a building to place")

    with col1:
        # Game grid
        st.subheader("Factory Grid")
        
        # Create a grid of buttons using containers for better layout
        grid_container = st.container()
        for row in range(game.GRID_SIZE):
            cols = st.columns(game.GRID_SIZE)
            for col, column in enumerate(cols):
                building_id = st.session_state.grid[row][col]
                building = game.buildings.get(building_id) if building_id else None
                
                # Style the button based on what's there
                label = building.icon if building else "⬜"
                with column:
                    if st.button(
                        label,
                        key=f"cell_{row}_{col}",
                        help=building.name if building else "Empty space",
                        use_container_width=True
                    ):
                        handle_cell_click(row, col, game)
    
    # Update resources periodically
    game.update_resources()
    
    # Force a rerun every second to update resources
    time.sleep(0.1)
    st.rerun()

if __name__ == "__main__":
    main()