import os
import streamlit as st
import pathlib
from PIL import Image
import google.generativeai as genai
from dotenv import load_dotenv
load_dotenv()

# Configure the API key directly in the script
API_KEY = os.environ.get("GOOGLE_API_KEY")
genai.configure(api_key=API_KEY)

# Generation configuration
generation_config = {
    "temperature": 0.8,
    "top_p": 0.95,
    "top_k": 64,
    "max_output_tokens": 50000,
    "response_mime_type": "text/plain",
}

# Safety settings
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
MODEL_NAME = "gemini-1.5-flash-latest"

# Framework selection (e.g., Tailwind, Bootstrap, etc.)
framework = "Regular CSS use flex grid etc"  # Change this to "Bootstrap" or any other framework as needed

# Create the model
model = genai.GenerativeModel(
    model_name=MODEL_NAME,
    safety_settings=safety_settings,
    generation_config=generation_config,
)

# Start a chat session
chat_session = model.start_chat(history=[])

# Function to send a message to the model
def send_message_to_model(message, image_path):
    image_input = {
        'mime_type': 'image/jpeg',
        'data': pathlib.Path(image_path).read_bytes()
    }
    response = chat_session.send_message([message, image_input])
    return response.text

# Streamlit app
def main():
    st.title("Image to HTML Code AI 👨‍💻 ")
    st.subheader('Made by [Pushpendra](https://github.com/pushparmar).')

    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        try:
            if uploaded_file.size > 5 * 1024 * 1024:  # Limit file size to 5MB
                st.error("File size exceeds the limit of 5MB.")
                return
            if uploaded_file.type not in ["image/jpeg", "image/png"]:
                st.error("Invalid file type. Please upload a JPEG or PNG image.")
                return

            # Load and display the image
            image = Image.open(uploaded_file)
            st.image(image, caption='Uploaded Image.', use_container_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
            temp_image_path = pathlib.Path("temp_image.jpg")
            image.save(temp_image_path, format="JPEG")
            styleframework = st.selectbox("Select CSS Framework", ["Regular CSS use flex grid etc", "Bootstrap (Disabled)", "Tailwind (Disabled)"])
            uiFrameWork = st.selectbox("Select UI Framework", ["Regular HTML Code", "React.js (Disabled)", "Angular.js (Disabled)", "Vue.js (Disabled)"])

            # Generate UI description
            if st.button("Code UI"):
                placeholder = st.empty()
                placeholder.write("🧑‍💻 A parmar code agent is Looking at your UI...")

               
                import time
                time.sleep(4)
                st.write("🧑‍💻 A parmar code agent Generating Code for your image")
                import time
                time.sleep(2)
                detailed_prompt = f"""
                Analyze the uploaded image and generate a detailed {uiFrameWork} file with inline CSS that accurately represents the UI. 
                - Identify all UI elements (e.g., buttons, text fields, images, etc.) and their bounding boxes in the format: [object name (y_min, x_min, y_max, x_max)].
                - Describe the color, font, and styling of each element.
                - Ensure the {uiFrameWork} code is responsive and mobile-first, using {styleframework} CSS for layout and styling.
                - If {uiFrameWork} is selected, generate components with selected framework syntax and include inline CSS for styling.
                - Include detailed inline CSS for each element to match the original design as closely as possible.
                - Avoid using external libraries or frameworks unless specified.
                - Do not include explanations or comments in the code.
                - ONLY return the {uiFrameWork} code with inline CSS.
                """

                if "Disabled" in uiFrameWork or "Disabled" in styleframework:
                    st.warning("This option is currently unavailable.")
                    return
                else:
                    st.write(f"You selected: {uiFrameWork}")
                
                    generated_html = send_message_to_model(detailed_prompt, temp_image_path)
                    

                    # Save the refined HTML to a file
                    with open("index.html", "w") as file:
                        file.write(generated_html)
                    st.success("HTML file 'index.html' has been created.")

                    # Provide download link for HTML
                    st.download_button(label="Download HTML", data=generated_html, file_name="index.html", mime="text/html")

                    # Provide a link to open the HTML file in a new tab
                    st.markdown(f'<a href="index.html" target="_blank">Open HTML in a new tab</a>', unsafe_allow_html=True)
        except Exception as e:
            st.error(f"An error occurred: {e}")

if __name__ == "__main__":
    main()