import streamlit as st
import google.generativeai as genai
from PIL import Image
import os

# Configure GenAI with API Key
genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))
model = genai.GenerativeModel('gemini-1.5-pro-latest')

# Input prompt for model
input_prompt = """
As an HTML and CSS expert, your task is to create complete HTML and CSS code based on the provided screenshot, ensuring clear and functional markup. Provide a main.html file with inline CSS that replicates the exact color and style as shown in the given screenshot.
Output structure:
Start and end with (```).
"""

# Function to generate response from model
def generate_response(input_prompt, image):
    response = model.generate_content([input_prompt, image[0]])
    return response.text

# Function to process uploaded image
def input_image_setup(uploaded_file):
    if uploaded_file is not None:
        bytes_data = uploaded_file.getvalue()
        image_parts = [
            {
                "mime_type": uploaded_file.type,
                "data": bytes_data
            }
        ]
        return image_parts
    else:
        raise FileNotFoundError("No file uploaded")

# Streamlit App Interface
st.title("Screenshot to HTML Code 📃")
st.text("Upload your demo webpage image here:")

upload_file = st.file_uploader('', type=['jpg', 'jpeg', 'png'])
if upload_file is not None:
    image = Image.open(upload_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)

submit = st.button('Create a Webpage')
if submit:
    image_data = input_image_setup(uploaded_file=upload_file)
    with st.spinner("Building the Webpage..."):
        response = generate_response(input_prompt, image_data)
        
        # Extract code from the response
        code = response.strip("```")
        
        # Display the code
        st.subheader("Generated Code:")
        st.code(code, language='html')
        
    # Render the HTML Output with larger dimensions
    if st.button("Render This Code"):
        st.subheader("Rendered HTML Output:")
        st.components.v1.html(code, height=1000, width=1500, scrolling=True)