app.py

import os
import time
import threading
import queue
from typing import List, Dict, Any, Optional
import logging
from urllib.parse import urlparse

import gradio as gr
import torch
import psycopg2
import zlib
import numpy as np
from sentence_transformers import SentenceTransformer, util
from sklearn.preprocessing import normalize

# Рекомендуется использовать python-dotenv для загрузки переменных окружения
# from dotenv import load_dotenv
# load_dotenv()

# Настройка логирования
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Настройки базы данных PostgreSQL
DATABASE_URL = os.environ.get("DATABASE_URL")
if DATABASE_URL is None:
    raise ValueError("DATABASE_URL environment variable not set.")

parsed_url = urlparse(DATABASE_URL)
db_params = {
    "host": parsed_url.hostname,
    "port": parsed_url.port,
    "database": parsed_url.path.lstrip("/"),
    "user": parsed_url.username,
    "password": parsed_url.password,
    "sslmode": "require"
}

# Загружаем модель
model_name = "BAAI/bge-m3"
logging.info(f"Загрузка модели {model_name}...")
model = SentenceTransformer(model_name)
logging.info("Модель загружена успешно.")

# Имена таблиц
embeddings_table = "movie_embeddings"
query_cache_table = "query_cache"

# Максимальный размер таблицы кэша запросов в байтах (50MB)
MAX_CACHE_SIZE = 50 * 1024 * 1024

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

# Очередь для необработанных фильмов
movies_queue: queue.Queue = queue.Queue()

# Флаги состояния
processing_complete = False
search_in_progress = False

# Блокировка для доступа к базе данных
db_lock = threading.Lock()

# Размер пакета для обработки эмбеддингов
batch_size = 32

def get_db_connection():
    """Устанавливает соединение с базой данных."""
    try:
        conn = psycopg2.connect(**db_params)
        return conn
    except Exception as e:
        logging.error(f"Ошибка подключения к базе данных: {e}")
        return None

def setup_database():
    """Настраивает базу данных: создает расширение, таблицы и индексы."""
    with get_db_connection() as conn:
        if conn is None:
            return
        try:
            with conn.cursor() as cur:
                # Создаем расширение pgvector если его нет
                cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
                
                # Удаляем существующие таблицы если они есть
                # cur.execute(f"DROP TABLE IF EXISTS {embeddings_table}, {query_cache_table};")
                
                # Создаем таблицу для хранения эмбеддингов фильмов
                cur.execute(f"""
                CREATE TABLE {embeddings_table} (
                    movie_id INTEGER PRIMARY KEY,
                    embedding_crc32 BIGINT,
                    string_crc32 BIGINT,
                    model_name TEXT,
                    embedding vector(1024)
                );
                CREATE INDEX ON {embeddings_table} (string_crc32);
                """)
                
                # Создаем таблицу для кэширования запросов
                cur.execute(f"""
                CREATE TABLE {query_cache_table} (
                    query_crc32 BIGINT PRIMARY KEY,
                    query TEXT,
                    model_name TEXT,
                    embedding vector(1024),
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
                );
                CREATE INDEX ON {query_cache_table} (query_crc32);
                CREATE INDEX ON {query_cache_table} (created_at);
                """)
            
            conn.commit()
            logging.info("База данных успешно настроена.")
        except Exception as e:
            logging.error(f"Ошибка при настройке базы данных: {e}")
            conn.rollback()

# Настраиваем базу данных при запуске
setup_database()

def calculate_crc32(text: str) -> int:
    """Вычисляет CRC32 для строки."""
    return zlib.crc32(text.encode('utf-8')) & 0xFFFFFFFF

def encode_string(text: str) -> np.ndarray:
    """Кодирует строку в эмбеддинг."""
    embedding = model.encode(text, convert_to_tensor=True, normalize_embeddings=True)
    return embedding.cpu().numpy()

