fckbot/bots/handlers/file_handler.py

# -*- coding: utf-8 -*-
"""
Обработчик файлов: скачивание файлов по URL (из OOB), сохранение во временную директорию,
постановка задачи в очередь индексации.

Ключевое изменение для групповой работы:
- Если файл пришёл в групповой чат (room_jid не None), он автоматически индексируется
  в базу знаний комнаты, без необходимости команды !learn.
- Для личных чатов – поведение зависит от режима обучения (learning_mode) и флага --global.

ДОБАВЛЕНО:
- Удалён прямой вызов GigaChat для изображений (теперь это делает RAG-сервер)
- Удалён прямой вызов транскрибации (теперь это делает RAG-сервер)
"""

import os
import asyncio
import logging
import urllib.parse
import requests
from datetime import datetime
import re

logger = logging.getLogger(__name__)


class FileHandler:
    """
    Отвечает за скачивание файлов, переданных через ссылку в сообщении (XEP-0066 OOB).
    """

    def __init__(self, bot):
        self.bot = bot

    async def _wait_for_queue(self):
        """
        Ожидает создания очереди индексации (защита от гонки при старте бота).

        Проблема: очередь `self.bot.indexing_queue` создаётся в `session_start` после
        подключения к БД. Если сообщение с файлом придёт ДО этого момента,
        `indexing_queue` будет `None`, и `put()` вызовет исключение.
        Решение: ждём, пока очередь станет доступна, проверяя каждые 0.5 секунды.
        """
        while self.bot.indexing_queue is None:
            logger.debug("Очередь индексации ещё не создана, ожидание...")
            await asyncio.sleep(0.5)

    async def handle_file(self, msg, url: str, user_jid: str, room_jid: str = None):
        """
        Скачивает файл и создаёт задачу для indexing_worker.

        Аргументы:
            msg: stanza сообщения (нужна для получения user, ответных сообщений).
            url (str): HTTP/HTTPS ссылка на файл (из атрибута url элемента oob).
            user_jid (str): JID отправителя (без ресурса).
            room_jid (str | None): JID комнаты, если сообщение из groupchat; иначе None.

        Логика:
            1. Извлечь имя файла из URL, декодировать, сгенерировать уникальное имя во временной папке.
            2. Скачать файл (синхронно, в отдельном потоке, с таймаутом из конфига).
            3. Определить режим индексации:
                - Если room_jid не None -> mode = 'room', is_global = False (всегда).
                - Если личный чат -> проверить learning_mode пользователя.
            4. Извлечь --title из тела сообщения (если есть).
            5. Положить задачу в очередь индексации (asyncio.Queue).
            6. Ответить пользователю, что файл принят в обработку.
        """
        self.bot.metrics_handler.increment("files_processed")
        try:
            # --- 1. Извлечение имени файла из URL ---
            raw_name = url.split('/')[-1].split('?')[0]
            original_fname = urllib.parse.unquote(raw_name)
            ext = os.path.splitext(original_fname)[1].lower()
            fname = f"{msg['from'].user}_{int(datetime.now().timestamp())}{ext}"
            target_path = os.path.join(self.bot.config.temp_dir, fname)

            # --- 2. Скачивание файла ---
            # HEAD-запрос для получения размера
            try:
                head_resp = requests.head(url, timeout=self.bot.config.http_download_timeout)
                content_length = head_resp.headers.get('Content-Length')
                if content_length is not None and int(content_length) > self.bot.config.max_file_size_mb * 1024 * 1024:
                    reply = msg.reply(f"❌ Файл превышает максимальный размер {self.bot.config.max_file_size_mb} МБ.")
                    if reply:
                        reply.send()
                    return
            except Exception:
                pass

            # Скачивание с прерыванием по лимиту
            def download_with_limit():
                r = requests.get(url, stream=True, timeout=self.bot.config.http_download_timeout)
                if r.status_code != 200:
                    return False
                downloaded = 0
                limit = self.bot.config.max_file_size_mb * 1024 * 1024
                with open(target_path, 'wb') as f:
                    for chunk in r.iter_content(8192):
                        f.write(chunk)
                        downloaded += len(chunk)
                        if downloaded > limit:
                            os.unlink(target_path)
                            raise Exception("File too large")
                return True

            try:
                success = await asyncio.to_thread(download_with_limit)
                if not success:
                    reply = msg.reply("❌ Не удалось загрузить файл. Проверьте доступность ссылки.")
                    if reply:
                        reply.send()
                    return
            except Exception as e:
                if "File too large" in str(e):
                    reply = msg.reply(f"❌ Файл превышает максимальный размер {self.bot.config.max_file_size_mb} МБ.")
                else:
                    reply = msg.reply("❌ Ошибка при скачивании файла.")
                if reply:
                    reply.send()
                return

            # --- 3. Определение режима индексации ---
            if room_jid:
                mode = 'room'
                is_global = False
                title = None
            else:
                mode_info = self.bot.learning_mode.get(user_jid)
                if mode_info is None:
                    mode = None
                    is_global = False
                    title = None
                else:
                    if isinstance(mode_info, tuple):
                        mode, saved_room = mode_info
                        is_global = (mode == 'global')
                    else:
                        mode = mode_info
                        is_global = (mode == 'global')
                    title = None

            # --- 4. Парсинг --title из тела сообщения ---
            raw_body = msg['body'].strip() if msg['body'] else ""
            if raw_body:
                match = re.search(r'--title "([^"]+)"', raw_body)
                if match:
                    title = match.group(1)

            # --- 5. Создание задачи для очереди ---
            await self._wait_for_queue()
            task = {
                "user_jid": user_jid,
                "room_jid": room_jid,
                "file_path": target_path,
                "original_fname": original_fname,
                "title": title,
                "is_global": is_global,
                "mode": mode,
                "ext": ext
            }
            await self.bot.indexing_queue.put(task)

            # --- 6. Уведомление пользователя ---
            if room_jid:
                reply = msg.reply(f"📥 Файл '{original_fname}' принят в обработку для КОМНАТНОЙ БЗ.")
            else:
                scope = "глобальную" if is_global else "личную"
                reply = msg.reply(f"📥 Файл '{original_fname}' принят в обработку для {scope} БЗ.")
            if reply:
                reply.send()

        except Exception as e:
            logger.exception(f"Ошибка при скачивании файла: {e}")
            self.bot.metrics_handler.increment("errors_total")
            reply = msg.reply("❌ Ошибка при обработке файла. Проверьте логи.")
            if reply:
                reply.send()