fckbot/rag/functions/file_processor.py

# -*- coding: utf-8 -*-
"""
Обработка файлов: отправка пользователю через HTTP Upload и хирургическая замена в DOCX.
Функции: send_file_to_user, handle_ai_file_request.

АДАПТАЦИЯ: эти функции уже работают через HTTP (загружают файлы в ejabberd).
Они не требуют изменений, так как работают с локальными файлами.
Но я добавлю комментарии для ясности.
"""

import asyncio
import os
import re
import shutil
import time
import random
import string
import logging

logger = logging.getLogger(__name__)


async def send_file_to_user(bot, msg, file_path, prefix="", suffix=None):
    """
    Отправляет файл пользователю через HTTP Upload (XEP-0363).

    Аргументы:
        bot: экземпляр бота
        msg: объект XMPP-сообщения
        file_path: путь к файлу на диске
        prefix: текст перед ссылкой
        suffix: суффикс для имени файла

    Возвращает:
        bool: True при успехе, False при ошибке
    """
    if not os.path.exists(file_path):
        logger.error(f"Файл не найден: {file_path}")
        return False

    original_name = os.path.basename(file_path)
    name, ext = os.path.splitext(original_name)

    # Санитизация суффикса (только буквы, цифры, подчёркивание, дефис)
    if suffix:
        suffix = re.sub(r'[^a-zA-Zа-яА-ЯёЁ0-9_-]', '', suffix)
        suffix = suffix.strip('.')
        if not suffix or suffix == '_':
            suffix = '_sanitized'
        if not suffix.startswith('_'):
            suffix = '_' + suffix
        new_name = f"{name}{suffix}{ext}"
        new_path = os.path.join(os.path.dirname(file_path), new_name)
        os.rename(file_path, new_path)
        file_path = new_path

    try:
        # Загружаем файл через XEP-0363 (HTTP Upload)
        result = await bot.plugin['xep_0363'].upload_file(file_path)
        if isinstance(result, dict):
            file_url = result.get('get_url', result.get('url', str(result)))
        else:
            file_url = str(result)
    except Exception as e:
        logger.error(f"Ошибка загрузки файла через XEP-0363: {e}", exc_info=True)
        error_reply = msg.reply("❌ Не удалось загрузить файл. Проверьте настройки HTTP Upload.")
        if error_reply:
            await error_reply.send()
        return False

    reply = msg.reply(f"{prefix}\n📄 {file_url}")
    if reply:
        reply.send()
    else:
        logger.error("Не удалось создать reply для ссылки на файл")
    return True


async def handle_ai_file_request(bot, msg, ai_response, user_jid, room_jid=None, body=None, suffix=None):
    """
    Обрабатывает теги [SURGICAL_REPLACE], [REWRITE], [FILE] из ответа AI.

    Поддерживает три режима:
        - SURGICAL: замена текста в существующем DOCX
        - REWRITE: создание нового DOCX из текста
        - FILE: создание нового DOCX (общий случай)
    """
    key = (user_jid, room_jid) if room_jid is not None else user_jid

    surgical_match = re.search(r'\[SURGICAL_REPLACE\](.*?)\[/SURGICAL_REPLACE\]', ai_response, re.DOTALL)
    rewrite_match = re.search(r'\[REWRITE\](.*?)\[/REWRITE\]', ai_response, re.DOTALL)
    file_match = re.search(r'\[FILE\](.*?)\[/FILE\]', ai_response, re.DOTALL)

    # ---- SURGICAL_REPLACE ----
    if surgical_match:
        content = surgical_match.group(1).strip()
        replacements = {}
        for line in content.split('\n'):
            line = line.strip()
            if '|||' in line:
                old, new = line.split('|||', 1)
                replacements[old.strip()] = new.strip()
        if not replacements:
            reply = msg.reply("⚠️ Не найдено пар 'старое ||| новое' в теге SURGICAL_REPLACE.")
            if reply:
                reply.send()
            return

        original_path = bot._last_file_path.get(key)
        if not original_path or not os.path.exists(original_path):
            reply = msg.reply("❌ Нет загруженного документа для замены. Сначала отправьте файл (без !learn).")
            if reply:
                reply.send()
            return

        new_path = await asyncio.to_thread(bot.files.surgical_replace, original_path, replacements)
        if new_path:
            await send_file_to_user(bot, msg, new_path, prefix="✅ Замена выполнена:", suffix="surgical")
            try:
                os.unlink(new_path)
            except Exception as e:
                logger.warning(f"Не удалось удалить временный файл {new_path}: {e}")
        else:
            reply = msg.reply("❌ Ошибка при хирургической замене. Проверьте формат документа.")
            if reply:
                reply.send()
        return

    # ---- REWRITE ----
    if rewrite_match:
        content = rewrite_match.group(1).strip()
        title = "Сгенерированный документ"
        title_match = re.search(r'\[TITLE\](.*?)\[/TITLE\]', ai_response, re.DOTALL)
        if title_match:
            title = title_match.group(1).strip()
        new_path = await asyncio.to_thread(bot.files.create_docx, content, title, user_jid, mode="REWRITE")
        if new_path:
            await send_file_to_user(bot, msg, new_path, prefix="📄 Документ создан:", suffix="rewrite")
            try:
                os.unlink(new_path)
            except Exception as e:
                logger.warning(f"Не удалось удалить временный файл {new_path}: {e}")
        else:
            reply = msg.reply("❌ Не удалось создать документ.")
            if reply:
                reply.send()
        return

    # ---- FILE ----
    if file_match:
        content = file_match.group(1).strip()
        title = "Файл"
        title_match = re.search(r'\[TITLE\](.*?)\[/TITLE\]', ai_response, re.DOTALL)
        if title_match:
            title = title_match.group(1).strip()
        new_path = await asyncio.to_thread(bot.files.create_docx, content, title, user_jid, mode="FILE")
        if new_path:
            await send_file_to_user(bot, msg, new_path, prefix="📄 Файл создан:", suffix="file")
            try:
                os.unlink(new_path)
            except Exception as e:
                logger.warning(f"Не удалось удалить временный файл {new_path}: {e}")
        else:
            reply = msg.reply("❌ Не удалось создать файл.")
            if reply:
                reply.send()
        return

    logger.debug(f"Ответ AI не содержит тегов файлов: {ai_response[:200]}")