import pygame
import sys
import json
from enum import Enum
from dataclasses import dataclass, asdict
from typing import Dict, List, Optional, Tuple
import random

# Initialize Pygame
pygame.init()

# Constants
WINDOW_WIDTH = 1024
WINDOW_HEIGHT = 768
GRID_SIZE = 32
GRID_WIDTH = WINDOW_WIDTH // GRID_SIZE
GRID_HEIGHT = WINDOW_HEIGHT // GRID_SIZE

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GREEN = (34, 139, 34)
BROWN = (139, 69, 19)
GRAY = (128, 128, 128)
BLUE = (30, 144, 255)
YELLOW = (255, 215, 0)
RED = (255, 0, 0)
LIGHT_GREEN = (144, 238, 144)

class BuildingType(Enum):
    HOUSE = "house"
    FARM = "farm"
    SHOP = "shop"
    WELL = "well"
    ROAD = "road"

@dataclass
class Building:
    x: int
    y: int
    building_type: BuildingType
    level: int = 1
    
    def get_color(self):
        colors = {
            BuildingType.HOUSE: BROWN,
            BuildingType.FARM: LIGHT_GREEN,
            BuildingType.SHOP: BLUE,
            BuildingType.WELL: GRAY,
            BuildingType.ROAD: BLACK
        }
        return colors.get(self.building_type, WHITE)

@dataclass
class Resources:
    gold: int = 100
    population: int = 0
    food: int = 50
    happiness: int = 50

