import streamlit as st
import openai

# Function to call OpenAI GPT model for prompt processing
def get_diagram_code(prompt, diagram_type, api_key):
    try:
        # Set the OpenAI API key dynamically based on user input
        openai.api_key = api_key

        response = openai.Completion.create(
            engine="gpt-4",
            prompt=f"Generate a {diagram_type} diagram in Mermaid.js syntax based on the following prompt: {prompt}",
            max_tokens=500
        )
        return response.choices[0].text.strip()
    except Exception as e:
        st.error(f"Error: {e}")
        return None

# Streamlit App UI
def main():
    st.title("Generate Diagrams using GPT-4 and Mermaid.js")

    # User input for OpenAI API key
    api_key = st.text_input("Enter your OpenAI API Key:", type="password")
    
    # Check if API key is provided
    if not api_key:
        st.warning("Please enter your OpenAI API Key to continue.")
        return

    # User prompt for diagram type
    prompt = st.text_area("Enter your prompt for the diagram:", "")
    
    diagram_types = ["UML Diagram", "ER Diagram", "State Diagram", "Class Diagram", "Sequence Diagram"]
    diagram_choice = st.selectbox("Select the type of diagram to generate:", diagram_types)

    if st.button("Generate Diagram"):
        if prompt:
            diagram_code = get_diagram_code(prompt, diagram_choice, api_key)
            
            if diagram_code:
                # Render Mermaid code to Streamlit
                st.code(diagram_code, language='mermaid')

                # Display Mermaid diagram using iframe (mermaid.live)
                diagram_url = f"https://mermaid-js.github.io/mermaid-live-editor/#/edit/{diagram_code}"
                st.markdown(f'<iframe src="{diagram_url}" width="800" height="600"></iframe>', unsafe_allow_html=True)

            else:
                st.error("Failed to generate diagram code.")
        else:
            st.error("Please enter a prompt.")

if __name__ == "__main__":
    main()