import streamlit as st
import openai
import urllib.parse

# 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

        # For GPT-4 and other chat models, use the v1/chat/completions endpoint
        response = openai.ChatCompletion.create(
            model="gpt-4",  # Make sure you're using the right model
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": f"Generate a {diagram_type} diagram in Mermaid.js syntax based on the following prompt: {prompt}"}
            ],
            max_tokens=500
        )

        return response['choices'][0]['message']['content'].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 MermaidFlow")

    # 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')

                # Encode the Mermaid code for use in the MermaidFlow editor URL
                encoded_code = urllib.parse.quote(diagram_code)  # URL encode the Mermaid code

                # MermaidFlow Editor URL with the encoded Mermaid code
                mermaidflow_url = f"https://www.mermaidflow.app/editor?code={encoded_code}"

                # Embed MermaidFlow Editor in iframe
                st.markdown(f'<iframe src="{mermaidflow_url}" width="100%" height="600px"></iframe>', unsafe_allow_html=True)

                # Provide a direct link to MermaidFlow editor for convenience
                st.markdown(f"Click the link to open and edit the diagram directly: [MermaidFlow Editor](https://www.mermaidflow.app/editor?code={encoded_code})")

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

if __name__ == "__main__":
    main()