import streamlit as st
import os
from PIL import Image
import random
from io import BytesIO

# Function to arrange images dynamically
def arrange_images(image_files, canvas_size=(3000, 3000)):
    if not image_files:
        return None

    # Initialize layout
    positions = []  # Keeps track of image positions (x1, y1, x2, y2)
    canvas = Image.new("RGBA", canvas_size, "white")

    def get_center(pos):
        """Calculate center of a bounding box (x1, y1, x2, y2)."""
        return ((pos[0] + pos[2]) // 2, (pos[1] + pos[3]) // 2)

    def does_overlap(new_box, existing_boxes):
        """Check if a new bounding box overlaps any existing boxes."""
        for box in existing_boxes:
            if (
                new_box[0] < box[2]
                and new_box[2] > box[0]
                and new_box[1] < box[3]
                and new_box[3] > box[1]
            ):
                return True
        return False

    # Place the first image at the center of the canvas
    first_img_path = os.path.join(map_dir, image_files[0])
    with Image.open(first_img_path) as img:
        width, height = img.size
        x1 = (canvas_size[0] - width) // 2
        y1 = (canvas_size[1] - height) // 2
        x2, y2 = x1 + width, y1 + height
        positions.append((x1, y1, x2, y2))
        canvas.paste(img, (x1, y1))

    # Place remaining images
    for img_file in image_files[1:]:
        placed = False
        img_path = os.path.join(map_dir, img_file)

        with Image.open(img_path) as img:
            width, height = img.size

            while not placed:
                # Choose a random existing image to connect to
                target_box = random.choice(positions)
                target_center = get_center(target_box)

                # Randomly choose a side of the target image
                side = random.choice(["top", "bottom", "left", "right"])

                # Calculate the new image's position based on the chosen side
                if side == "top":
                    x1 = target_center[0] - width // 2
                    y1 = target_box[1] - height
                elif side == "bottom":
                    x1 = target_center[0] - width // 2
                    y1 = target_box[3]
                elif side == "left":
                    x1 = target_box[0] - width
                    y1 = target_center[1] - height // 2
                elif side == "right":
                    x1 = target_box[2]
                    y1 = target_center[1] - height // 2

                x2, y2 = x1 + width, y1 + height

                # Check for overlap
                if not does_overlap((x1, y1, x2, y2), positions):
                    # Place the image
                    positions.append((x1, y1, x2, y2))
                    canvas.paste(img, (x1, y1))
                    placed = True

    # Save canvas to memory
    buffer = BytesIO()
    canvas.save(buffer, format="PNG")
    buffer.seek(0)
    return buffer, canvas

# Streamlit App
st.title("Dynamic Dungeon Map Generator")
st.write("Automatically generates dungeon maps by arranging PNG images dynamically.")

# Directory for images
map_dir = "."  # Top-level directory

# Canvas size and scroll offsets
canvas_size = 3000
scroll_offset = {"x": 0, "y": 0}

# Initialize session state
if "layout_image" not in st.session_state:
    st.session_state["layout_image"] = None
if "canvas" not in st.session_state:
    st.session_state["canvas"] = None
if "saved" not in st.session_state:
    st.session_state["saved"] = False

# Generate map if not saved and layout_image is empty
if st.session_state["layout_image"] is None or not st.session_state["saved"]:
    image_files = [f for f in os.listdir(map_dir) if f.endswith(".png")]
    if image_files:
        layout_image, canvas = arrange_images(image_files, canvas_size=(canvas_size, canvas_size))
        st.session_state["layout_image"] = layout_image
        st.session_state["canvas"] = canvas

# Display map
if st.session_state["layout_image"] is not None:
    st.image(
        st.session_state["layout_image"],
        caption="Generated Dungeon Map Layout",
        use_column_width=True,
        output_format="PNG",
        clamp=True,
    )

# Save map button
if st.button("💾 Save Map"):
    save_path = "saved_dungeon_map.png"
    st.session_state["canvas"].save(save_path)
    st.session_state["saved"] = True
    st.success(f"Map saved as {save_path}")

# Regenerate map button
if st.button("🗺️ Regenerate Map"):
    image_files = [f for f in os.listdir(map_dir) if f.endswith(".png")]
    if image_files:
        layout_image, canvas = arrange_images(image_files, canvas_size=(canvas_size, canvas_size))
        st.session_state["layout_image"] = layout_image
        st.session_state["canvas"] = canvas
        st.session_state["saved"] = False  # Mark as unsaved
        st.experimental_rerun()

# WASD Controls for scrolling
st.sidebar.title("Scroll Map")
scroll_buttons = st.sidebar.columns(3)
if scroll_buttons[0].button("⬅️"):
    scroll_offset["x"] -= 100
if scroll_buttons[1].button("⬆️"):
    scroll_offset["y"] -= 100
if scroll_buttons[2].button("➡️"):
    scroll_offset["x"] += 100
if st.sidebar.button("⬇️"):
    scroll_offset["y"] += 100

# Sidebar for file upload
st.sidebar.title("Options")
uploaded_file = st.sidebar.file_uploader("Upload a PNG file", type="png")

if uploaded_file:
    # Save uploaded file to the directory
    with open(os.path.join(map_dir, uploaded_file.name), "wb") as f:
        f.write(uploaded_file.getbuffer())
    st.sidebar.success("File uploaded successfully! Refresh the app to include it in the dungeon map.")