# Make a virtual environment with compulsory GPU!!
# python -m venv .venv
# Activate the virtual environment
# Windows: .venv\Scripts\activate
# Linux/Mac: source .venv/bin/activate

# Install packages
# pip install -r requirements.txt

import streamlit as st
from generate_caption import generate_caption
from PIL import Image
import io

# Set page config
st.set_page_config(
    page_title="AI Image Caption & Prompt Generator",
    page_icon="🖼️",
    layout="wide"
)

# Title and description
st.title("🖼️ AI Image Caption & Prompt Generator")

# Create two columns for layout
col1, col2 = st.columns([1, 1])

with col1:
    st.header("📤 Upload & Configure")

    # File uploader
    uploaded_file = st.file_uploader(
        "Choose an image file",
        type=['png', 'jpg', 'jpeg', 'gif', 'bmp', 'tiff', 'webp'],
        help="Supported formats: PNG, JPG, JPEG, GIF, BMP, TIFF, WebP"
    )

    # Caption options
    st.subheader("Caption Options")

    caption_type = st.selectbox(
        "Caption Type",
        options=["MidJourney", "Descriptive", "Training Prompt"],
        index=0,
        help="Choose the style of caption you want to generate"
    )

    caption_length = st.selectbox(
        "Caption Length",
        options=["short", "any", "long"],
        index=0,
        help="Select the desired length of the caption"
    )

    # Generate button
    generate_btn = st.button("🎯 Generate Caption", type="primary", use_container_width=True)

with col2:
    st.header("Preview & Results")

    if uploaded_file is not None:
        # Display uploaded image
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image", use_container_width=True)

        # Generate caption when button is clicked
        if generate_btn:
            with st.spinner("Generating caption... This may take a moment."):
                try:
                    # Generate caption
                    prompt_used, caption = generate_caption(
                        image,
                        caption_type=caption_type,
                        caption_length=caption_length,
                        extra_options=None
                    )

                    # Display results
                    st.success("Caption generated successfully!")

                    # Caption result
                    st.subheader("📝 Generated Caption")
                    st.write(f"{caption}")

                    # Copy to clipboard button
                    st.code(caption, language=None)

                    # Additional info
                    with st.expander("ℹ️ Generation Details"):
                        st.write(f"**Caption Type:** {caption_type}")
                        st.write(f"**Caption Length:** {caption_length}")
                        if prompt_used:
                            st.write(f"**Prompt Used:** {prompt_used}")

                except Exception as e:
                    st.error(f"Error generating caption: {str(e)}")
                    st.info("Please make sure you have installed all required dependencies and your GPU is properly configured.")

    else:
        st.markdown("""
        ### How to use:
        1. Upload an image using the file uploader
        2. Select your preferred caption type and length
        3. Click 'Generate Caption' to create your AI caption

        ### Caption Types:
        - **MidJourney**: Optimized for AI art generation prompts
        - **Descriptive**: Detailed description of the image content
        - **Training Prompt**: Formatted for AI model training
        """)

# Sidebar with additional information
with st.sidebar:
    st.header("🔧 System Requirements")
    st.markdown("""
    **Required Setup:**
    - GPU-enabled environment
    - Virtual environment activated
    - All dependencies installed via `requirements.txt`

    **Supported Image Formats:**
    - PNG, JPG, JPEG
    - GIF, BMP, TIFF, WebP
    """)

    st.header("💡 Tips")
    st.markdown("""
    - Higher quality images produce better captions
    - Different caption types serve different purposes
    - Short captions are more focused, long ones more detailed
    """)

    st.header("⚙️ Setup Instructions")
    with st.expander("Click to view setup commands"):
        st.code("""
# Create virtual environment
python -m venv .venv

# Activate virtual environment
# Windows:
.venv\\Scripts\\activate
# Linux/Mac:
source .venv/bin/activate

# Install dependencies
pip install -r requirements.txt

# Run Streamlit app
streamlit run app.py
        """, language="bash")

# Footer
st.markdown("Built by [Aditya Singh]")