def get_movies_without_embeddings() -> List[Dict[str, Any]]:
    """Получает список фильмов, для которых нужно создать эмбеддинги."""
    with get_db_connection() as conn:
        if conn is None:
            return []
        try:
            with conn.cursor() as cur:
                # Получаем список ID фильмов, которые уже есть в базе
                cur.execute(f"SELECT movie_id FROM {embeddings_table}")
                existing_ids = {row[0] for row in cur.fetchall()}
                
                # Фильтруем только те фильмы, которых нет в базе
                movies_to_process = [movie for movie in movies_data if movie['id'] not in existing_ids]
                
            logging.info(f"Найдено {len(movies_to_process)} фильмов для обработки.")
            return movies_to_process
        except Exception as e:
            logging.error(f"Ошибка при получении списка фильмов для обработки: {e}")
            return []

def get_embedding_from_db(conn, table_name: str, crc32_column: str, crc32_value: int, model_name: str) -> Optional[np.ndarray]:
    """Получает эмбеддинг из базы данных."""
    try:
        with conn.cursor() as cur:
            cur.execute(f"SELECT embedding FROM {table_name} WHERE {crc32_column} = %s AND model_name = %s", (crc32_value, model_name))
            result = cur.fetchone()
            if result and result[0]:
                # Нормализуем эмбеддинг после извлечения из БД
                return normalize(np.array(result[0]).reshape(1, -1))[0]
    except Exception as e:
        logging.error(f"Ошибка при получении эмбеддинга из БД: {e}")
    return None

def insert_embedding(conn, table_name: str, movie_id: int, embedding_crc32: int, string_crc32: int, embedding: np.ndarray) -> bool:
    """Вставляет эмбеддинг в базу данных."""
    try:
        # Нормализуем эмбеддинг перед сохранением
        normalized_embedding = normalize(embedding.reshape(1, -1))[0]
        with conn.cursor() as cur:
            cur.execute(f"""
            INSERT INTO {table_name} (movie_id, embedding_crc32, string_crc32, model_name, embedding)
            VALUES (%s, %s, %s, %s, %s)
            ON CONFLICT (movie_id) DO NOTHING
            """, (movie_id, embedding_crc32, string_crc32, model_name, normalized_embedding.tolist()))
        conn.commit()
        return True
    except Exception as e:
        logging.error(f"Ошибка при вставке эмбеддинга: {e}")
        conn.rollback()
        return False

def process_movies():
    """Обрабатывает фильмы, создавая для них эмбеддинги."""
    global processing_complete
    logging.info("Начало обработки фильмов.")
    
    # Получаем список фильмов, которые нужно обработать
    movies_to_process = get_movies_without_embeddings()
    if not movies_to_process:
        logging.info("Все фильмы уже обработаны.")
        processing_complete = True
        return

    # Добавляем фильмы в очередь
    for movie in movies_to_process:
        movies_queue.put(movie)

    with get_db_connection() as conn:
        if conn is None:
            processing_complete = True
            return
        try:
            while not movies_queue.empty():
                if search_in_progress:
                    time.sleep(1)
                    continue

                batch = []
                while not movies_queue.empty() and len(batch) < batch_size:
                    try:
                        movie = movies_queue.get_nowait()
                        batch.append(movie)
                    except queue.Empty:
                        break

                if not batch:
                    break

                logging.info(f"Обработка пакета из {len(batch)} фильмов...")
                for movie in batch:
                    embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
                    string_crc32 = calculate_crc32(embedding_string)

                    # Проверяем существующий эмбеддинг
                    existing_embedding = get_embedding_from_db(conn, embeddings_table, "string_crc32", string_crc32, model_name)
                    if existing_embedding is None:
                        embedding = encode_string(embedding_string)
                        embedding_crc32 = calculate_crc32(str(embedding.tolist()))
                        if insert_embedding(conn, embeddings_table, movie['id'], embedding_crc32, string_crc32, embedding):
                            logging.info(f"Сохранен эмбеддинг для '{movie['name']}'")
                        else:
                            logging.error(f"Ошибка сохранения эмбеддинга для '{movie['name']}'")
                    else:
                        logging.info(f"Эмбеддинг для '{movie['name']}' уже существует")

        except Exception as e:
            logging.error(f"Ошибка при обработке фильмов: {e}")
        finally:
            processing_complete = True
            logging.info("Обработка фильмов завершена")

