import streamlit as st
from PIL import Image
import io
import base64
import time
from datetime import datetime

# --------------------------------------------------
# Configuration
# --------------------------------------------------
st.set_page_config(page_title="🤖 Chat IA - Analyseur d'Images", page_icon="🖼️", layout="wide")

# --------------------------------------------------
# CSS
# --------------------------------------------------
st.markdown("""
<style>
.message-user {background: linear-gradient(135deg, #4ade80, #22d3ee); color: white; padding: 10px; border-radius: 15px; margin: 10px 0; margin-left: 20%;}
.message-ai {background: #f8fafc; padding: 10px; border-radius: 15px; margin: 10px 0; margin-right: 20%; border-left: 4px solid #667eea;}
.uploaded-image {max-width: 100%; border-radius: 10px; margin-top: 5px;}
</style>
""", unsafe_allow_html=True)

# --------------------------------------------------
# State init
# --------------------------------------------------
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
if 'pending_image' not in st.session_state:
    st.session_state.pending_image = None
if 'captioner' not in st.session_state:
    st.session_state.captioner = None
if 'model_loaded' not in st.session_state:
    st.session_state.model_loaded = False

# --------------------------------------------------
# Model loader
# --------------------------------------------------
@st.cache_resource
def _load_pipeline():
    from transformers import pipeline
    return pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")

def get_captioner():
    if not st.session_state.model_loaded or st.session_state.captioner is None:
        st.session_state.captioner = _load_pipeline()
        st.session_state.model_loaded = True
    return st.session_state.captioner

# --------------------------------------------------
# Utils
# --------------------------------------------------

def add_message(sender, content, image=None):
    st.session_state.chat_history.append({
        'sender': sender,
        'content': content,
        'image': image,
        'timestamp': datetime.now().strftime("%H:%M")
    })

def display_chat():
    for m in st.session_state.chat_history:
        img_html = f'<img src="data:image/png;base64,{m["image"]}" class="uploaded-image"/>' if m.get('image') else ''
        if m['sender'] == 'user':
            st.markdown(f"<div class='message-user'>{m['content']}{img_html}<div style='font-size:0.8em;opacity:0.6'>{m['timestamp']}</div></div>", unsafe_allow_html=True)
        else:
            st.markdown(f"<div class='message-ai'>{m['content']}{img_html}<div style='font-size:0.8em;opacity:0.6'>{m['timestamp']}</div></div>", unsafe_allow_html=True)

def analyze_image(pil_image):
    captioner = get_captioner()
    result = captioner(pil_image)[0]['generated_text']
    return result

# --------------------------------------------------
# UI
# --------------------------------------------------

st.markdown("## 🤖 Assistant IA - Analyseur d'Images")

# Chat display
chat_container = st.container()
with chat_container:
    display_chat()

# Upload image (stored in pending_image until send)
uploaded_file = st.file_uploader("Uploader une image à envoyer avec votre message", type=["png", "jpg", "jpeg"], key="uploader_image")
if uploaded_file:
    image = Image.open(uploaded_file).convert("RGB")
    buffer = io.BytesIO()
    image.save(buffer, format='PNG')
    img_base64 = base64.b64encode(buffer.getvalue()).decode()
    st.session_state.pending_image = {'pil': image, 'base64': img_base64, 'name': uploaded_file.name}

# Text input + send
col1, col2 = st.columns([4, 1])
with col1:
    user_message = st.text_area("Votre message", key="user_input", height=80)
with col2:
    st.markdown("<br>", unsafe_allow_html=True)
    if st.button("📤 Envoyer", use_container_width=True):
        if user_message.strip() or st.session_state.pending_image:
            # Add user message
            if st.session_state.pending_image:
                add_message('user', user_message.strip() or f"🖼️ {st.session_state.pending_image['name']}", image=st.session_state.pending_image['base64'])
                # Analyze image
                analysis = analyze_image(st.session_state.pending_image['pil'])
                add_message('ai', f"🔍 **Analyse de l'image :** {analysis}")
                st.session_state.pending_image = None
            else:
                add_message('user', user_message.strip())
                add_message('ai', "(Pas d'image à analyser)")
            st.session_state.user_input = ""
            st.experimental_rerun()

# Clear history
if st.button("🗑️ Effacer l'historique"):
    st.session_state.chat_history = []
    st.session_state.pending_image = None
    st.experimental_rerun()