class Game:
    def __init__(self):
        self.screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
        pygame.display.set_caption("Tiny Village")
        self.clock = pygame.time.Clock()
        self.running = True
        
        # Game state
        self.resources = Resources()
        self.buildings: List[Building] = []
        self.selected_building_type = BuildingType.HOUSE
        self.grid = [[None for _ in range(GRID_WIDTH)] for _ in range(GRID_HEIGHT)]
        
        # Building costs
        self.building_costs = {
            BuildingType.HOUSE: {'gold': 50, 'food': 10},
            BuildingType.FARM: {'gold': 30, 'food': 5},
            BuildingType.SHOP: {'gold': 80, 'food': 15},
            BuildingType.WELL: {'gold': 60, 'food': 8},
            BuildingType.ROAD: {'gold': 10, 'food': 0}
        }
        
        # Font for UI
        self.font = pygame.font.Font(None, 24)
        self.small_font = pygame.font.Font(None, 18)
        
        # Game timer for resource generation
        self.last_update = pygame.time.get_ticks()
        
    def handle_events(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                self.running = False
            elif event.type == pygame.KEYDOWN:
                self.handle_keypress(event.key)
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 1:  # Left click
                    self.handle_click(event.pos)
    
    def handle_keypress(self, key):
        # Building selection keys
        if key == pygame.K_1:
            self.selected_building_type = BuildingType.HOUSE
        elif key == pygame.K_2:
            self.selected_building_type = BuildingType.FARM
        elif key == pygame.K_3:
            self.selected_building_type = BuildingType.SHOP
        elif key == pygame.K_4:
            self.selected_building_type = BuildingType.WELL
        elif key == pygame.K_5:
            self.selected_building_type = BuildingType.ROAD
        elif key == pygame.K_s:
            self.save_game()
        elif key == pygame.K_l:
            self.load_game()
    
    def handle_click(self, pos):
        grid_x = pos[0] // GRID_SIZE
        grid_y = pos[1] // GRID_SIZE
        
        # Check if click is within game area (not UI area)
        if grid_x < GRID_WIDTH and grid_y < GRID_HEIGHT - 3:  # Reserve bottom 3 rows for UI
            self.place_building(grid_x, grid_y)
    
    def place_building(self, x, y):
        # Check if position is already occupied
        if self.grid[y][x] is not None:
            return False
        
        # Check if player has enough resources
        costs = self.building_costs[self.selected_building_type]
        if self.resources.gold < costs['gold'] or self.resources.food < costs['food']:
            return False
        
        # Place the building
        building = Building(x, y, self.selected_building_type)
        self.buildings.append(building)
        self.grid[y][x] = building
        
        # Deduct resources
        self.resources.gold -= costs['gold']
        self.resources.food -= costs['food']
        
        return True
    
    def update_resources(self):
        current_time = pygame.time.get_ticks()
        if current_time - self.last_update > 2000:  # Update every 2 seconds
            # Calculate resource generation based on buildings
            gold_income = 0
            food_income = 0
            population_change = 0
            happiness_change = 0
            
            for building in self.buildings:
                if building.building_type == BuildingType.HOUSE:
                    population_change += 2
                    happiness_change += 1
                    gold_income += 5
                elif building.building_type == BuildingType.FARM:
                    food_income += 10
                elif building.building_type == BuildingType.SHOP:
                    gold_income += 15
                    happiness_change += 2
                elif building.building_type == BuildingType.WELL:
                    happiness_change += 3
            
            # Apply changes
            self.resources.gold += gold_income
            self.resources.food += food_income
            self.resources.population = max(0, self.resources.population + population_change)
            self.resources.happiness = max(0, min(100, self.resources.happiness + happiness_change))
            
            # Food consumption
            food_consumption = self.resources.population // 2
            self.resources.food = max(0, self.resources.food - food_consumption)
            
            self.last_update = current_time
    
    def draw_grid(self):
        # Draw grid lines
        for x in range(0, WINDOW_WIDTH, GRID_SIZE):
            pygame.draw.line(self.screen, GRAY, (x, 0), (x, WINDOW_HEIGHT - 100))
        for y in range(0, WINDOW_HEIGHT - 100, GRID_SIZE):
            pygame.draw.line(self.screen, GRAY, (0, y), (WINDOW_WIDTH, y))
    
    def draw_buildings(self):
        for building in self.buildings:
            rect = pygame.Rect(
                building.x * GRID_SIZE + 2,
                building.y * GRID_SIZE + 2,
                GRID_SIZE - 4,
                GRID_SIZE - 4
            )
            pygame.draw.rect(self.screen, building.get_color(), rect)
            
            # Draw building type indicator
            text = building.building_type.value[0].upper()
            text_surface = self.small_font.render(text, True, WHITE)
            text_rect = text_surface.get_rect(center=rect.center)
            self.screen.blit(text_surface, text_rect)
    
    def draw_ui(self):
        # UI background
        ui_rect = pygame.Rect(0, WINDOW_HEIGHT - 100, WINDOW_WIDTH, 100)
        pygame.draw.rect(self.screen, BLACK, ui_rect)
        
        # Resource display
        y_offset = WINDOW_HEIGHT - 90
        resources_text = [
            f"Gold: {self.resources.gold}",
            f"Population: {self.resources.population}",
            f"Food: {self.resources.food}",
            f"Happiness: {self.resources.happiness}/100"
        ]
        
        for i, text in enumerate(resources_text):
            color = WHITE
            if "Food: 0" in text:
                color = RED
            elif "Happiness: 0" in text:
                color = RED
            
            text_surface = self.font.render(text, True, color)
            self.screen.blit(text_surface, (10 + i * 150, y_offset))
        
        # Building selection display
        y_offset = WINDOW_HEIGHT - 60
        building_options = [
            "1: House (50g, 10f)",
            "2: Farm (30g, 5f)",
            "3: Shop (80g, 15f)",
            "4: Well (60g, 8f)",
            "5: Road (10g, 0f)"
        ]
        
        for i, text in enumerate(building_options):
            color = WHITE
            building_type = list(BuildingType)[i]
            if building_type == self.selected_building_type:
                color = YELLOW
            
            # Check if player can afford this building
            costs = self.building_costs[building_type]
            if self.resources.gold < costs['gold'] or self.resources.food < costs['food']:
                color = RED
            
            text_surface = self.small_font.render(text, True, color)
            self.screen.blit(text_surface, (10, y_offset + i * 20))
        
        # Instructions
        instructions = [
            "Click to place selected building",
            "Keys 1-5: Select building type",
            "S: Save game, L: Load game"
        ]
        
        for i, instruction in enumerate(instructions):
            text_surface = self.small_font.render(instruction, True, WHITE)
            self.screen.blit(text_surface, (WINDOW_WIDTH - 250, WINDOW_HEIGHT - 90 + i * 20))
    
    def save_game(self):
        game_data = {
            'resources': asdict(self.resources),
            'buildings': [asdict(building) for building in self.buildings]
        }
        
        # Convert enum values to strings for JSON serialization
        for building_data in game_data['buildings']:
            building_data['building_type'] = building_data['building_type'].value
        
        try:
            with open('tiny_village_save.json', 'w') as f:
                json.dump(game_data, f, indent=2)
            print("Game saved successfully!")
        except Exception as e:
            print(f"Error saving game: {e}")
    
    def load_game(self):
        try:
            with open('tiny_village_save.json', 'r') as f:
                game_data = json.load(f)
            
            # Load resources
            self.resources = Resources(**game_data['resources'])
            
            # Load buildings
            self.buildings = []
            self.grid = [[None for _ in range(GRID_WIDTH)] for _ in range(GRID_HEIGHT)]
            
            for building_data in game_data['buildings']:
                building_data['building_type'] = BuildingType(building_data['building_type'])
                building = Building(**building_data)
                self.buildings.append(building)
                self.grid[building.y][building.x] = building
            
            print("Game loaded successfully!")
        except FileNotFoundError:
            print("No save file found!")
        except Exception as e:
            print(f"Error loading game: {e}")
    
    def run(self):
        while self.running:
            self.handle_events()
            self.update_resources()
            
            # Draw everything
            self.screen.fill(GREEN)
            self.draw_grid()
            self.draw_buildings()
            self.draw_ui()
            
            pygame.display.flip()
            self.clock.tick(60)
        
        pygame.quit()
        sys.exit()

if __name__ == "__main__":
    game = Game()
    game.run()