def get_movie_embeddings(conn) -> Dict[str, np.ndarray]:
    """Загружает все эмбеддинги фильмов из базы данных."""
    movie_embeddings = {}
    try:
        with conn.cursor() as cur:
            cur.execute(f"SELECT movie_id, embedding FROM {embeddings_table}")
            for movie_id, embedding in cur.fetchall():
                # Находим название фильма по ID
                for movie in movies_data:
                    if movie['id'] == movie_id:
                        movie_embeddings[movie['name']] = normalize(np.array(embedding).reshape(1, -1))[0]
                        break
        logging.info(f"Загружено {len(movie_embeddings)} эмбеддингов фильмов.")
    except Exception as e:
        logging.error(f"Ошибка при загрузке эмбеддингов фильмов: {e}")
    return movie_embeddings

def clean_query_cache(conn):
    """Очищает устаревшие записи из кэша запросов."""
    try:
        with conn.cursor() as cur:
            # Получаем общий размер кэша
            cur.execute(f"SELECT pg_total_relation_size('{query_cache_table}')")
            total_size = cur.fetchone()[0]
            
            if total_size > MAX_CACHE_SIZE:
                # Удаляем старые записи, пока размер не станет меньше максимального
                cur.execute(f"""
                DELETE FROM {query_cache_table}
                WHERE ctid IN (
                    SELECT ctid
                    FROM {query_cache_table}
                    ORDER BY created_at ASC
                    LIMIT (SELECT COUNT(*) / 2 FROM {query_cache_table})
                )
                """)
                conn.commit()
                logging.info("Кэш запросов очищен.")
    except Exception as e:
        logging.error(f"Ошибка при очистке кэша запросов: {e}")
        conn.rollback()

def search_movies(query: str, top_k: int = 10) -> List[Dict[str, Any]]:
    """Выполняет поиск фильмов по запросу."""
    global search_in_progress
    search_in_progress = True
    
    try:
        with get_db_connection() as conn:
            if conn is None:
                return []

            clean_query_cache(conn)

            query_crc32 = calculate_crc32(query)
            query_embedding = get_embedding_from_db(conn, query_cache_table, "query_crc32", query_crc32, model_name)

            if query_embedding is None:
                query_embedding = encode_string(query)
                insert_embedding(conn, query_cache_table, -1, -1, query_crc32, query_embedding)

            movie_embeddings = get_movie_embeddings(conn)

            # Вычисляем косинусное сходство
            similarities = util.cos_sim(query_embedding, list(movie_embeddings.values()))[0]

            # Сортируем результаты
            top_results = sorted(zip(similarities, movie_embeddings.keys()), key=lambda x: x[0], reverse=True)[:top_k]

            results = []
            for score, movie_name in top_results:
                movie = next((m for m in movies_data if m['name'] == movie_name), None)
                if movie:
                    results.append({
                        "name": movie['name'],
                        "year": movie['year'],
                        "genres": movie['genresList'],
                        "description": movie['description'],
                        "score": float(score)
                    })

        return results
    except Exception as e:
        logging.error(f"Ошибка при поиске фильмов: {e}")
        return []
    finally:
        search_in_progress = False

# Запускаем обработку фильмов в отдельном потоке
threading.Thread(target=process_movies, daemon=True).start()

# Создаем интерфейс Gradio
def gradio_search(query: str) -> str:
    results = search_movies(query)
    output = ""
    for movie in results:
        output += f"Название: {movie['name']} ({movie['year']})\n"
        output += f"Жанры: {', '.join(movie['genres'])}\n"
        output += f"Описание: {movie['description']}\n"
        output += f"Релевантность: {movie['score']:.2f}\n\n"
    return output

iface = gr.Interface(
    fn=gradio_search,
    inputs="text",
    outputs="text",
    title="Поиск фильмов",
    description="Введите запрос для поиска фильмов"
)

iface.launch()