import streamlit as st
from rag import rag_chain, retriever

# === Настройка страницы ===
st.set_page_config(
    page_title="Поиск фильмов",
    page_icon="🎬",
    layout="centered"
)

# === Заголовок и подзаголовок ===
st.title("🎬 RAG-поиск по фильмам")
st.caption("Спросите что-нибудь о фильмах — я найду похожие и расскажу!")

# === Инициализация истории сообщений ===
if "messages" not in st.session_state:
    st.session_state.messages = []

# === Отображение истории чата ===
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.write(msg["content"])
        if "posters" in msg and msg["posters"]:
            posters = [url.strip() for url in msg["posters"] if url.strip().startswith("http")]
            cols = st.columns(min(5, len(posters)))
            for i, url in enumerate(posters):
                cols[i % len(cols)].image(url, width=130)

# === Обработка нового запроса ===
if prompt := st.chat_input("Например: «Порекомендуй фильм про космос»"):
    # Добавляем сообщение пользователя
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.write(prompt)

    # Получаем релевантные документы
    with st.spinner("🔍 Ищу подходящие фильмы..."):
        docs = retriever.invoke(prompt)

    # Извлекаем постеры из метаданных
    posters = [
        doc.metadata.get("poster_url", "").strip()
        for doc in docs
        if doc.metadata.get("poster_url", "").strip().startswith("http")
    ]

    # Генерируем ответ через RAG
    with st.spinner("🧠 Формулирую ответ..."):
        response = rag_chain.invoke(prompt)

    # Сохраняем ответ ассистента с постерами
    st.session_state.messages.append({
        "role": "assistant",
        "content": response,
        "posters": posters
    })

    # Отображаем ответ
    with st.chat_message("assistant"):
        st.write(response)
        if posters:
            st.markdown("##### 🎞️ Найденные фильмы:")
            cols = st.columns(min(5, len(posters)))
            for i, url in enumerate(posters):
                cols[i % len(cols)].image(url, width=130)