import streamlit as st
import networkx as nx
import matplotlib.pyplot as plt
import openai
import os

# Set OpenAI API key from Hugging Face secrets
openai.api_key = os.getenv("OPENAI_API_KEY")

st.set_page_config(page_title="AI Mind Map Generator", layout="wide")
st.title("🧠 AI Mind Map Generator")
st.write("Enter a topic and generate a mind map using GPT-3.5!")

topic = st.text_input("Enter a topic to generate mind map:", "")

def generate_mindmap(topic):
    prompt = f"Create a simple, structured mind map for the topic '{topic}'. Use bullet points to show sub-topics and indentation for hierarchy."
    response = openai.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.7
    )
    return response.choices[0].message.content

def build_graph_from_text(text):
    G = nx.Graph()
    parents = []
    prev_indent = -1

    for line in text.splitlines():
        if not line.strip():
            continue

        indent = len(line) - len(line.lstrip(" "))
        node = line.strip("-• ").strip()

        if indent == 0:
            parents = [node]
        else:
            level = indent // 2
            if level < len(parents):
                parents = parents[:level]
            parents.append(node)

        G.add_node(node)
        if len(parents) > 1:
            G.add_edge(parents[-2], node)

    return G

if topic:
    with st.spinner("Generating mind map..."):
        mindmap_text = generate_mindmap(topic)
        G = build_graph_from_text(mindmap_text)

        fig, ax = plt.subplots(figsize=(10, 6))
        pos = nx.spring_layout(G, seed=42)
        nx.draw(G, pos, with_labels=True, node_color="skyblue", edge_color="gray", node_size=2000, font_size=10)
        st.pyplot(fig)

        with st.expander("🔍 View raw mind map text"):
            st.text(mindmap_text)