import gradio as gr
from sentence_transformers import SentenceTransformer, util
import json
import os

# Загружаем модель
model_name = "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1"
model = SentenceTransformer(model_name)

# Имя файла для сохранения эмбеддингов
embeddings_file = f"movie_embeddings_{model_name.replace('/', '_')}.json"

# Загружаем данные из файла movies.json
try:
    with open("movies.json", "r", encoding="utf-8") as f:
        movies_data = json.load(f)
except FileNotFoundError:
    print("Ошибка: Файл movies.json не найден.")
    movies_data = []

# Проверяем, есть ли сохраненные эмбеддинги
if os.path.exists(embeddings_file):
    with open(embeddings_file, "r", encoding="utf-8") as f:
        movie_embeddings_loaded = json.load(f)
    print("Загружены эмбеддинги из файла.")
else:
    movie_embeddings_loaded = {}

# Создаем словарь с описаниями фильмов и эмбеддингами
movie_descriptions = {}
movie_embeddings = {}

for movie in movies_data:
    title = movie["name"]
    # Формируем строку для эмбеддинга из всех данных фильма
    embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
    movie_descriptions[title] = embedding_string

    # Проверяем, есть ли эмбеддинг для фильма в загруженных
    if title in movie_embeddings_loaded:
        movie_embeddings[title] = movie_embeddings_loaded[title]
    else:
        # Создаем эмбеддинг и добавляем в словарь
        embedding = model.encode(embedding_string, convert_to_tensor=True).tolist()
        movie_embeddings[title] = embedding

# Сохраняем эмбеддинги в файл, если были созданы новые
if len(movie_embeddings_loaded) < len(movie_embeddings):
    with open(embeddings_file, "w", encoding="utf-8") as f:
        json.dump(movie_embeddings, f, ensure_ascii=False, indent=4)
    print("Эмбеддинги сохранены в файл.")

# Конвертируем эмбеддинги в тензоры
if len(movie_embeddings) > 0:
    movie_embeddings_tensor = {
        title: util.pytorch_cos_sim(
            model.encode(query, convert_to_tensor=True),
            model.encode(embedding_string, convert_to_tensor=True)
        )
        for title, embedding_string in movie_descriptions.items()
    }
else:
    movie_embeddings_tensor = None

def search_movies(query, top_k=3):
    """
    Ищет наиболее похожие фильмы по запросу.

    Args:
        query: Текстовый запрос.
        top_k: Количество возвращаемых результатов.

    Returns:
        Строку с результатами поиска в формате HTML.
    """
    if movie_embeddings_tensor is None:
        return "<p>Ошибка: Данные фильмов не загружены.</p>"

    # Сортируем фильмы по убыванию сходства с запросом
    sorted_movies = sorted(
        movie_embeddings_tensor.items(),
        key=lambda item: util.pytorch_cos_sim(
            model.encode(query, convert_to_tensor=True),
            model.encode(movie_descriptions[item[0]], convert_to_tensor=True)
        )[0][0],
        reverse=True
    )

    results_html = ""
    for title, _ in sorted_movies[:top_k]:
        # Ищем полное описание фильма в исходных данных
        for movie in movies_data:
            if movie["name"] == title:
                description = movie["description"]
                year = movie["year"]
                genres = movie["genresList"]
                score = util.pytorch_cos_sim(
                    model.encode(query, convert_to_tensor=True),
                    model.encode(movie_descriptions[title], convert_to_tensor=True)
                )[0][0].item()
                break

        results_html += f"<h3><b>{title} ({year})</b></h3>"
        results_html += f"<p><b>Жанры:</b> {genres}</p>"
        results_html += f"<p><b>Описание:</b> {description}</p>"
        results_html += f"<p><b>Сходство:</b> {score:.4f}</p>"
        results_html += "<hr>"

    return results_html

# Создаем интерфейс Gradio
iface = gr.Interface(
    fn=search_movies,
    inputs=gr.Textbox(label="Введите запрос:"),
    outputs=gr.HTML(label="Результаты поиска:"),
    title="Поиск фильмов по описанию",
    description="Введите запрос, и система найдет наиболее похожие фильмы по их описаниям.",
    examples=[
        ["Фильм про ограбление"],
        ["Комедия 2019 года"],
        ["Фантастика про космос"],
    ],
)

# Запускаем приложение
iface.launch()