import streamlit as st
from PIL import Image

# Custom CSS for enhanced styling
st.markdown("""
    <style>
        body {
            background-color: #f7f9fc;
            font-family: 'Arial', sans-serif;
        }
        .main {
            background-color: #ffffff;
            padding: 30px;
            border-radius: 15px;
            box-shadow: 0px 0px 15px rgba(0, 0, 0, 0.1);
        }
        .title {
            color: #2c3e50;
            font-size: 40px;
            font-weight: 700;
            text-align: center;
            margin-bottom: 40px;
        }
        .upload-area {
            background-color: #e8f4f8;
            padding: 20px;
            border-radius: 10px;
            border: 2px dashed #1abc9c;
            text-align: center;
            margin-bottom: 20px;
        }
        .upload-area:hover {
            background-color: #d1ecf1;
        }
        .generate-button {
            background-color: #1abc9c;
            color: white;
            border-radius: 5px;
            padding: 10px;
            font-weight: bold;
            cursor: pointer;
            text-align: center;
            margin-top: 20px;
        }
        .generate-button:hover {
            background-color: #16a085;
        }
        .result-section {
            margin-top: 40px;
        }
        .result-title {
            color: #34495e;
            font-size: 24px;
            font-weight: 600;
            margin-bottom: 20px;
        }
        .success {
            background-color: #dff0d8;
            color: #3c763d;
            padding: 15px;
            border-radius: 5px;
        }
    </style>
""", unsafe_allow_html=True)

# Main UI container
st.markdown('<div class="main">', unsafe_allow_html=True)
st.markdown('<div class="title">Florence-2 Image Captioning Demo</div>', unsafe_allow_html=True)

# Task prompt input
task_prompt = st.text_input("Task Prompt", value="Describe the image in detail:")

# Image upload area
st.markdown('<div class="upload-area">', unsafe_allow_html=True)
uploaded_image = st.file_uploader("Upload your image here", type=["jpg", "jpeg", "png"])
st.markdown('</div>', unsafe_allow_html=True)

# Additional text input (optional)
text_input = st.text_area("Additional Text Input (Optional)", height=150)

# Display the uploaded image and process it if available
if uploaded_image is not None:
    image = Image.open(uploaded_image)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Generate Caption button
    if st.button("Generate Caption", key="generate"):
        # Assuming `run_example` function is defined
        result = run_example(task_prompt, image, text_input)

        # Display the generated caption
        st.markdown('<div class="result-section">', unsafe_allow_html=True)
        st.markdown('<div class="result-title">Generated Caption</div>', unsafe_allow_html=True)
        st.markdown(f'<div class="success">{result["text"]}</div>', unsafe_allow_html=True)

        # Display bounding boxes or polygons if available
        if "bboxes" in result:
            st.markdown("### Detected Objects")
            fig = plot_bbox(image, result)
            st.pyplot(fig)

        if "polygons" in result:
            st.markdown("### Image with Polygons")
            processed_image = draw_polygons(image.copy(), result)
            st.image(processed_image, caption="Image with Polygons", use_column_width=True)

st.markdown('</div>', unsafe_allow_html=True)