app.py

import streamlit as st
import pathlib
from PIL import Image
import google.generativeai as genai
import os # Import os for environment variables

# --- Configuration ---
try:
    # Ensure GOOGLE_API_KEY is set in Streamlit secrets or as an environment variable
    API_KEY = st.secrets["GOOGLE_API_KEY"]
except KeyError:
    st.error("Google API Key not found. Please set it in Streamlit secrets or as an environment variable (`GOOGLE_API_KEY`).")
    st.stop() # Stop the app if API key is missing

genai.configure(api_key=API_KEY)

# Generation configuration for the Gemini model
GENERATION_CONFIG = {
    "temperature": 1,
    "top_p": 0.95,
    "top_k": 64,
    "max_output_tokens": 8192,
    "response_mime_type": "text/plain",
}

# Safety settings for the Gemini model
SAFETY_SETTINGS = [
    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
]

# Model name constant
MODEL_NAME = "gemini-1.5-pro-latest"

# --- Model Initialization ---
@st.cache_resource # Cache the model to avoid re-initializing on every rerun
def load_gemini_model():
    """Loads and caches the Google GenerativeModel."""
    return genai.GenerativeModel(
        model_name=MODEL_NAME,
        safety_settings=SAFETY_SETTINGS,
        generation_config=GENERATION_CONFIG,
    )

model = load_gemini_model()

# Initialize chat session in Streamlit's session state
# This ensures the chat history persists across reruns for a single user
if "chat_session" not in st.session_state:
    st.session_state.chat_session = model.start_chat(history=[])

# --- Helper Function for Model Communication ---
def send_message_to_model(message: str, image_path: pathlib.Path) -> str:
    """Sends a message and an image to the Gemini model and returns the response."""
    image_input = {
        'mime_type': 'image/jpeg',
        'data': image_path.read_bytes()
    }
    try:
        response = st.session_state.chat_session.send_message([message, image_input])
        return response.text
    except Exception as e:
        st.error(f"Error communicating with the Gemini model: {e}")
        st.exception(e) # Display full traceback for debugging
        return "An error occurred during AI model communication. Please try again or check your API key."

# --- Streamlit App ---
def main():
    """Main function to run the Streamlit application."""
    st.set_page_config(page_title="Gemini 1.5 Pro: Images to Code", layout="wide") # Set a wider layout
    st.title("Gemini 1.5 Pro: Images to Code 👨‍💻")
    st.markdown('Made with ❤️ by [KhulnaSoft](https://x.com/khulnasoft)')

    st.info("Upload an image of a UI design, and I'll generate the corresponding HTML and CSS code for you!")

    # Framework selection using a selectbox
    framework_options = {
        "Regular CSS (Flexbox/Grid)": "Regular CSS use flex grid etc",
        "Bootstrap": "Bootstrap",
        "Tailwind CSS": "Tailwind CSS",
        "Materialize CSS": "Materialize CSS"
    }
    selected_framework_name = st.selectbox(
        "Choose your preferred CSS framework:",
        options=list(framework_options.keys()),
        help="This will influence the CSS generated within your HTML file."
    )
    framework = framework_options[selected_framework_name]

    uploaded_file = st.file_uploader("Upload a UI image (JPG, JPEG, PNG):", type=["jpg", "jpeg", "png"])

    # temp_image_path is declared outside the try block to ensure it's accessible for cleanup
    temp_image_path = pathlib.Path("temp_image.jpg")

    if uploaded_file is not None:
        try:
            # Load and display the image
            image = Image.open(uploaded_file)
            st.image(image, caption='Uploaded UI Image', use_column_width=True)

            # Convert image to RGB mode if it has an alpha channel
            if image.mode == 'RGBA':
                image = image.convert('RGB')

            # Save the uploaded image temporarily
            image.save(temp_image_path, format="JPEG")

            st.markdown("---") # Visual separator

            # Button to trigger the generation process
            if st.button("Generate UI Code", help="Click to initiate the multi-step code generation."):
                st.subheader("Code Generation Process:")

                # Step 1: Generate initial UI description
                with st.spinner("Step 1/4: Describing your UI elements and colors..."):
                    prompt = "Describe this UI in accurate details. When you reference a UI element put its name and bounding box in the format: [object name (y_min, x_min, y_max, x_max)]. Also Describe the color of the elements."
                    description = send_message_to_model(prompt, temp_image_path)
                    st.success("UI Description Generated!")
                    with st.expander("See Initial UI Description"):
                        st.text(description)

                # Step 2: Refine the description
                with st.spinner("Step 2/4: Refining description with visual comparison..."):
                    refine_prompt = f"Compare the described UI elements with the provided image and identify any missing elements or inaccuracies. Also Describe the color of the elements. Provide a refined and accurate description of the UI elements based on this comparison. Here is the initial description: {description}"
                    refined_description = send_message_to_model(refine_prompt, temp_image_path)
                    st.success("UI Description Refined!")
                    with st.expander("See Refined UI Description"):
                        st.text(refined_description)

                # Step 3: Generate initial HTML
                with st.spinner("Step 3/4: Generating initial HTML with CSS..."):
                    html_prompt = (
                        f"Create an HTML file based on the following UI description, using the UI elements described in the previous response. "
                        f"Include {framework} CSS within the HTML file to style the elements. "
                        f"Make sure the colors used are the same as the original UI. "
                        f"The UI needs to be responsive and mobile-first, matching the original UI as closely as possible. "
                        f"Do not include any explanations or comments. Avoid using ```html and ``` at the end. "
                        f"ONLY return the HTML code with inline CSS. Here is the refined description: {refined_description}"
                    )
                    initial_html = send_message_to_model(html_prompt, temp_image_path)
                    st.success("Initial HTML Generated!")
                    st.subheader("Initial Generated HTML:")
                    st.code(initial_html, language='html')

                # Step 4: Refine HTML
                with st.spinner("Step 4/4: Refining the generated HTML code..."):
                    refine_html_prompt = (
                        f"Validate the following HTML code based on the UI description and image and provide a refined version of the HTML code with {framework} CSS that improves accuracy, responsiveness, and adherence to the original design. "
                        f"ONLY return the refined HTML code with inline CSS. Avoid using ```html and ``` at the end. "
                        f"Here is the initial HTML: {initial_html}"
                    )
                    refined_html = send_message_to_model(refine_html_prompt, temp_image_path)
                    st.success("HTML Refined Successfully!")
                    st.subheader("Refined Generated HTML:")
                    st.code(refined_html, language='html')

                st.markdown("---") # Final separator
                st.success("All steps completed! Your `index.html` file is ready for download.")

                # Save the refined HTML to a file and provide download link
                with open("index.html", "w", encoding="utf-8") as file: # Specify encoding
                    file.write(refined_html)

                st.download_button(
                    label="Download index.html",
                    data=refined_html,
                    file_name="index.html",
                    mime="text/html"
                )
                st.info("You can open the downloaded `index.html` file in your web browser to view the generated UI.")

        except Exception as e:
            st.error(f"An unexpected error occurred: {e}")
            st.exception(e) # Displays the full traceback

        finally:
            # Clean up the temporary image file whether an error occurred or not
            if temp_image_path.exists():
                os.remove(temp_image_path)
                # st.success("Temporary image file removed.") # Can uncomment for debugging
    else:
        st.write("Please upload an image to start generating UI code.")

if __name__ == "__main__":
    main()