import os
import shutil
import re

# Define paths
mod_dir = "."  # Current directory
mod_file = os.path.join(mod_dir, "mod.py")
backup_file = os.path.join(mod_dir, "mod.py.backup")

# Create backup
print(f"Creating backup of original mod.py to {backup_file}")
shutil.copy2(mod_file, backup_file)

# Read the original file
with open(mod_file, 'r') as f:
    content = f.read()

# Define the town generation code to inject
town_code = """
    def _generate_town_around_castle(self, chunk, center_x, center_y, castle_size):
        \"\"\"
        Generate a detailed town surrounding the grand castle
        
        Args:
            chunk: 2D array of chunk blocks
            center_x, center_y: Castle center coordinates
            castle_size: Size of the castle
        \"\"\"
        # Town generation parameters
        town_radius = castle_size + 10  # Town extends beyond castle
        
        # Clear area for town (keeping castle intact)
        for dx in range(-town_radius, town_radius + 1):
            for dy in range(-town_radius, town_radius + 1):
                x = center_x + dx
                y = center_y + dy
                
                # Skip if out of bounds
                if not (0 <= x < len(chunk[0]) and 0 <= y < len(chunk)):
                    continue
                
                # Don't overwrite castle (castle is within castle_size/2 of center)
                if abs(dx) <= castle_size//2 + 2 and abs(dy) <= castle_size//2 + 2:
                    continue
                    
                # Set to grass to ensure buildable terrain for town
                chunk[y][x] = 'grass'
        
        # Create main roads
        self._create_town_roads(chunk, center_x, center_y, castle_size, town_radius)
        
        # Generate various buildings
        self._generate_town_buildings(chunk, center_x, center_y, castle_size, town_radius)
        
        # Add decorative elements
        self._add_town_decorations(chunk, center_x, center_y, castle_size, town_radius)
        
        # Add town NPCs (using spawners)
        self._add_town_npcs(chunk, center_x, center_y, castle_size, town_radius)
        
    def _create_town_roads(self, chunk, center_x, center_y, castle_size, town_radius):
        \"\"\"Create road network through the town\"\"\"
        # Road material
        road_block = 'cobblestone'
        half_castle = castle_size // 2
        
        # Create four main roads extending from castle doors (N, E, S, W)
        for direction in [(0, -1), (1, 0), (0, 1), (-1, 0)]:  # N, E, S, W
            dx, dy = direction
            
            # Start from the castle wall
            start_x = center_x + dx * half_castle
            start_y = center_y + dy * half_castle
            
            # End at town border
            end_x = center_x + dx * town_radius
            end_y = center_y + dy * town_radius
            
            # Create straight road
            for i in range(town_radius):
                x = start_x + dx * i
                y = start_y + dy * i
                
                # Skip if out of bounds
                if not (0 <= x < len(chunk[0]) and 0 <= y < len(chunk)):
                    continue
                
                # Main road is 3 blocks wide
                for offset in range(-1, 2):
                    if dx == 0:  # North-South road
                        road_x = x + offset
                        road_y = y
                    else:  # East-West road
                        road_x = x
                        road_y = y + offset
                    
                    if 0 <= road_x < len(chunk[0]) and 0 <= road_y < len(chunk):
                        chunk[road_y][road_x] = road_block
        
        # Create circular road around the castle
        ring_radius = half_castle + 5
        for angle in range(0, 360, 5):  # Every 5 degrees
            radian = math.radians(angle)
            x = int(center_x + ring_radius * math.cos(radian))
            y = int(center_y + ring_radius * math.sin(radian))
            
            if 0 <= x < len(chunk[0]) and 0 <= y < len(chunk):
                chunk[y][x] = road_block
                
                # Make road 3 blocks wide
                for offset in range(1, 2):
                    inner_x = int(center_x + (ring_radius-offset) * math.cos(radian))
                    inner_y = int(center_y + (ring_radius-offset) * math.sin(radian))
                    
                    if 0 <= inner_x < len(chunk[0]) and 0 <= inner_y < len(chunk):
                        chunk[inner_y][inner_x] = road_block
        
        # Add smaller connecting roads
        quadrants = [
            (1, 1),   # Northeast
            (1, -1),  # Southeast
            (-1, 1),  # Northwest
            (-1, -1)  # Southwest
        ]
        
        for quad_x, quad_y in quadrants:
            # Create 2-3 smaller roads in each quadrant
            for _ in range(2):
                # Random starting point on main road
                if random.random() < 0.5:  # Start on N/S road
                    start_x = center_x + quad_x * random.randint(3, half_castle - 2)
                    start_y = center_y + quad_y * half_castle
                    direction = (0, quad_y)  # Continue in N/S direction
                else:  # Start on E/W road
                    start_x = center_x + quad_x * half_castle
                    start_y = center_y + quad_y * random.randint(3, half_castle - 2)
                    direction = (quad_x, 0)  # Continue in E/W direction
                
                # Create side road
                road_length = random.randint(5, town_radius - half_castle)
                for i in range(road_length):
                    x = start_x + direction[0] * i
                    y = start_y + direction[1] * i
                    
                    if not (0 <= x < len(chunk[0]) and 0 <= y < len(chunk)):
                        break
                    
                    # Side roads are 2 blocks wide
                    chunk[y][x] = road_block
                    
                    # Add width to road
                    width_dir = (direction[1], direction[0])  # Perpendicular
                    side_x = x + width_dir[0]
                    side_y = y + width_dir[1]
                    
                    if 0 <= side_x < len(chunk[0]) and 0 <= side_y < len(chunk):
                        chunk[side_y][side_x] = road_block
    
    def _generate_town_buildings(self, chunk, center_x, center_y, castle_size, town_radius):
        \"\"\"Generate various town buildings\"\"\"
        half_castle = castle_size // 2
        
        # Define building types with their sizes and colors
        building_types = {
            'bank': {
                'size': (7, 5),
                'wall': 'brick',
                'floor': 'marble',
                'roof': 'gold',
                'min_dist': half_castle + 6,
                'max_dist': half_castle + 14,
                'door_direction': (0, 1)  # Door faces south
            },
            'shop': {
                'size': (5, 6),
                'wall': 'wood',
                'floor': 'wood',
                'roof': 'roof',
                'min_dist': half_castle + 7,
                'max_dist': half_castle + 15,
                'door_direction': (1, 0)  # Door faces east
            },
            'smithy': {
                'size': (6, 6),
                'wall': 'stone',
                'floor': 'stone',
                'roof': 'roof',
                'min_dist': half_castle + 8,
                'max_dist': half_castle + 16,
                'door_direction': (0, 1)  # Door faces south
            },
            'bakery': {
                'size': (5, 5),
                'wall': 'brick',
                'floor': 'wood',
                'roof': 'roof',
                'min_dist': half_castle + 6,
                'max_dist': half_castle + 13,
                'door_direction': (1, 0)  # Door faces east
            },
            'wizard_tower': {
                'size': (4, 4),
                'height': 8,
                'wall': 'cobblestone',
                'floor': 'obsidian',
                'roof': 'tower_top',
                'min_dist': half_castle + 10,
                'max_dist': town_radius - 5,
                'door_direction': (-1, 0)  # Door faces west
            },
            'inn': {
                'size': (8, 6),
                'wall': 'wood',
                'floor': 'wood',
                'roof': 'roof',
                'min_dist': half_castle + 7,
                'max_dist': half_castle + 16,
                'door_direction': (0, 1)  # Door faces south
            },
            'stable': {
                'size': (7, 5),
                'wall': 'wood',
                'floor': 'dirt',
                'roof': 'wood',
                'min_dist': half_castle + 8,
                'max_dist': town_radius - 3,
                'door_direction': (0, 1)  # Door faces south
            },
            'barracks': {
                'size': (9, 7),
                'wall': 'stone',
                'floor': 'stone',
                'roof': 'roof',
                'min_dist': half_castle + 10,
                'max_dist': town_radius - 4,
                'door_direction': (0, -1)  # Door faces north (toward castle)
            },
            'houses': {
                'count': 15,  # Generate multiple houses
                'size_range': (4, 6),
                'wall': 'wood',
                'floor': 'wood',
                'roof': 'roof',
                'min_dist': half_castle + 7,
                'max_dist': town_radius - 2,
                'door_direction': None  # Random
            },
            'tents': {
                'count': 8,
                'size_range': (3, 4),
                'wall': 'banner_red',
                'floor': 'dirt',
                'roof': None,
                'min_dist': half_castle + 5,
                'max_dist': town_radius - 1,
                'door_direction': None  # Random
            }
        }
        
        # Position for each special building (angle from center)
        positions = {
            'bank': 45,          # Northeast
            'shop': 120,         # Southeast
            'smithy': 210,       # Southwest
            'bakery': 270,       # West
            'wizard_tower': 330, # Northwest
            'inn': 165,          # South
            'stable': 75,        # East-Northeast
            'barracks': 300      # West-Northwest
        }
        
        # First place the special buildings
        for building_name, angle in positions.items():
            building = building_types[building_name]
            
            # Calculate building center position
            distance = random.randint(building['min_dist'], building['max_dist'])
            radian = math.radians(angle)
            building_center_x = int(center_x + distance * math.cos(radian))
            building_center_y = int(center_y + distance * math.sin(radian))
            
            # Determine building size
            if 'size' in building:
                width, height = building['size']
            else:
                min_size, max_size = building['size_range']
                width = random.randint(min_size, max_size)
                height = random.randint(min_size, max_size)
            
            # Build the special building
            if building_name == 'wizard_tower':
                self._build_tower(chunk, building_center_x, building_center_y, 
                               building['wall'], building['floor'], building['roof'],
                               width, building['height'], building['door_direction'])
            else:
                self._build_house(chunk, building_center_x, building_center_y, 
                               building['wall'], building['floor'], building['roof'],
                               width, height, building['door_direction'])
        
        # Then place the random houses and tents
        for building_name in ['houses', 'tents']:
            building = building_types[building_name]
            
            for _ in range(building['count']):
                # Find a suitable spot
                for attempt in range(20):  # Try 20 times
                    angle = random.randint(0, 359)
                    distance = random.randint(building['min_dist'], building['max_dist'])
                    radian = math.radians(angle)
                    
                    building_center_x = int(center_x + distance * math.cos(radian))
                    building_center_y = int(center_y + distance * math.sin(radian))
                    
                    # Check if spot is suitable (on grass, not on road)
                    if (0 <= building_center_x < len(chunk[0]) and 
                        0 <= building_center_y < len(chunk) and
                        chunk[building_center_y][building_center_x] == 'grass'):
                        
                        # Determine building size
                        min_size, max_size = building['size_range']
                        width = random.randint(min_size, max_size)
                        height = random.randint(min_size, max_size)
                        
                        # Determine door direction
                        if building['door_direction'] is None:
                            door_direction = random.choice([(0, 1), (1, 0), (0, -1), (-1, 0)])
                        else:
                            door_direction = building['door_direction']
                        
                        # Build the house or tent
                        if building_name == 'tents':
                            self._build_tent(chunk, building_center_x, building_center_y, 
                                          building['wall'], building['floor'], 
                                          width, height, door_direction)
                        else:
                            self._build_house(chunk, building_center_x, building_center_y, 
                                           building['wall'], building['floor'], building['roof'],
                                           width, height, door_direction)
                        break
    
    def _build_house(self, chunk, center_x, center_y, wall_block, floor_block, roof_block, width, height, door_direction):
        \"\"\"Build a house with specified parameters\"\"\"
        half_width = width // 2
        half_height = height // 2
        
        # Build walls and floor
        for dx in range(-half_width, half_width + 1):
            for dy in range(-half_height, half_height + 1):
                x = center_x + dx
                y = center_y + dy
                
                # Skip if out of bounds
                if not (0 <= x < len(chunk[0]) and 0 <= y < len(chunk)):
                    continue
                
                # Walls
                if dx == -half_width or dx == half_width or dy == -half_height or dy == half_height:
                    chunk[y][x] = wall_block
                else:
                    # Floor
                    chunk[y][x] = floor_block
        
        # Add door in specified direction
        if door_direction:
            door_x = center_x + door_direction[0] * half_width
            door_y = center_y + door_direction[1] * half_height
            
            if 0 <= door_x < len(chunk[0]) and 0 <= door_y < len(chunk):
                chunk[door_y][door_x] = 'door_closed'
        
        # Add roof
        if roof_block:
            for dx in range(-half_width, half_width + 1):
                for dy in range(-half_height, half_height + 1):
                    x = center_x + dx
                    y = center_y + dy
                    
                    # Skip if out of bounds
                    if not (0 <= x < len(chunk[0]) and 0 <= y < len(chunk)):
                        continue
                    
                    # Roof (skip edges to make it look better)
                    if abs(dx) < half_width and abs(dy) < half_height:
                        # Add at one block above floor level
                        if y > 0 and y - 1 < len(chunk):
                            chunk[y-1][x] = roof_block
        
        # Add lantern
        lantern_x = center_x
        lantern_y = center_y
        if 0 <= lantern_x < len(chunk[0]) and 0 <= lantern_y < len(chunk):
            chunk[lantern_y][lantern_x] = 'lantern'
    
    def _build_tower(self, chunk, center_x, center_y, wall_block, floor_block, roof_block, width, height, door_direction):
        \"\"\"Build a tall tower with specified parameters\"\"\"
        radius = width // 2
        
        # Create circular tower
        for y_offset in range(height):
            for dx in range(-radius-1, radius+2):
                for dy in range(-radius-1, radius+2):
                    x = center_x + dx
                    y = center_y + dy
                    
                    # Skip if out of bounds
                    if not (0 <= x < len(chunk[0]) and 0 <= y < len(chunk)):
                        continue
                    
                    # Calculate distance from center
                    distance = math.sqrt(dx*dx + dy*dy)
                    
                    # For the first level (ground)
                    if y_offset == 0:
                        if distance <= radius:
                            # Interior/floor
                            chunk[y][x] = floor_block
                        elif distance <= radius + 0.5:
                            # Wall
                            chunk[y][x] = wall_block
                    else:
                        # We can only represent the top of the tower in 2D
                        # So for the top level, add tower walls and roof
                        if y_offset == height - 1:
                            if distance <= radius:
                                # Interior
                                chunk[y][x] = floor_block
                            elif distance <= radius + 0.5:
                                # Wall
                                chunk[y][x] = wall_block
                            elif distance <= radius + 1.5 and (dx == 0 or dy == 0 or abs(dx) == abs(dy)):
                                # Roof decorations at cardinal points and diagonals
                                chunk[y][x] = roof_block
        
        # Add door in specified direction
        if door_direction:
            door_x = center_x + door_direction[0] * (radius + 1)
            door_y = center_y + door_direction[1] * (radius + 1)
            
            if 0 <= door_x < len(chunk[0]) and 0 <= door_y < len(chunk):
                chunk[door_y][door_x] = 'door_closed'
        
        # Add light source
        chunk[center_y][center_x] = 'lantern'
    
    def _build_tent(self, chunk, center_x, center_y, wall_block, floor_block, width, height, door_direction):
        \"\"\"Build a tent with specified parameters\"\"\"
        half_width = width // 2
        half_height = height // 2
        
        # Build tent floor
        for dx in range(-half_width, half_width + 1):
            for dy in range(-half_height, half_height + 1):
                x = center_x + dx
                y = center_y + dy
                
                # Skip if out of bounds
                if not (0 <= x < len(chunk[0]) and 0 <= y < len(chunk)):
                    continue
                
                # Floor
                chunk[y][x] = floor_block
                
                # Walls only on perimeter
                if dx == -half_width or dx == half_width or dy == -half_height or dy == half_height:
                    chunk[y][x] = wall_block
        
        # Add entrance
        if door_direction:
            door_x = center_x + door_direction[0] * half_width
            door_y = center_y + door_direction[1] * half_height
            
            if 0 <= door_x < len(chunk[0]) and 0 <= door_y < len(chunk):
                chunk[door_y][door_x] = 'air'  # Tents have an open entry
    
    def _add_town_decorations(self, chunk, center_x, center_y, castle_size, town_radius):
        \"\"\"Add decorative elements throughout the town\"\"\"
        half_castle = castle_size // 2
        
        # Add lampposts along roads
        for direction in [(0, -1), (1, 0), (0, 1), (-1, 0)]:  # N, E, S, W
            dx, dy = direction
            
            # Start from just outside the castle
            start_x = center_x + dx * (half_castle + 1)
            start_y = center_y + dy * (half_castle + 1)
            
            # Place lanterns every 6 blocks along roads
            for i in range(1, town_radius - half_castle, 6):
                x = start_x + dx * i
                y = start_y + dy * i
                
                # Offset to the side of the road
                if dx == 0:  # North-South road
                    lamp_x = x + 2
                    lamp_y = y
                else:  # East-West road
                    lamp_x = x
                    lamp_y = y + 2
                
                if 0 <= lamp_x < len(chunk[0]) and 0 <= lamp_y < len(chunk):
                    chunk[lamp_y][lamp_x] = 'lantern'
        
        # Add a market square in one quadrant
        market_distance = half_castle + 10
        market_radian = math.radians(135)  # Southeast
        market_center_x = int(center_x + market_distance * math.cos(market_radian))
        market_center_y = int(center_y + market_distance * math.sin(market_radian))
        market_size = 6
        
        # Create market square
        for dx in range(-market_size, market_size + 1):
            for dy in range(-market_size, market_size + 1):
                x = market_center_x + dx
                y = market_center_y + dy
                
                # Skip if out of bounds
                if not (0 <= x < len(chunk[0]) and 0 <= y < len(chunk)):
                    continue
                
                # Market square has a special floor
                if abs(dx) < market_size and abs(dy) < market_size:
                    chunk[y][x] = 'stone'
                
                # Add market stalls (wooden structures) around the perimeter
                if (abs(dx) == market_size - 1 or abs(dy) == market_size - 1) and random.random() < 0.3:
                    chunk[y][x] = 'wood'
        
        # Add a fountain in the center of the market
        chunk[market_center_y][market_center_x] = 'water'
        for dx, dy in [(0, 1), (1, 0), (0, -1), (-1, 0), (1, 1), (1, -1), (-1, 1), (-1, -1)]:
            x = market_center_x + dx
            y = market_center_y + dy
            if 0 <= x < len(chunk[0]) and 0 <= y < len(chunk):
                chunk[y][x] = 'stone'
        
        # Add a town square near the castle
        square_distance = half_castle + 4
        square_radian = math.radians(0)  # North
        square_center_x = int(center_x + square_distance * math.cos(square_radian))
        square_center_y = int(center_y + square_distance * math.sin(square_radian))
        square_size = 4
        
        # Create town square
        for dx in range(-square_size, square_size + 1):
            for dy in range(-square_size, square_size + 1):
                x = square_center_x + dx
                y = square_center_y + dy
                
                # Skip if out of bounds
                if not (0 <= x < len(chunk[0]) and 0 <= y < len(chunk)):
                    continue
                
                # Town square has decorative flooring
                chunk[y][x] = 'castle_floor'
        
        # Add statue in the center of the town square
        chunk[square_center_y][square_center_x] = 'castle_pillar'
        
        # Add a park or garden
        garden_distance = half_castle + 12
        garden_radian = math.radians(240)  # Southwest
        garden_center_x = int(center_x + garden_distance * math.cos(garden_radian))
        garden_center_y = int(center_y + garden_distance * math.sin(garden_radian))
        garden_size = 5
        
        # Create garden
        for dx in range(-garden_size, garden_size + 1):
            for dy in range(-garden_size, garden_size + 1):
                x = garden_center_x + dx
                y = garden_center_y + dy
                
                # Skip if out of bounds
                if not (0 <= x < len(chunk[0]) and 0 <= y < len(chunk)):
                    continue
                
                # Garden has special grass and trees
                dist_from_center = abs(dx) + abs(dy)
                if dist_from_center < garden_size:
                    if dist_from_center > garden_size // 2 and random.random() < 0.3:
                        # Trees around the edge
                        chunk[y][x] = 'leaves'
                    else:
                        # Garden floor
                        chunk[y][x] = 'grass'
        
        # Add random flowers and details throughout town
        for _ in range(50):
            dx = random.randint(-town_radius, town_radius)
            dy = random.randint(-town_radius, town_radius)
            
            x = center_x + dx
            y = center_y + dy
            
            # Skip if out of bounds or too close to castle
            if not (0 <= x < len(chunk[0]) and 0 <= y < len(chunk)) or (abs(dx) < half_castle and abs(dy) < half_castle):
                continue
            
            # Only place on grass
            if chunk[y][x] == 'grass':
                # 50% chance for flowers, 50% for other decorative elements
                if random.random() < 0.5:
                    # Simple flowers - use colored blocks
                    if random.random() < 0.5:
                        chunk[y][x] = 'banner_red'  # Red flowers
                    else:
                        chunk[y][x] = 'banner_blue'  # Blue flowers
                else:
                    # Other decorative elements
                    if random.random() < 0.3:
                        chunk[y][x] = 'wood'  # Small wooden structure
                    elif random.random() < 0.5:
                        chunk[y][x] = 'stone'  # Stone decoration
    
    def _add_town_npcs(self, chunk, center_x, center_y, castle_size, town_radius):
        \"\"\"Add NPCs throughout the town using spawners\"\"\"
        half_castle = castle_size // 2
        
        # Add villager spawners in key locations
        spawn_locations = [
            # Town square
            (center_x, center_y - (half_castle + 4)),
            
            # Market
            (center_x + int(math.cos(math.radians(135)) * (half_castle + 10)), 
             center_y + int(math.sin(math.radians(135)) * (half_castle + 10))),
            
            # Random street corners
            (center_x + half_castle // 2, center_y + half_castle + 3),
            (center_x - half_castle - 6, center_y - half_castle // 2),
            (center_x + half_castle + 8, center_y - half_castle - 4),
            
            # Near the bank
            (center_x + int(math.cos(math.radians(45)) * (half_castle + 8)), 
             center_y + int(math.sin(math.radians(45)) * (half_castle + 8))),
            
            # Near shops
            (center_x + int(math.cos(math.radians(120)) * (half_castle + 9)), 
             center_y + int(math.sin(math.radians(120)) * (half_castle + 9))),
        ]
        
        # Create spawners at designated locations
        for spawn_x, spawn_y in spawn_locations:
            if 0 <= spawn_x < len(chunk[0]) and 0 <= spawn_y < len(chunk):
                # Only place on valid blocks (not walls)
                if chunk[spawn_y][spawn_x] in ['grass', 'castle_floor', 'stone', 'cobblestone']:
                    chunk[spawn_y][spawn_x] = 'spawner'
"""

# Define the modification point in the _generate_grand_castle method
pattern_to_modify = r'([ ]{8}# Return castle data\n[ ]{8}return \{.*?center_y': center_y,.*?\})'

# Inject call to the new town generation method
inject_call = """        # Generate town around the castle
        self._generate_town_around_castle(chunk, center_x, center_y, castle_size)
        
        # Return castle data
        return {
            'size': castle_size,
            'center_x': center_x,
            'center_y': center_y"""

# Replace the pattern with our modified version
modified_content = re.sub(pattern_to_modify, inject_call, content, flags=re.DOTALL)

# Add the new town generation methods to the mod.py file
# Find the end of the last method in the class
last_method_pattern = r'([ ]{4}def [^(]+\([^)]+\):.*?\n)(?=[ ]{4}# Optional save and load functions)'
modified_content = re.sub(last_method_pattern, r'\1' + town_code, modified_content, flags=re.DOTALL)

# Write the modified content back to the file
with open(mod_file, 'w') as f:
    f.write(modified_content)