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

# Adjust Streamlit layout to wide mode
st.set_page_config(layout="wide")

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

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

    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))
        room_details.append(f"Room 1: {image_files[0]} at center")

    # Place remaining images
    for idx, img_file in enumerate(image_files[1:], start=2):
        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:
                target_box = random.choice(positions)
                target_center = get_center(target_box)
                side = random.choice(["top", "bottom", "left", "right"])

                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

                if not does_overlap((x1, y1, x2, y2), positions):
                    positions.append((x1, y1, x2, y2))
                    canvas.paste(img, (x1, y1))
                    room_details.append(f"Room {idx}: {img_file} at ({x1}, {y1})")
                    placed = True

    buffer = BytesIO()
    canvas.save(buffer, format="PNG")
    buffer.seek(0)
    return buffer, canvas, room_details

# Function to create a base64 link
def create_base64_download_link(file_content, filename, file_type):
    b64 = base64.b64encode(file_content.encode() if file_type == "txt" else file_content).decode()
    return f"[📥 Download {filename}](data:application/{file_type};base64,{b64})"

# Sidebar Title
st.sidebar.markdown("#### 🏰 Dynamic Dungeon Map Generator")

# Directory for images
map_dir = "."
canvas_size = 3000

# 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 "room_details" not in st.session_state:
    st.session_state["room_details"] = []
if "saved_files" not in st.session_state:
    st.session_state["saved_files"] = []

# Generate map if layout_image is empty
if st.session_state["layout_image"] is None:
    image_files = [f for f in os.listdir(map_dir) if f.endswith(".png")]
    if image_files:
        layout_image, canvas, room_details = 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["room_details"] = room_details

# Sidebar Controls
if st.sidebar.button("💾 Save Map"):
    now = datetime.datetime.now()
    filename_png = f"dungeon_map_{now.strftime('%Y%m%d_%H%M%S')}.png"
    filename_txt = f"room_layout_{now.strftime('%Y%m%d_%H%M%S')}.txt"

    # Save the PNG map
    st.session_state["canvas"].save(filename_png)
    st.sidebar.success(f"Map saved as {filename_png}")

    # Save room details as a text file
    room_details_content = "\n".join(st.session_state["room_details"])
    with open(filename_txt, "w") as f:
        f.write(room_details_content)
    st.sidebar.success(f"Room layout saved as {filename_txt}")

    # Generate base64 download links
    with open(filename_png, "rb") as f:
        png_base64 = create_base64_download_link(f.read(), filename_png, "octet-stream")
    txt_base64 = create_base64_download_link(room_details_content, filename_txt, "txt")

    # Add files to save history
    st.session_state["saved_files"].append((png_base64, txt_base64))

if st.sidebar.button("🗺️ Regenerate Map"):
    image_files = [f for f in os.listdir(map_dir) if f.endswith(".png")]
    if image_files:
        layout_image, canvas, room_details = 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["room_details"] = room_details
        st.rerun()

# Display save history in the sidebar
st.sidebar.markdown("### 📜 Save History")
for idx, (png_link, txt_link) in enumerate(st.session_state["saved_files"], start=1):
    st.sidebar.markdown(f"{idx}. {png_link} | {txt_link}")

# Zoom Controls
st.sidebar.title("Zoom")
show_zoomed = st.sidebar.checkbox("Show Zoomed Version")
zoom_level = st.sidebar.slider("Zoom Level", min_value=0.1, max_value=2.0, value=1.0, step=0.1)

if show_zoomed and st.session_state["canvas"] is not None:
    zoomed_canvas = st.session_state["canvas"].resize(
        (int(canvas_size * zoom_level), int(canvas_size * zoom_level)),
        resample=Image.Resampling.LANCZOS,
    )
    buffer = BytesIO()
    zoomed_canvas.save(buffer, format="PNG")
    buffer.seek(0)
    st.image(
        buffer,
        caption=f"Zoomed Dungeon Map Layout (Zoom Level: {zoom_level}x)",
        use_container_width=False,
        output_format="PNG",
    )
else:
    st.image(
        st.session_state["layout_image"],
        caption="Generated Dungeon Map Layout",
        use_container_width=True,
        output_format="PNG",
    )