import streamlit as st
from PIL import Image
import time
from chatbot import chat_with_image

USER_AVATAR = "👤"
BOT_AVATAR = '<img src="https://img.icons8.com/emoji/48/robot-emoji.png" width="20"/>'

st.set_page_config(page_title="Gemini 2.5 Image Chatbot", layout="wide")

st.sidebar.header("🌐 Language Settings")
selected_language = st.sidebar.selectbox(
    "Select response language",
    [
        "Afrikaans", "Albanian", "Amharic", "Arabic", "Armenian", "Assamese", "Azerbaijani",
        "Basque", "Belarusian", "Bengali", "Bosnian", "Bulgarian", "Burmese",
        "Catalan", "Cebuano", "Chinese", "Corsican", "Croatian", "Czech",
        "Danish", "Dutch", "English", "Esperanto", "Estonian",
        "Finnish", "French", "Frisian",
        "Galician", "Georgian", "German", "Greek", "Gujarati",
        "Haitian Creole", "Hausa", "Hawaiian", "Hebrew", "Hindi", "Hmong", "Hungarian",
        "Icelandic", "Igbo", "Indonesian", "Irish", "Italian",
        "Japanese", "Javanese", "Kannada", "Kazakh", "Khmer", "Kinyarwanda", "Korean", "Kurdish",
        "Kyrgyz", "Lao", "Latin", "Latvian", "Lithuanian", "Luxembourgish",
        "Macedonian", "Malagasy", "Malay", "Malayalam", "Maltese", "Maori", "Marathi", "Mongolian",
        "Nepali", "Norwegian", "Nyanja", "Odia", "Pashto", "Persian", "Polish", "Portuguese",
        "Punjabi", "Quechua", "Romanian", "Russian", "Samoan", "Scots Gaelic", "Serbian", "Sesotho",
        "Shona", "Sindhi", "Sinhala", "Slovak", "Slovenian", "Somali", "Spanish", "Sundanese",
        "Swahili", "Swedish", "Tagalog", "Tajik", "Tamil", "Tatar", "Telugu", "Thai",
        "Tigrinya", "Turkish", "Turkmen", "Ukrainian", "Urdu", "Uyghur", "Uzbek", "Vietnamese",
        "Welsh", "Xhosa", "Yiddish", "Yoruba", "Zulu"
    ],
    index=21
)

st.markdown("<h1 style='text-align: center;'>Welcome, Gemini 2.5 Flash Image Chatbot 2.0</h1>", unsafe_allow_html=True)

if "history" not in st.session_state:
    st.session_state.history = []
if "images" not in st.session_state:
    st.session_state.images = []
if "current_image_index" not in st.session_state:
    st.session_state.current_image_index = None
if "last_animated_index" not in st.session_state:
    st.session_state.last_animated_index = None

left_col, right_col = st.columns([1, 5])

with left_col:
    st.header("📁 Upload Image")
    uploaded_file = st.file_uploader(
        "Drop an image here",
        type=["jpg", "jpeg", "png"],
        accept_multiple_files=False,
        label_visibility="collapsed"
    )

    if uploaded_file:
        image = Image.open(uploaded_file).convert("RGB")
        st.session_state.images = [image]
        st.session_state.current_image_index = 0
        st.image(image, caption="Uploaded Image")
    else:
        st.session_state.images = []
        st.session_state.current_image_index = None

with right_col:
    if st.session_state.images and st.session_state.current_image_index is not None:
        current_image = st.session_state.images[st.session_state.current_image_index]

        st.markdown("### Chat With Image:")
        st.markdown(f"#### Language : `{selected_language}`")

        history = st.session_state.history

        latest_index = None
        for idx in reversed(range(len(history))):
            msg = history[idx]
            if (
                msg["sender"] == "Gemini" and 
                msg["image_index"] == st.session_state.current_image_index and 
                idx != st.session_state.last_animated_index
            ):
                latest_index = idx
                break

        for i, msg in enumerate(history):
            if msg["image_index"] != st.session_state.current_image_index:
                continue

            if msg["sender"] == "You":
                st.markdown(f"""
                <div style="background-color: #f1f1f1; padding: 10px; border-radius: 8px; margin-bottom: 10px;">
                <strong>{USER_AVATAR} :</strong> {msg['message']}
                </div>
                """, unsafe_allow_html=True)

            elif msg["sender"] == "Gemini":
                if i == latest_index:
                    st.session_state.last_animated_index = i  
                    placeholder = st.empty()
                    full_response = msg["message"]
                    for j in range(1, len(full_response) + 1):
                        placeholder.markdown(f"{BOT_AVATAR} : {full_response[:j]}", unsafe_allow_html=True)
                        time.sleep(0.002)
                else:
                    st.markdown(f"{BOT_AVATAR} : {msg['message']}", unsafe_allow_html=True)

        user_prompt = st.chat_input("Ask About Image")
        if user_prompt:
            full_prompt = f"{user_prompt}\n\nPlease respond in {selected_language}."
            try:
                gemini_response = chat_with_image(full_prompt, current_image)
            except Exception as e:
                gemini_response = f"Gemini failed to respond: {e}"

            st.session_state.history.append({
                "sender": "You",
                "message": user_prompt,
                "image_index": st.session_state.current_image_index
            })
            st.session_state.history.append({
                "sender": "Gemini",
                "message": gemini_response,
                "image_index": st.session_state.current_image_index
            })

            st.rerun()
    else:
        st.info("Please upload an image to start chatting.")