import streamlit as st
import openai
import urllib.parse
import io
from PIL import Image
import time  # Used for simulating progress updates

# 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-4o",  # 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=2048
        )

        return response['choices'][0]['message']['content'].strip()
    
    except Exception as e:
        st.error(f"Error: {e}")
        return None

# Adding custom CSS
def custom_css():
    st.markdown("""
        <style>
            body {
                background-color: #f0f4f8;
                font-family: 'Roboto', sans-serif;
            }
            .stButton>button {
                background-color: #4CAF50;
                color: white;
                font-size: 18px;
                padding: 12px 24px;
                border-radius: 8px;
                transition: background-color 0.3s ease;
            }
            .stButton>button:hover {
                background-color: #f1c232;
            }
            .stTextInput, .stTextArea {
                background-color: #fff;
                border: 2px solid #ddd;
                border-radius: 8px;
                padding: 12px;
                font-size: 16px;
                transition: border 0.3s ease;
            }
            .stTextInput:focus, .stTextArea:focus {
                border-color: #4CAF50;
            }
            .stImage {
                margin-bottom: 20px;
                border-radius: 10px;
            }
            h1 {
                color: #333;
                text-align: center;
                font-size: 36px;
            }
            .expandable-section {
                margin-top: 20px;
            }
            iframe {
                border: 0;
                border-radius: 10px;
            }
            .loading-spinner {
                display: none;
                animation: spin 2s infinite linear;
                border: 6px solid #f3f3f3;
                border-top: 6px solid #3498db;
                border-radius: 50%;
                width: 50px;
                height: 50px;
                margin: 0 auto;
            }
            @keyframes spin {
                0% { transform: rotate(0deg); }
                100% { transform: rotate(360deg); }
            }
        </style>
    """, unsafe_allow_html=True)

# Add anime.js CDN and custom JS for animations
def anime_js():
    st.markdown("""
        <script src="https://cdn.jsdelivr.net/npm/animejs@3.2.1/lib/anime.min.js"></script>
        <script>
            document.addEventListener('DOMContentLoaded', function() {
                // Example animation for the button when clicked
                const button = document.querySelector('.stButton>button');
                if (button) {
                    button.addEventListener('mouseover', function() {
                        anime({
                            targets: button,
                            scale: [1, 1.1],
                            duration: 300,
                            easing: 'easeInOutQuad'
                        });
                    });
                    button.addEventListener('mouseout', function() {
                        anime({
                            targets: button,
                            scale: [1.1, 1],
                            duration: 300,
                            easing: 'easeInOutQuad'
                        });
                    });
                }
                
                // Example animation for diagram generation code
                const diagramCode = document.querySelector('pre code');
                if (diagramCode) {
                    anime({
                        targets: diagramCode,
                        opacity: [0, 1],
                        translateY: [-50, 0],
                        duration: 1000,
                        easing: 'easeOutExpo'
                    });
                }
            });
        </script>
    """, unsafe_allow_html=True)

# Streamlit App UI
def main():
    custom_css()  # Apply custom CSS
    anime_js()  # Add anime.js functionality
    
    st.title("Generate Diagrams using AI and MermaidFlow")
    
    # Display a logo or icon with subtle animation
    image = Image.open("11.png")  # Path to your image file
    st.image(image, width=300) 

    # Display sample prompt
    with st.expander("Sample Prompt For Users"):
        st.markdown("""
            Create a UML diagram for a **Library Management System** that includes classes such as **Book** (with attributes like `bookID`, `title`, `author`, `publisher`, `genre`, `availabilityStatus` and methods like `checkAvailability()`, `updateAvailability()`), **Member** (with attributes such as `memberID`, `name`, `email`, `phoneNumber`, `membershipDate` and methods like `borrowBook()`, `returnBook()`, `reserveBook()`), **Staff** (with attributes such as `staffID`, `name`, `role`, `email`, `phoneNumber` and methods like `addBook()`, `removeBook()`, `updateBookInfo()`), **Transaction** (with attributes like `transactionID`, `transactionDate`, `dueDate`, `fineAmount` and methods like `calculateFine()`, `generateReceipt()`), and **Reservation** (with attributes like `reservationID`, `reservationDate`, `expirationDate` and methods like `cancelReservation()`, `checkReservationStatus()`). Define the relationships where a **Member** can borrow many **Books**, a **Staff** can manage many **Books**, a **Transaction** is linked to one **Book** and one **Member**, and a **Reservation** is made by a **Member** for a **Book**. Ensure appropriate visibility markers for methods and attributes, and include necessary relationships such as associations, multiplicities, and inheritance. Optionally, consider adding sequence or activity diagrams for processes like book borrowing.
        """)

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

    # Create a progress bar
    progress_bar = st.progress(0)

    # Instruction Expander with step-by-step guide
    with st.expander("Instruction to Generate Diagram"):
        st.markdown("""
                        **Copy and paste the Generated Mermaid Code** in the Live Editor and download the generated diagram by clicking on **Export as SVG** format to get a high-resolution generated image.
        """)

    if st.button("Generate Diagram"):
        if prompt:
            # Simulate the progress
            for i in range(100):
                progress_bar.progress(i + 1)
                time.sleep(0.05)  # Simulate processing time for loading bar

            diagram_code = get_diagram_code(prompt, diagram_choice, api_key)
            
            if diagram_code:
                # Hide the progress bar once the diagram is generated
                progress_bar.empty()

                # Render Mermaid code to Streamlit with animation
                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)

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

if __name__ == "__main__":
    main()