import streamlit as st
import time
from chatbot import generate_html_css_from_image

def main():
    st.set_page_config(page_title="Gemini 2.5 Flash HTML/CSS Chatbot", layout="wide")

    if "output_html" not in st.session_state:
        st.session_state.output_html = None
    if "last_image" not in st.session_state:
        st.session_state.last_image = None

    st.markdown("""
        <div style='text-align: center; padding: 10px 0;'>
            <h1>Welcome, Gemini 2.5 Flash HTML/CSS Chatbot</h1>
        </div>
        <div style='text-align: center;'>
            <p style='font-size: 18px; background-color: #e6f0ff; padding: 10px; border-radius: 8px; display: inline-block;'>
                📁 Please upload an image of a website UI
            </p>
        </div>
    """, unsafe_allow_html=True)

    left, right = st.columns([1, 5])

    with left:
        st.markdown(
            """
            <div style='
                border: 2px dashed #cccccc;
                border-radius: 10px;
                padding: 20px;
                text-align: center;
                background-color: #f9f9f9;
                width: 100%;
            '>
                <h3 style="margin-top: 0;">📁 Upload Image</h3>
                <p>Drag and drop or browse</p>
            </div>
            """,
            unsafe_allow_html=True
        )

        uploaded_image = st.file_uploader(
            label="Upload a website UI screenshot",
            type=["jpg", "jpeg", "png"],
            label_visibility="collapsed"
        )

        if uploaded_image:
            st.image(uploaded_image, caption="Uploaded UI Image", use_column_width=True)

    with right:
        if uploaded_image:
            current_name = getattr(uploaded_image, "name", "uploaded_image")

            if st.session_state.last_image != current_name or st.session_state.output_html is None:
                with st.spinner("Generating HTML + CSS using Gemini 2.5 Flash..."):
                    try:
                        output = generate_html_css_from_image(uploaded_image)
                        st.session_state.output_html = output
                        st.session_state.last_image = current_name
                    except Exception as e:
                        st.error(f"Error: {e}")
                        return

            output = st.session_state.get("output_html")
            if output:
                st.subheader("💬 Generated HTML + CSS:")
                placeholder = st.empty()
                typed_text = ""
                for char in output:
                    typed_text += char
                    placeholder.code(typed_text, language="html")
                    time.sleep(0.002)

                st.download_button(
                    "💾 Download HTML file",
                    data=output,
                    file_name="generated_ui.html",
                    mime="text/html"
                )
            else:
                st.error("No output generated.")
        else:
            st.session_state.output_html = None
            st.session_state.last_image = None

if __name__ == "__main__":
    main()