import asyncio
import logging
from typing import List, Dict, Set
from src.metrics import PROCESSING_TIME, PROCESSING_COUNT

logger = logging.getLogger("lightrag")

DELETE_TIMEOUT = 120

class DocumentProcessor:
    """Обработка документов в рамках проектов RAG"""
    
    def __init__(self):
        self.extract_metadata = None
        self.enrich_document = None
        self._init_dependencies()
    
    def _init_dependencies(self):
        """Ленивая инициализация зависимостей"""
        if self.extract_metadata is None:
            from src.fill_rag import extract_metadata, enrich_document
            self.extract_metadata = extract_metadata
            self.enrich_document = enrich_document
    
    async def process_project_documents(self, project_rag, documents: List[Dict], project_id: str) -> tuple[int, int]:
        """
        Обрабатывает документы проекта с логикой замены:
        1. Существующие документы - удаляем
        2. Все документы - вставляем заново
        
        Args:
            project_rag: экземпляр ProjectRAG
            documents: список документов для обработки
            project_id: ID проекта (для логирования)
            
        Returns:
            tuple[int, int]: (количество удаленных документов, количество вставленных документов)
        """
        if not documents:
            logger.info(f"Проект {project_id} не содержит документов")
            return 0, 0
        
        try:
            # 1. Определяем, какие документы уже существуют
            existing_ids = await self._find_existing_documents(project_rag, documents)
            
            # 2. УДАЛЯЕМ все существующие документы с timeout
            deleted_count = 0
            if existing_ids:
                deleted_count = await self._delete_existing_documents(project_rag, existing_ids)
                logger.info(f"🗑️ Удалено {deleted_count} существующих документов")
            
            # 3. ВСТАВЛЯЕМ ВСЕ документы заново (и ранее существующие, и новые)
            inserted_count = await self._insert_all_documents(project_rag, documents)
            
            return deleted_count, inserted_count
            
        except Exception as e:
            logger.error(f"❌ Ошибка при обработке документов проекта {project_id}: {e}")
            raise
    
    async def _find_existing_documents(self, project_rag, documents: List[Dict]) -> Set[str]:
        """
        Находит ID документов, которые уже существуют в LightRAG
        
        Returns:
            Set[str]: множество message_id существующих документов
        """
        # Собираем все message_id из входящих документов
        all_message_ids = [doc["message_id"] for doc in documents]
        
        # filter_keys возвращает ID, которых НЕТ в системе (новые)
        new_message_ids = await project_rag.rag_instance.doc_status.filter_keys(all_message_ids)
        new_ids_set = set(new_message_ids)
        
        # Существующие ID = все ID минус новые ID
        existing_ids = set(all_message_ids) - new_ids_set
        
        logger.info(f"📈 Найдено: {len(existing_ids)} существующих, {len(new_ids_set)} новых документов")
        
        return existing_ids

    async def _delete_existing_documents(self, project_rag, existing_ids: Set[str]) -> int:
        """
        Удаляет существующие документы по их ID с простым timeout
        
        Returns:
            int: количество успешно удаленных документов
        """
        
        deleted_count = 0
        
        logger.info(f"🗑️ Начинаем удаление {len(existing_ids)} документов (timeout: {DELETE_TIMEOUT}s)")
        
        for doc_id in existing_ids:
            try:
                await asyncio.wait_for(
                    project_rag.delete_by_doc_id(doc_id),
                    timeout=DELETE_TIMEOUT
                )
                deleted_count += 1
                logger.debug(f"🗑️ Удален документ {doc_id}")
                
            except asyncio.TimeoutError:
                logger.error(f"⏰ Timeout удаления документа {doc_id} после {DELETE_TIMEOUT}s - пропускаем")
                continue
                
            except Exception as e:
                logger.warning(f"⚠️ Ошибка удаления документа {doc_id}: {e}")
                continue
        
        return deleted_count

    async def _insert_all_documents(self, project_rag, documents: List[Dict]) -> int:
        """
        Вставляет ВСЕ документы с их message_id
        
        Returns:
            int: количество вставленных документов
        """
        if not documents:
            return 0
        
        try:
            enriched_contents = []
            message_ids = []
            
            for doc in documents:
                metadata = self.extract_metadata(doc["text"])
                enriched_content = self.enrich_document(doc["text"], metadata)
                enriched_contents.append(enriched_content)
                message_ids.append(doc["message_id"])
            
            await project_rag.insert(enriched_contents, ids=message_ids)
            
            logger.info(f"✅ Вставлено {len(documents)} документов")
            return len(documents)
            
        except Exception as e:
            logger.error(f"Ошибка при вставке документов: {e}")
            raise
    
    async def full_rebuild_project_database(self, project_rag):
        """
        Выполняет полную перестройку базы данных проекта:
        1. Удаляет базу данных PostgreSQL
        2. Удаляет рабочую директорию проекта
        """
        try:
            project_id = project_rag.project_id
            logger.info(f"🔥 Начинаем полную перестройку БД для проекта {project_id}")
            
            if hasattr(project_rag.rag_instance, 'close_connections'):
                await project_rag.rag_instance.close_connections()

            # Обнуляем ссылку ПЕРЕД удалением БД  
            project_rag.rag_instance = None

            # 1. Удаляем базу данных
            await project_rag.drop_database()
            logger.info(f"✅ База данных для проекта {project_id} удалена")
            
            # 2. Удаляем рабочую директорию
            project_rag.cleanup_working_directory()
            logger.info(f"✅ Рабочая директория для проекта {project_id} очищена")
            
            logger.info(f"🎉 Полная перестройка БД для проекта {project_id} завершена")
            
        except Exception as e:
            logger.error(f"❌ Ошибка при полной перестройке БД проекта {project_rag.project_id}: {e}")
            raise
    
    def record_processing_metrics(self, project_id: str, inserted_count: int, processing_time: float, success: bool = True):
        """Записывает метрики обработки документов"""
        status = "success" if success else "error"
        
        if inserted_count > 0 and success:
            time_per_doc = processing_time / inserted_count
            for _ in range(inserted_count):
                PROCESSING_TIME.observe(time_per_doc)
                PROCESSING_COUNT.labels(
                    project_id=project_id, 
                    status=status
                ).inc()
        else:
            # Для ошибок записываем общее время
            PROCESSING_TIME.observe(processing_time)
            PROCESSING_COUNT.labels(
                project_id=project_id, 
                status=status
            ).inc()
