# Основные библиотеки
import os
import re
import string
import warnings
import numpy as np
import pandas as pd
import torch
# Машинное обучение и обработка текста
from transformers import AutoTokenizer, AutoModel, AutoModelForSeq2SeqLM, MarianMTModel, MarianTokenizer
from sentence_transformers import SentenceTransformer, util
from sklearn.base import BaseEstimator, TransformerMixin
# FAISS для семантического поиска
import faiss
# Лемматизация и морфология
import pymorphy3
# Streamlit для создания веб-приложений
import streamlit as st
# Кастомные модули
from EmbeddingGenerator import EmbeddingGenerator
from TextAugmentation import TextAugmentation
# LangChain для интеграции GigaChat
from langchain_community.chat_models.gigachat import GigaChat

# ======= загружаем ранее рассчитанные эмбеддинги и объявляем классы=======

# Инициализация GigaChat с ключом и отключенной проверкой SSL
giga = GigaChat(
    credentials="ODk0NDE1ODEtYTJhMi00N2Y1LTk4YWItNGZlNzNkM2QwMDNiOjk5YmVjN2ZjLThmM2EtNDhjYy04OWQ2LWNkOTlhOTNkNGY3NQ==", 
    verify_ssl_certs=False
)

augmentor = TextAugmentation()
embedding_gen = EmbeddingGenerator()
df=pd.read_csv("movies_data_fixed.csv")
image_path = "image-2.png"














# Загружаем и отображаем картинку
st.image(image_path, use_container_width=True)

# Заголовок
st.markdown(
    """
    <div class="title">
        КиноКринж
    </div>
    """,
    unsafe_allow_html=True
)

# Добавляем окно ввода текста
user_input = st.text_area("Добавьте описание фильма", "", height=150)

# Слайдер для выбора количества фильмов
num_results = st.slider('Выберите количество фильмов', min_value=1, max_value=20, value=4)

# Выбор модели
model_option = st.selectbox('Выберите модель для обработки запроса:', ['cointegrated/rubert-tiny2','DeepPavlov/rubert-base-cased','all-MiniLM-L6-v2', 'paraphrase-MiniLM-L6-v2'])

if model_option!='DeepPavlov/rubert-base-cased':
    model = SentenceTransformer(model_option)


# ======= дополнительная фильтрация для аугментаций (убираем слишком непохожие) =======
def filter_paraphrases(original, paraphrases, threshold=0.8):
    original_embedding = model.encode(original)
    filtered = []
    for paraphrase in paraphrases:
        paraphrase_embedding = model.encode(paraphrase)
        similarity = util.cos_sim(original_embedding, paraphrase_embedding).item()
        if similarity >= threshold:
            filtered.append(paraphrase)
    return filtered
#======================СЕМПЛ======= =======










# Проверка наличия рекомендованных фильмов
if 'recommended_movies' not in st.session_state:
    st.session_state.recommended_movies = []

# Кнопка для поиска
if st.button('Найти фильм'):
    if user_input.strip():
        # Генерация эмбеддинга для запроса
        if model_option != 'DeepPavlov/rubert-base-cased' and model_option != 'cointegrated/rubert-tiny2':
            index = faiss.read_index('faiss_index.bin')
            query_embedding = model.encode([user_input]).astype("float32")
            faiss.normalize_L2(query_embedding)
        elif model_option == 'DeepPavlov/rubert-base-cased':
            index = faiss.read_index('pavlov3.bin')
            back_translate = augmentor.back_translate(user_input)
            augmented_query_pavlov = user_input + " " + back_translate
            query_embedding = embedding_gen.generate_embeddings(augmented_query_pavlov, method="pavlov")
        elif model_option == 'cointegrated/rubert-tiny2':
            index = faiss.read_index('rubert2.bin')
            paraphrase = augmentor.paraphrase(user_input, num_return_sequences=3)
            filtered_rubert = filter_paraphrases(user_input, paraphrase)
            augmented_query_rubert = user_input + " " + " ".join(filtered_rubert)
            query_embedding = embedding_gen.generate_embeddings(augmented_query_rubert, method="rubert_tiny2").reshape(1, -1)
            faiss.normalize_L2(query_embedding)

        # Поиск ближайших соседей
        distances, indices = index.search(query_embedding, num_results)

        # Отображение результатов
        st.write(f"Результаты поиска ({num_results} фильмов):")
        recommended_movies = []
        for idx, distance in zip(indices[0], distances[0]):
            recommended_movies.append({
                'title': df.iloc[idx]['movie_title'],
                'description': df.iloc[idx]['description'],
                'image_url': df.iloc[idx]['image_url'],
                'page_url': df.iloc[idx]['page_url'],
                'similarity': distance,
                'short_description': None,  # Содержимое краткого описания
                'is_short_description_shown': False  # Флаг для того, чтобы избежать повторного запроса
            })

        # Сохраняем результаты в session_state
        st.session_state.recommended_movies = recommended_movies

# Отображение рекомендованных фильмов
for idx, movie in enumerate(st.session_state.recommended_movies):
    st.write(f"### {movie['title']}")
    st.write(f"Описание: {movie['description']}")
    st.write(f"Схожесть: {movie['similarity']:.4f}")

    # Отображаем картинку постера
    if movie.get('image_url'):
        st.image(movie['image_url'], width=200)

    # Добавляем ссылку на страницу фильма
    if movie.get('page_url'):
        st.markdown(f"[Перейти на страницу фильма]({movie['page_url']})")

    # Генерируем уникальный ключ с использованием индекса
    button_key = f"short_description_button_{idx}"  # Уникальный ключ для кнопки
    if st.button(f"Получить краткое содержание для {movie['title']}", key=button_key):
        if not movie.get('is_short_description_shown', False):  # Проверяем состояние
            try:
                # Отправляем запрос в GigaChat
                prompt = f"{movie['title']} краткое содержание фильма не более 100 слов"
                response = giga.invoke(prompt)

                # Извлекаем описание из ответа
                description = response.content if response else "Описание не найдено."
                movie['short_description'] = description
                movie['is_short_description_shown'] = True

            except Exception as e:
                st.error(f"Ошибка при запросе в GigaChat: {e}")

    # Показываем краткое содержание
    if movie.get('short_description') and movie.get('is_short_description_shown', False):
        st.write(f"Краткое содержание для {movie['title']}: {movie['short_description']}")

    st.write("---")