fckbot/bots/workers/indexing_worker.py

# -*- coding: utf-8 -*-
"""
Фоновый воркер (worker) для индексации файлов и URL через RAG-сервер.
Читает задачи из asyncio.Queue и выполняет:
- Для URL: скачивание, извлечение текста, индексация через HTTP-запрос к серверу.
- Для файлов: извлечение текста, обработка изображений/аудио, индексация через HTTP-запрос.

Все операции с БД и Qdrant выполняются на сервере. Воркер только отправляет данные.
"""

import asyncio
import logging
import os
import hashlib
import shutil
from datetime import datetime

logger = logging.getLogger(__name__)


class IndexingWorker:
    """
    Фоновый воркер для индексации.
    Получает задачи из очереди и выполняет их, отправляя данные на RAG-сервер.
    """

    def __init__(self, bot):
        """
        Аргументы:
            bot: экземпляр Bot (содержит rag_client, config, file_handler и состояния)
        """
        self.bot = bot
        logger.info("IndexingWorker инициализирован (режим: HTTP-клиент RAG-сервера)")

    async def run(self):
        """Бесконечный цикл обработки задач из очереди."""
        while True:
            task = await self.bot.indexing_queue.get()
            try:
                if task.get("type") == "url":
                    await self._process_url(task)
                else:
                    await self._process_file(task)
            except Exception as e:
                logger.exception(f"Ошибка в воркере при обработке задачи: {e}")
                self.bot.metrics_handler.increment("errors_total")
            finally:
                self.bot.indexing_queue.task_done()

    # ==========================================================
    # Обработка URL
    # ==========================================================

    async def _process_url(self, task):
        """
        Обрабатывает URL-задачу.
        - Одиночный URL: загружает страницу/PDF, извлекает текст, индексирует через сервер.
        - Рекурсивный обход: через crawl_url, затем индексирует каждую страницу.
        """
        from core.utils.web_utils import fetch_any_url, crawl_url

        url = task["url"]
        user_jid = task["user_jid"]
        room_jid = task.get("room_jid")
        custom_title = task.get("title")
        is_global = task.get("is_global", False)
        recursive = task.get("recursive", False)
        depth = task.get("depth", 1)
        max_pages = task.get("max_pages", self.bot.config.ws_max_pages)

        try:
            if recursive:
                pages = await crawl_url(url, max_depth=depth, max_pages=max_pages, domain_restrict=True)
                if not pages:
                    await self._send_message(
                        user_jid,
                        f"❌ Не удалось загрузить ни одной страницы с {url}",
                        room_jid
                    )
                    return
                total = 0
                for page_url, page_title, page_text in pages:
                    if page_text:
                        doc_title = custom_title if custom_title else (page_title if page_title else page_url)
                        # Отправляем на сервер для индексации
                        result = await self.bot.rag_client.index_document(
                            file_name=page_url,
                            file_text=page_text,
                            user_jid=user_jid,
                            room_jid=room_jid,
                            is_global=is_global,
                            title=doc_title,
                            metadata={"url": page_url}
                        )
                        if result.get('doc_id'):
                            total += 1
                        else:
                            logger.error(f"Ошибка индексации {page_url}: {result.get('error')}")
                await self._send_message(
                    user_jid,
                    f"✅ Проиндексировано {total} страниц (рекурсивно).",
                    room_jid
                )
            else:
                title, text = await fetch_any_url(url)
                if text:
                    doc_title = custom_title if custom_title else (title if title else url)
                    result = await self.bot.rag_client.index_document(
                        file_name=url,
                        file_text=text,
                        user_jid=user_jid,
                        room_jid=room_jid,
                        is_global=is_global,
                        title=doc_title,
                        metadata={"url": url}
                    )
                    if result.get('doc_id'):
                        scope = "глобальную" if is_global else ("комнату" if room_jid else "личную")
                        await self._send_message(
                            user_jid,
                            f"✅ Страница '{doc_title}' добавлена в {scope} БЗ.",
                            room_jid
                        )
                    else:
                        await self._send_message(
                            user_jid,
                            f"❌ Ошибка индексации {url}: {result.get('error')}",
                            room_jid
                        )
                else:
                    await self._send_message(
                        user_jid,
                        f"❌ Не удалось загрузить или извлечь текст из {url}",
                        room_jid
                    )
        except Exception as e:
            logger.exception(f"Ошибка индексации URL {url}: {e}")
            await self._send_message(
                user_jid,
                f"❌ Ошибка при обработке URL: {e}",
                room_jid
            )

    # ==========================================================
    # Обработка файлов
    # ==========================================================

    async def _process_file(self, task):
        """
        Обрабатывает задачу с файлом:
        1. Извлекает текст (синхронно в потоке).
        2. Обрабатывает изображения/аудио (Vision/распознавание) – если есть.
        3. Сохраняет постоянную копию в data/ (локально, для истории и отладки).
        4. Отправляет текст и метаданные на RAG-сервер для индексации.
        5. Удаляет временный файл.
        """
        user_jid = task["user_jid"]
        room_jid = task.get("room_jid")
        file_path = task["file_path"]
        original_fname = task["original_fname"]
        title = task.get("title")
        is_global = task.get("is_global", False)
        mode = task.get("mode")
        ext = task.get("ext", "").lower()

        try:
            # --- 1. Извлечение текста (синхронная операция в потоке) ---
            result = await asyncio.to_thread(self.bot.files.process_any_file, file_path)
            if isinstance(result, tuple) and len(result) == 2:
                first, second = result
                if isinstance(first, list):
                    files_list = first
                    total_count = second
                else:
                    files_list = [(original_fname, first)] if first else []
                    total_count = second
            else:
                files_list = [(original_fname, result)] if result else []
                total_count = 1 if result else 0

            if total_count == 0:
                await self._send_message(
                    user_jid,
                    f"⚠️ Не удалось извлечь текст из {original_fname}.",
                    room_jid
                )
                return

            # Сохраняем текст для локального использования (например, !summary)
            full_text = "\n\n".join([t for _, t in files_list if t])
            self.bot.last_file_texts[user_jid] = full_text

            # --- 2. Сохранение постоянной копии в data/ (локально) ---
            # Это нужно для команд !summary и хирургической замены, которые работают с файлом локально.
            # В будущем можно перенести и их на сервер, но пока оставляем.
            with open(file_path, 'rb') as f:
                file_hash = hashlib.sha256(f.read()).hexdigest()
            short_hash = file_hash[:12]
            safe_user = user_jid.replace('@', '_at_').replace('/', '_')
            safe_fname = original_fname.replace('/', '_').replace('\\', '_')
            dest_filename = f"user_{short_hash}_{safe_user}_{safe_fname}"
            dest_path = self.bot.config.data_dir / dest_filename

            if not dest_path.exists():
                shutil.copy2(file_path, dest_path)
                logger.debug(f"Создана постоянная копия: {dest_path}")
            else:
                logger.debug(f"Постоянная копия уже существует: {dest_path}")

            key = (user_jid, room_jid) if room_jid is not None else user_jid
            self.bot._last_file_path[key] = str(dest_path)

            # --- 3. Специальная обработка: изображения (Vision API) ---
            # Если это изображение и включена опция vision, отправляем на анализ.
            # Это можно делать отдельно, не блокируя индексацию.
            if ext in ('.jpg', '.jpeg', '.png', '.bmp', '.tiff') and self.bot.config.vision:
                file_id = await self.bot.giga.upload_file(file_path)  # giga всё ещё локально? нужно будет перенести и это.
                # Это временно, пока мы не вынесли всё на сервер. Пока оставим как есть.
                if file_id:
                    answer = await self.bot.giga.chat(
                        history=[],
                        query="Проанализируй это изображение. Если есть текст – прочитай его.",
                        system_prompt=None,
                        file_id=file_id,
                        temperature=self.bot.config.ai_temperature
                    )
                    await self._send_message(user_jid, answer, room_jid)
                else:
                    await self._send_message(user_jid, "⚠️ Не удалось передать изображение в GigaChat.", room_jid)
                # После обработки изображения мы не индексируем его как документ (обычно OCR делает это),
                # но можно и проиндексировать текст. Здесь мы просто возвращаемся.
                return

            # --- 4. Специальная обработка: аудио (голос → текст) ---
            # Аналогично, используем локальный SaluteSpeech, пока не вынесли.
            if ext in ('.ogg', '.wav', '.mp3', '.amr', '.m4a') and self.bot.config.voice_recognition:
                first_text = files_list[0][1] if files_list else ""
                if "--- РАСПОЗНАННЫЙ ГОЛОС ---" in first_text:
                    voice_text = first_text.split("--- РАСПОЗНАННЫЙ ГОЛОС ---")[1].strip()
                    await self._send_message(user_jid, f"🎤 Распознано: {voice_text}", room_jid)
                    # Можно также ответить на голос, но это отдельный сценарий.
                    # Мы не индексируем аудио как документ (разве что текст распознавания).
                    # Если хотим индексировать текст голоса, нужно передать его как документ.
                    # Например, можно создать виртуальный документ с текстом.
                else:
                    await self._send_message(user_jid, f"⚠️ {first_text}", room_jid)
                return

            # --- 5. Индексация документов через RAG-сервер ---
            # Отправляем на сервер каждый извлечённый файл (включая вложенные из архивов)
            if mode in ('personal', 'global') or (mode == 'room' and room_jid):
                saved = 0
                for sub_name, sub_text in files_list:
                    if sub_text:
                        result = await self.bot.rag_client.index_document(
                            file_name=sub_name,
                            file_text=sub_text,
                            user_jid=user_jid,
                            room_jid=room_jid,
                            is_global=is_global,
                            title=title,
                            # file_hash=file_hash,  # можно передать, если нужно
                            update_if_exists=True
                        )
                        if result.get('doc_id'):
                            saved += 1
                        else:
                            logger.error(f"Ошибка индексации {sub_name}: {result.get('error')}")
                if saved:
                    if room_jid:
                        scope = f"комнату {room_jid}"
                    elif is_global:
                        scope = "глобальную"
                    else:
                        scope = "личную"
                    await self._send_message(
                        user_jid,
                        f"🎓 {saved} файлов добавлено в {scope} БЗ.",
                        room_jid
                    )
                else:
                    await self._send_message(
                        user_jid,
                        f"⚠️ Не удалось проиндексировать {original_fname}. Проверьте логи.",
                        room_jid
                    )
            else:
                # Если режим обучения не включён (или не совпадает), просто показываем текст
                preview = full_text[:500] + ("..." if len(full_text) > 500 else "")
                await self._send_message(
                    user_jid,
                    f"📄 Файл `{original_fname}` прочитан.\n\n{preview}",
                    room_jid
                )
        except Exception as e:
            logger.exception(f"Ошибка при обработке файла {original_fname}: {e}")
            self.bot.metrics_handler.increment("errors_total")
            await self._send_message(
                user_jid,
                f"❌ Ошибка при обработке файла '{original_fname}'. Проверьте логи.",
                room_jid
            )
        finally:
            # Удаляем временный файл (постоянная копия уже сохранена)
            try:
                if os.path.exists(file_path):
                    os.unlink(file_path)
                    logger.debug(f"Временный файл удалён: {file_path}")
            except Exception as e:
                logger.error(f"Не удалось удалить временный файл {file_path}: {e}")

    # ==========================================================
    # Вспомогательные методы
    # ==========================================================

    async def _send_message(self, user_jid: str, body: str, room_jid: str = None):
        """
        Отправляет сообщение пользователю (в личный чат или комнату).
        Используется для уведомлений о результатах индексации.
        """
        if room_jid:
            self.bot.send_message(mto=room_jid, mbody=body, mtype='groupchat')
        else:
            self.bot.send_message(mto=user_jid, mbody=body, mtype='chat')