import os
import pygame
from os.path import join

# Initialize pygame and set up the display
pygame.init()
pygame.display.set_caption("Level Editor")

WIDTH, HEIGHT = 1000, 800
FPS = 60
BLOCK_SIZE = 96

window = pygame.display.set_mode((WIDTH, HEIGHT))

# Utility functions
def get_block(size):
    path = join("assets", "Terrain", "Terrain.png")
    image = pygame.image.load(path).convert_alpha()
    surface = pygame.Surface((size, size), pygame.SRCALPHA, 32)
    rect = pygame.Rect(96, 0, size, size)
    surface.blit(image, (0, 0), rect)
    return pygame.transform.scale2x(surface)

def load_sprite_sheets(dir1, dir2, width, height, direction=False):
    path = join("assets", dir1, dir2)
    images = [f for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]

    all_sprites = {}

    for image in images:
        sprite_sheet = pygame.image.load(join(path, image)).convert_alpha()

        sprites = []
        for i in range(sprite_sheet.get_width() // width):
            surface = pygame.Surface((width, height), pygame.SRCALPHA, 32)
            rect = pygame.Rect(i * width, 0, width, height)
            surface.blit(sprite_sheet, (0, 0), rect)
            sprites.append(pygame.transform.scale2x(surface))

        if direction:
            all_sprites[image.replace(".png", "") + "_right"] = sprites
            all_sprites[image.replace(".png", "") + "_left"] = flip(sprites)
        else:
            all_sprites[image.replace(".png", "")] = sprites

    return all_sprites

# Object classes
class Object(pygame.sprite.Sprite):
    def __init__(self, x, y, width, height, name=None):
        super().__init__()
        self.rect = pygame.Rect(x, y, width, height)
        self.image = pygame.Surface((width, height), pygame.SRCALPHA)
        self.width = width
        self.height = height
        self.name = name

    def draw(self, win):
        win.blit(self.image, (self.rect.x, self.rect.y))

class Block(Object):
    def __init__(self, x, y, size):
        super().__init__(x, y, size, size)
        block = get_block(size)
        self.image.blit(block, (0, 0))
        self.mask = pygame.mask.from_surface(self.image)

class Fire(Object):
    SPRITES = load_sprite_sheets("Traps", "Fire", 16, 32)

    def __init__(self, x, y, width, height):
        super().__init__(x, y, width, height, "fire")
        self.image = self.SPRITES["off"][0]
        self.mask = pygame.mask.from_surface(self.image)
        self.animation_name = "off"

    def on(self):
        self.animation_name = "on"
        self.image = self.SPRITES["on"][0]

    def off(self):
        self.animation_name = "off"
        self.image = self.SPRITES["off"][0]

class Finish(Object):
    def __init__(self, x, y, width, height):
        super().__init__(x, y, width, height, "finish")
        finish_image = pygame.Surface((width, height), pygame.SRCALPHA)
        finish_image.fill((0, 255, 0))  # Green finish line
        self.image.blit(finish_image, (0, 0))
        self.mask = pygame.mask.from_surface(self.image)

# Main editor function
def level_editor(window):
    clock = pygame.time.Clock()
    objects = []

    selected_object = "block"

    run = True
    while run:
        clock.tick(FPS)

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                run = False

            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_s:
                    save_level(objects)
                if event.key == pygame.K_1:
                    selected_object = "block"
                if event.key == pygame.K_2:
                    selected_object = "fire"
                if event.key == pygame.K_3:
                    selected_object = "finish"

            if event.type == pygame.MOUSEBUTTONDOWN:
                x, y = pygame.mouse.get_pos()
                x = x - (x % BLOCK_SIZE)
                y = y - (y % BLOCK_SIZE)

                if selected_object == "block":
                    objects.append(Block(x, y, BLOCK_SIZE))
                elif selected_object == "fire":
                    fire = Fire(x, y, 16 * 2, 32 * 2)
                    fire.on()
                    objects.append(fire)
                elif selected_object == "finish":
                    objects.append(Finish(x, y, 50, 50))

        window.fill((255, 255, 255))
        for obj in objects:
            obj.draw(window)

        pygame.display.update()

    pygame.quit()

def save_level(objects):
    with open("level.txt", "w") as f:
        for obj in objects:
            f.write(f"{obj.name} {obj.rect.x} {obj.rect.y} {obj.width} {obj.height}\n")

if __name__ == "__main__":
    level_editor(window)