import os
import streamlit as st
import requests
import re
from dotenv import load_dotenv
import tempfile

# Load environment variables (for FLOWISE_API token)
load_dotenv()

# Define API settings
API_URL = "https://nakheeltech.com:8030/api/v1/prediction/c1681ef1-8f47-4004-b4ab-594fbbd3eb3f"
headers = {"Authorization": f"Bearer {os.getenv('FLOWISE_API')}"}

# Function to send a query to the API
def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()

# Extract only Python code from text
def extract_python_code(text):
    code_pattern = r"```(?:python)?(.*?)```"
    match = re.search(code_pattern, text, re.DOTALL)
    if match:
        return match.group(1).strip()
    return "No Python code found in response."

# Streamlit UI
st.title("Architecture Diagram Generator")

# Get input from the user
user_input = st.text_area("Describe your architecture:", 
                          placeholder="Enter a description of your architecture here...")

if st.button("Generate Diagram"):
    if user_input:
        # Send the user's input to the API
        payload = {"question":"make sure to import from diagrams import Diagram" + user_input + " architecture"}
        output = query(payload)
        print(output)
        
        # Check and display the Python code from the API response
        if "text" in output:
            python_code = extract_python_code(output["text"])
            st.code(python_code, language="python")

            # Save the generated code to a temporary file and execute it
            with tempfile.TemporaryDirectory() as tmpdirname:
                temp_file = os.path.join(tmpdirname, "architecture_diagram.py")
                with open(temp_file, "w") as f:
                    f.write(python_code)

                # Change directory to temp and execute the code
                original_dir = os.getcwd()
                os.chdir(tmpdirname)
                try:
                    exec(open(temp_file).read())
                    diagram_path = os.path.join(tmpdirname, "diagram.jpg")

                    # Display the generated image
                    if os.path.exists(diagram_path):
                        st.image(diagram_path, caption="Generated Architecture Diagram")
                    else:
                        st.error("No diagram image was generated.")
                except Exception as e:
                    st.error(f"Error executing the Python code: {e}")
                finally:
                    os.chdir(original_dir)  # Restore the original directory
        else:
            st.write("Response:", output)
    else:
        st.warning("Please enter an architecture description.")