import streamlit as st
import google.generativeai as genai
import os
from dotenv import load_dotenv
from graphviz import Digraph

load_dotenv()

genai.configure(api_key=("AIzaSyDziGvuT1woHnH4_S3L_zQZV55Yj-123A8"))

# gemini function for generating mind maps and diagrams
def get_gemini_response(input):
    model = genai.GenerativeModel('gemini-pro')
    response = model.generate_content(input)
    return response

# Function to generate mind map from content
def generate_mind_map(content):
    mind_map_prompt = f"""
    ### Generate a mind map from the following educational content:
    ### Content:
    {content}
    ### Mind Map:
    """
    response = get_gemini_response(mind_map_prompt)
    return response

# Function to parse mind map data
def parse_mind_map_data(data):
    lines = data.split('\n')
    mind_map_structure = {}
    current_parent = None
    current_sub_parent = None
    
    for line in lines:
        stripped_line = line.strip()
        if stripped_line.startswith('**') and stripped_line.endswith('**'):
            current_parent = stripped_line.strip('**').strip()
            mind_map_structure[current_parent] = {}
        elif stripped_line.startswith('- '):
            current_sub_parent = stripped_line.strip('- ').strip()
            if current_parent:
                mind_map_structure[current_parent][current_sub_parent] = []
        elif stripped_line.startswith('   - '):
            sub_child = stripped_line.strip('   - ').strip()
            if current_parent and current_sub_parent:
                mind_map_structure[current_parent][current_sub_parent].append(sub_child)
    
    return mind_map_structure

# Function to display mind map using Graphviz
def display_mind_map(mind_map_response):
    st.subheader("Generated Mind Map")
    if mind_map_response and mind_map_response.candidates:
        mind_map_data = mind_map_response.candidates[0].content.parts[0].text if mind_map_response.candidates[0].content.parts else ""
        st.write("Mind Map Data:", mind_map_data)  # Debugging output

        # Parse the mind map data into a hierarchical structure
        mind_map_structure = parse_mind_map_data(mind_map_data)
        
        # Initialize a new Digraph
        dot = Digraph()

        # Add nodes and edges to the graph based on the structure
        for parent, children in mind_map_structure.items():
            dot.node(parent, parent)
            for child, sub_children in children.items():
                dot.node(child, child)
                dot.edge(parent, child)
                for sub_child in sub_children:
                    dot.node(sub_child, sub_child)
                    dot.edge(child, sub_child)

        # Render the graph
        st.graphviz_chart(dot.source)
    else:
        st.write("No response received from the model.")

## Streamlit App
st.title("AI-Powered Educational Content Visualization")
st.text("Paste your engineering educational content to generate a mind map and diagram.")

content_input = st.text_area("Paste your educational content here", height=300)
submit_content = st.button('Generate Mind Map and Diagram')

if submit_content:
    if content_input:
        mind_map_response = generate_mind_map(content_input)
        
        # Display the generated mind map
        display_mind_map(mind_map_response)