import streamlit as st
import os
import random
from PIL import Image
import networkx as nx
import matplotlib.pyplot as plt

# Function to calculate aspect ratio
def calculate_aspect_ratio(image_path):
    with Image.open(image_path) as img:
        width, height = img.size
        gcd = lambda a, b: a if not b else gcd(b, a % b)
        divisor = gcd(width, height)
        return width // divisor, height // divisor

# Function to classify room types based on aspect ratio
def classify_room(aspect_ratio):
    if aspect_ratio == (16, 9) or aspect_ratio == (19, 6):
        return "Hallway"
    elif aspect_ratio[0] < 8 and aspect_ratio[1] < 8:
        return "Door"
    else:
        return "Room"

# Function to generate a layout
def generate_dungeon_map(image_files):
    dungeon = {"Hallways": [], "Doors": [], "Rooms": []}

    for img_file in image_files:
        img_path = os.path.join(map_dir, img_file)
        aspect_ratio = calculate_aspect_ratio(img_path)
        room_type = classify_room(aspect_ratio)

        if room_type == "Hallway":
            dungeon["Hallways"].append(img_file)
        elif room_type == "Door":
            dungeon["Doors"].append(img_file)
        else:
            dungeon["Rooms"].append(img_file)

    # Generate connections between elements
    layout = nx.Graph()
    for i, hallway in enumerate(dungeon["Hallways"]):
        layout.add_node(f"Hallway {i+1}", type="Hallway")
        if dungeon["Rooms"]:
            room = dungeon["Rooms"].pop(0)
            layout.add_node(room, type="Room")
            layout.add_edge(f"Hallway {i+1}", room)

        if dungeon["Doors"]:
            door = dungeon["Doors"].pop(0)
            layout.add_node(door, type="Door")
            layout.add_edge(f"Hallway {i+1}", door)

    return layout, dungeon

# Function to visualize the dungeon layout
def visualize_dungeon(layout, output_path="dungeon_tree.png"):
    plt.figure(figsize=(12, 8))
    pos = nx.spring_layout(layout)  # Positions for nodes
    types = nx.get_node_attributes(layout, 'type')

    # Draw nodes with different colors
    color_map = {
        "Hallway": "blue",
        "Door": "green",
        "Room": "orange"
    }
    colors = [color_map[types[node]] for node in layout.nodes]

    nx.draw(
        layout, pos, with_labels=True, node_color=colors, node_size=2000,
        font_size=10, font_color="white", font_weight="bold", edge_color="gray"
    )
    plt.savefig(output_path)
    plt.close()

# Streamlit App
st.title("Dynamic Dungeon Map Generator")
st.write("Automatically generates dungeon maps by analyzing PNG images in a directory.")

# Directory for images
map_dir = "maps"  # Replace with the actual directory path

# Scan the directory for .png files
if os.path.exists(map_dir):
    image_files = [f for f in os.listdir(map_dir) if f.endswith(".png")]

    if image_files:
        st.subheader("Generate Dungeon Map")

        # Add emoji buttons
        if st.button("🔄 Regenerate Dungeon"):
            layout, dungeon = generate_dungeon_map(image_files)
            visualize_dungeon(layout)

            st.image("dungeon_tree.png", caption="Dungeon Layout Tree")

            # Display classified rooms
            st.subheader("Classified Rooms")
            st.write("### Hallways")
            st.write(", ".join(dungeon["Hallways"]) if dungeon["Hallways"] else "No hallways found.")

            st.write("### Doors")
            st.write(", ".join(dungeon["Doors"]) if dungeon["Doors"] else "No doors found.")

            st.write("### Rooms")
            st.write(", ".join(dungeon["Rooms"]) if dungeon["Rooms"] else "No rooms found.")
        else:
            st.write("Click the button above to generate a dungeon map.")
    else:
        st.write("No PNG files found in the specified directory.")
else:
    st.write("The specified directory does not exist. Please check the path.")

# 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.")