import streamlit as st
import openai
import tempfile
import os
from pymermaid import mermaid

# 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

# Function to render Mermaid.js diagram using pymermaid
def render_mermaid(diagram_code):
    # Create a temporary file to save the mermaid code
    with tempfile.NamedTemporaryFile(delete=False, suffix=".mmd") as mmd_file:
        mmd_file.write(diagram_code.encode('utf-8'))
        mmd_file.close()

        # Generate diagram image using pymermaid
        output_image_path = mmd_file.name.replace(".mmd", ".png")
        mermaid.render(mmd_file.name, output_image_path)
        
        return output_image_path

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

                # Save diagram as image
                try:
                    image_path = render_mermaid(diagram_code)
                    st.image(image_path, caption="Generated Diagram", use_column_width=True)
                    
                    # Provide download option
                    with open(image_path, "rb") as img_file:
                        st.download_button(
                            label="Download Diagram Image",
                            data=img_file,
                            file_name="diagram.png",
                            mime="image/png"
                        )
                    
                    # Clean up the temporary files
                    os.remove(image_path)
                except Exception as e:
                    st.error(f"Error generating diagram image: {e}")
            else:
                st.error("Failed to generate diagram code.")
        else:
            st.error("Please enter a prompt.")

if __name__ == "__main__":
    main()