import pygame
import json
import sys

# Editor settings.
SCREEN_WIDTH, SCREEN_HEIGHT = 1000, 800
BACKGROUND_COLOR = (30, 30, 30)
FPS = 60

# Colors for different element types.
COLORS = {
    "obstacle": (150, 50, 50),
    "road": (100, 100, 100),
    "checkpoint": (200, 200, 0),
    "finish_line": (0, 0, 255)
}

# Element types.
ELEMENT_TYPES = ["obstacle", "road", "checkpoint", "finish_line"]

def main():
    pygame.init()
    screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
    pygame.display.set_caption("Advanced Map Editor: Roads, Checkpoints & Finish Line")
    clock = pygame.time.Clock()
    font = pygame.font.SysFont(None, 24)
    
    # A list to store drawn elements as tuples: (type, rect dict)
    elements = []
    current_type = "obstacle"  # Default drawing type.
    
    drawing = False
    start_pos = (0, 0)
    current_rect = None

    instructions = [
        "Press O: Obstacle, R: Road, C: Checkpoint, F: Finish Line",
        "Click & drag to draw. Press S to save, U to undo last element, ESC to exit."
    ]

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    pygame.quit()
                    sys.exit()
                elif event.key == pygame.K_o:
                    current_type = "obstacle"
                elif event.key == pygame.K_r:
                    current_type = "road"
                elif event.key == pygame.K_c:
                    current_type = "checkpoint"
                elif event.key == pygame.K_f:
                    current_type = "finish_line"
                elif event.key == pygame.K_s:
                    # Organize elements by type.
                    map_data = {etype: [] for etype in ELEMENT_TYPES}
                    for etype, rect in elements:
                        map_data[etype].append(rect)
                    with open('map.json', 'w') as f:
                        json.dump(map_data, f, indent=4)
                    print("Map saved to map.json")
                elif event.key == pygame.K_u:
                    if elements:
                        removed = elements.pop()
                        print(f"Removed last element of type {removed[0]}")
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 1:  # Left click starts drawing.
                    drawing = True
                    start_pos = event.pos
            elif event.type == pygame.MOUSEBUTTONUP:
                if event.button == 1 and drawing:
                    drawing = False
                    end_pos = event.pos
                    x = min(start_pos[0], end_pos[0])
                    y = min(start_pos[1], end_pos[1])
                    width = abs(end_pos[0] - start_pos[0])
                    height = abs(end_pos[1] - start_pos[1])
                    if width > 5 and height > 5:
                        rect_dict = {"x": x, "y": y, "width": width, "height": height}
                        elements.append((current_type, rect_dict))
                    current_rect = None
            elif event.type == pygame.MOUSEMOTION:
                if drawing:
                    end_pos = event.pos
                    x = min(start_pos[0], end_pos[0])
                    y = min(start_pos[1], end_pos[1])
                    width = abs(end_pos[0] - start_pos[0])
                    height = abs(end_pos[1] - start_pos[1])
                    current_rect = pygame.Rect(x, y, width, height)
                    
        screen.fill(BACKGROUND_COLOR)
        
        # Draw saved elements.
        for etype, rect_dict in elements:
            rect = pygame.Rect(rect_dict["x"], rect_dict["y"], rect_dict["width"], rect_dict["height"])
            pygame.draw.rect(screen, COLORS.get(etype, (255,255,255)), rect)
        
        # Draw current preview rectangle.
        if current_rect:
            pygame.draw.rect(screen, COLORS.get(current_type, (255,255,255)), current_rect, 2)
        
        # Display instructions.
        y_offset = 10
        for line in instructions:
            surf = font.render(line, True, (255,255,255))
            screen.blit(surf, (10, y_offset))
            y_offset += 24
        type_text = font.render(f"Current Element Type: {current_type.upper()}", True, (255,255,255))
        screen.blit(type_text, (10, y_offset))
        
        pygame.display.flip()
        clock.tick(FPS)

if __name__ == "__main__":
    main()