app/data/files_management.py

import os
import json
import hashlib
import logging
from openai import OpenAI
from dotenv import load_dotenv

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

load_dotenv()
client = OpenAI(api_key=os.getenv("API_KEY_openai"))

# ==== Управление этапами через флаги ====
RUN_CREATE_VECTOR_STORE = False
RUN_UPLOAD_FILES = False
RUN_GENERATE_TEMPLATE = True
RUN_UPDATE_ATTRIBUTES = False
RUN_DELETE_FILE = False

# ==== Параметры ====
VECTOR_STORE_NAME = "Logoped_KB"
FILES_DIR = "materials/vector_store"
ATTRIBUTES_FILE = "attributes_template.json"
VS_ID = os.getenv("VECTOR_STORE_ID")
FILE_NAME_DELETE = ""



# 1. Создание векторного хранилища
def create_vector_store(name: str) -> str:
    """Создает новое векторное хранилище, если его еще нет, и возвращает его ID"""

    # Получаем список всех векторных хранилищ
    existing_stores = client.vector_stores.list()

    # Проверяем, существует ли хранилище с таким же именем
    for store in existing_stores:
        if store.name == name:
            logger.info(f"Хранилище с именем '{name}' уже существует, его ID: {store.id}")
            return store.id  # Возвращаем существующий ID хранилища

    # Если хранилище не найдено, создаем новое
    vector_store = client.vector_stores.create(name=name)
    logger.info(f"Создано хранилище: {vector_store.id}")
    return vector_store.id


# 2. Загрузка файлов с проверкой дубликатов
def upload_files(vector_store_id: str, files_dir: str):
    """
    Загружает новые файлы в хранилище, пропуская уже существующие
    Возвращает список file_id новых файлов
    """
    existing = load_existing_files()
    new_files = []

    for filename in os.listdir(files_dir):
        file_path = os.path.join(files_dir, filename)
        file_hash = calculate_hash(file_path)

        # Если файл с таким именем и незменённым хэшем уже существует – пропускаем
        if filename in existing and existing[filename]['hash'] == file_hash:
            logger.info(f"Файл {filename} уже загружен, пропуск")
            continue

        try:
            with open(file_path, "rb") as f:
                # Загрузка файла в OpenAI (файл попадет в список "files")
                uploaded_file = client.files.create(file=f, purpose="assistants")

            # Добавление файла в векторное хранилище с (пока пустыми) атрибутами
            # (Если атрибуты вы укажете в update_attributes позже,
            # здесь можно передать {} или заранее известные значения.)
            client.vector_stores.files.create(
                vector_store_id=vector_store_id,
                file_id=uploaded_file.id,
                attributes={}   # атрибуты можно обновить позже
            )

            # Сохраняем сведения о файле локально
            existing[filename] = {
                'file_id': uploaded_file.id,
                'hash': file_hash,
                'attributes': {}   # и здесь храните текущие атрибуты
            }
            new_files.append(uploaded_file.id)

        except Exception as e:
            logger.error(f"Ошибка загрузки {filename}: {str(e)}")

    save_existing_files(existing)
    return new_files


# 3. Генерация шаблона метаданных
def generate_attributes_template(vector_store_id: str):
    """Создает JSON-файл с текущими метаданными для ручного редактирования"""
    files_info = load_existing_files()
    vector_files = client.vector_stores.files.list(vector_store_id=vector_store_id)

    template = {}
    for v_file in vector_files:
        filename = next((k for k, v in files_info.items() if v['file_id'] == v_file.id), None)
        if filename:
            template[v_file.id] = {
                "filename": filename,
                "attributes": files_info[filename].get('attributes', {})
            }

    with open("attributes_template.json", "w") as f:
        json.dump(template, f, indent=2, ensure_ascii=False)
    logger.info("Шаблон метаданных сгенерирован")


# 4. Обновление метаданных
def update_attributes(vector_store_id: str, attributes_file: str):
    """Обновляет метаданные файлов на основе редактированного шаблона"""
    with open(attributes_file, "r") as f:
        attributes = json.load(f)

    existing = load_existing_files()
    updated = 0

    for file_id, data in attributes.items():
        filename = data['filename']
        if filename not in existing:
            continue

        try:
            # Обновляем метаданные через API
            client.vector_stores.files.update(
                vector_store_id=vector_store_id,
                file_id=file_id,
                attributes=data['attributes']
            )

            # Обновляем локальную копию
            existing[filename]['attributes'] = data['attributes']
            updated += 1

        except Exception as e:
            logger.error(f"Ошибка обновления {filename}: {str(e)}")

    save_existing_files(existing)
    logger.info(f"Обновлено {updated} записей")


# Вспомогательные функции
def calculate_hash(file_path: str) -> str:
    """Вычисляет SHA-256 хеш файла для отслеживания изменений"""
    sha = hashlib.sha256()
    with open(file_path, "rb") as f:
        while True:
            data = f.read(65536)
            if not data:
                break
            sha.update(data)
    return sha.hexdigest()


def load_existing_files() -> dict:
    """Загружает информацию о ранее загруженных файлах"""
    try:
        with open("existing_files.json", "r") as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return {}


def save_existing_files(data: dict):
    """Сохраняет информацию о загруженных файлах"""
    with open("existing_files.json", "w") as f:
        json.dump(data, f, indent=2)


def delete_file_by_filename(filename: str):
    """
    Удаляет файл из OpenAI по имени (если он есть в existing_files.json)
    """
    existing = load_existing_files()
    if filename not in existing:
        print(f"Файл '{filename}' не найден в existing_files.json")
        return

    file_id = existing[filename]['file_id']
    vector_store_id = VS_ID
    try:
        client.vector_stores.files.delete(
            vector_store_id=vector_store_id,
            file_id=file_id
        )
        print(f"Файл '{filename}' c ID {file_id} успешно удалён")

        # Удаляем из локального кэша
        del existing[filename]
        save_existing_files(existing)
    except Exception as e:
        print(f"Ошибка при удалении файла '{file_id}': {str(e)}")

if __name__ == "__main__":

    # ==== Запуск этапов ====
    if RUN_CREATE_VECTOR_STORE:
        VS_ID = create_vector_store(VECTOR_STORE_NAME)

    if RUN_UPLOAD_FILES:
        new_files = upload_files(VS_ID, FILES_DIR)
    else:
        new_files = []

    if RUN_GENERATE_TEMPLATE:
        generate_attributes_template(VS_ID)

    if RUN_UPDATE_ATTRIBUTES:
        update_attributes(VS_ID, ATTRIBUTES_FILE)

    if RUN_DELETE_FILE:
        delete_file_by_filename(FILE_NAME_DELETE)