import pygame
import button
import csv
import os

pygame.init()

clock = pygame.time.Clock()
FPS = 60

# Game window
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 640
LOWER_MARGIN = 100
SIDE_MARGIN = 300

screen = pygame.display.set_mode((SCREEN_WIDTH + SIDE_MARGIN, SCREEN_HEIGHT + LOWER_MARGIN))
pygame.display.set_caption('Level Editor')

# Define game variables
ROWS = 16
MAX_COLS = 150
TILE_SIZE = SCREEN_HEIGHT // ROWS
TILE_TYPES = 21
level = 0
current_tile = 0
scroll_left = False
scroll_right = False
scroll = 0
scroll_speed = 1
clicked = False

# Load images safely
def load_image(filename, scale=None):
    if os.path.exists(filename):
        img = pygame.image.load(filename).convert_alpha()
        if scale:
            img = pygame.transform.scale(img, scale)
        return img
    else:
        print(f"Warning: Image '{filename}' not found.")
        return pygame.Surface((TILE_SIZE, TILE_SIZE))  # return blank tile

pine1_img = load_image('pine1.png')
pine2_img = load_image('pine2.png')
mountain_img = load_image('mountain.png')
sky_img = load_image('sky_cloud.png')

# Store tile images
img_list = []
for x in range(TILE_TYPES):
    img = load_image(f'{x}.png', (TILE_SIZE, TILE_SIZE))
    img_list.append(img)

save_img = load_image('save_btn.png')
load_img = load_image('load_btn.png')

# Define colours
GREEN = (144, 201, 120)
WHITE = (255, 255, 255)
RED = (200, 25, 25)

# Define font
font = pygame.font.SysFont('Futura', 30)

# Create empty tile list
world_data = [[-1] * MAX_COLS for _ in range(ROWS)]

# Create ground layer
for tile in range(MAX_COLS):
    world_data[ROWS - 1][tile] = 0

# Draw text
def draw_text(text, font, text_col, x, y):
    img = font.render(text, True, text_col)
    screen.blit(img, (x, y))

# Draw background
def draw_bg():
    screen.fill(GREEN)
    width = sky_img.get_width()
    for x in range(5):
        screen.blit(sky_img, ((x * width) - scroll * 0.5, 0))
        screen.blit(mountain_img, ((x * width) - scroll * 0.6, SCREEN_HEIGHT - mountain_img.get_height() - 300))
        screen.blit(pine1_img, ((x * width) - scroll * 0.7, SCREEN_HEIGHT - pine1_img.get_height() - 150))
        screen.blit(pine2_img, ((x * width) - scroll * 0.8, SCREEN_HEIGHT - pine2_img.get_height()))

# Draw grid
def draw_grid():
    for c in range(MAX_COLS + 1):
        pygame.draw.line(screen, WHITE, (c * TILE_SIZE - scroll, 0), (c * TILE_SIZE - scroll, SCREEN_HEIGHT))
    for r in range(ROWS + 1):
        pygame.draw.line(screen, WHITE, (0, r * TILE_SIZE), (SCREEN_WIDTH, r * TILE_SIZE))

# Draw world tiles
def draw_world():
    for y, row in enumerate(world_data):
        for x, tile in enumerate(row):
            if tile >= 0 and tile < len(img_list):
                screen.blit(img_list[tile], (x * TILE_SIZE - scroll, y * TILE_SIZE))

# Create buttons
save_button = button.Button(SCREEN_WIDTH // 2, SCREEN_HEIGHT + LOWER_MARGIN - 50, save_img, 1)
load_button = button.Button(SCREEN_WIDTH // 2 + 200, SCREEN_HEIGHT + LOWER_MARGIN - 50, load_img, 1)

# Tile selection buttons
button_list = []
button_col = 0
button_row = 0
for i, img in enumerate(img_list):
    tile_button = button.Button(SCREEN_WIDTH + (75 * button_col) + 50, 75 * button_row + 50, img, 1)
    button_list.append(tile_button)
    button_col += 1
    if button_col == 3:
        button_row += 1
        button_col = 0

# Main loop
run = True
while run:
    clock.tick(FPS)

    draw_bg()
    draw_grid()
    draw_world()

    draw_text(f'Level: {level}', font, WHITE, 10, SCREEN_HEIGHT + LOWER_MARGIN - 90)
    draw_text('Press UP or DOWN to change level', font, WHITE, 10, SCREEN_HEIGHT + LOWER_MARGIN - 60)

    if save_button.draw(screen):
        with open(f'level{level}_data.csv', 'w', newline='') as file:
            writer = csv.writer(file)
            writer.writerows(world_data)

    if load_button.draw(screen):
        scroll = 0
        try:
            with open(f'level{level}_data.csv', newline='') as file:
                reader = csv.reader(file)
                for x, row in enumerate(reader):
                    for y, tile in enumerate(row):
                        world_data[x][y] = int(tile)
        except Exception as e:
            print(f"Failed to load level: {e}")

    pygame.draw.rect(screen, GREEN, (SCREEN_WIDTH, 0, SIDE_MARGIN, SCREEN_HEIGHT))

    # Tile selection
    for i, tile_button in enumerate(button_list):
        if tile_button.draw(screen):
            current_tile = i
    pygame.draw.rect(screen, RED, button_list[current_tile].rect, 3)

    # Scroll map
    if scroll_left and scroll > 0:
        scroll -= 5 * scroll_speed
    if scroll_right and scroll < (MAX_COLS * TILE_SIZE - SCREEN_WIDTH):
        scroll += 5 * scroll_speed

    # Mouse tile placement
    mouse_pos = pygame.mouse.get_pos()
    mx = (mouse_pos[0] + scroll) // TILE_SIZE
    my = mouse_pos[1] // TILE_SIZE
    if 0 <= mx < MAX_COLS and 0 <= my < ROWS and mouse_pos[0] < SCREEN_WIDTH and mouse_pos[1] < SCREEN_HEIGHT:
        if pygame.mouse.get_pressed()[0] and not clicked:
            world_data[my][mx] = current_tile
            clicked = True
        if pygame.mouse.get_pressed()[2] and not clicked:
            world_data[my][mx] = -1
            clicked = True

    if not pygame.mouse.get_pressed()[0] and not pygame.mouse.get_pressed()[2]:
        clicked = False

    # Event handling
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            run = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP:
                level += 1
            if event.key == pygame.K_DOWN and level > 0:
                level -= 1
            if event.key == pygame.K_LEFT:
                scroll_left = True
            if event.key == pygame.K_RIGHT:
                scroll_right = True
            if event.key == pygame.K_RSHIFT:
                scroll_speed = 5
        if event.type == pygame.KEYUP:
            if event.key == pygame.K_LEFT:
                scroll_left = False
            if event.key == pygame.K_RIGHT:
                scroll_right = False
            if event.key == pygame.K_RSHIFT:
                scroll_speed = 1

    pygame.display.update()

pygame.quit()