"""Модуль получения контекста для пайплайна обработки запросов."""
import logging
import time
from typing import Dict, List, Optional, Tuple
import requests
import os

from utils.db_logger import db_logger

from models.embeddings import EmbeddingService
from models.reranker import RerankerService
from storage.milvus_client import MilvusTableSearch

logger = logging.getLogger(__name__)


class ContextRetriever:
    """Сервис для получения релевантного контекста из различных источников
    данных."""

    def __init__(
        self,
        embedding_service: EmbeddingService,
        reranker_service: RerankerService,
        table_search: MilvusTableSearch,
    ) -> None:
        """Инициализация сервиса получения контекста.

        Args:
            embedding_service: EmbeddingService
            reranker_service: RerankerService
            table_search: MilvusTableSearch
        """
        self.embedding_service = embedding_service
        self.reranker_service = reranker_service
        self.table_search = table_search

        self.lightrag_url = os.getenv('URL_LIGHTRAG')


    def retrieve_document_context(
            self,
            query: str,
            embedder_name: str,
            reranker_name: str,
            reference_project: str,
            attached_link: str,
            embedder_out_chunks: int,
            reranker_out_chunks: int,
            attached_out_chunks: int,
        ) -> Tuple[List[Tuple[str]], Dict[str, float]]:
            """Получение контекста из текстовых документов с поддержкой attached_link.

            Выполняет поиск релевантных документов:
            1. Основной поиск БЕЗ attached страницы (с реранкером)
            2. Дополнительный поиск ТОЛЬКО по attached странице (без реранкера)
            3. Объединение результатов

            Returns:
                - список кортежей (документ, page_info, score)
                - информация о времени (эмбеддер и реранкер)
            """
            timing_data = {}
            final_ranked_results = []

            # === ОСНОВНОЙ ПОИСК (БЕЗ attached страницы) ===
            start = time.perf_counter()
            
            documents = self.embedding_service.load_milvus_database_with_link_filter(
                query=query,
                embedder_name=embedder_name,
                top_k=embedder_out_chunks,
                reference_project=reference_project,
                attached_link=attached_link,
                include_attached=False  # БЕЗ attached страницы
            )
            
            timing_data['embedding'] = time.perf_counter() - start

            # Подготавливаем данные для реранкера
            documents_content = [doc['content'] for doc in documents]
            document_page_infos = [(doc['page_link'], doc['page_title']) for doc in documents]

            # Применяем реранкер к основным результатам
            start = time.perf_counter()
            main_ranked_results = self.reranker_service.rank_documents(
                query=query,
                documents=documents_content,
                document_page_infos=document_page_infos,
                model_name=reranker_name,
                top_k=reranker_out_chunks,
            )
            timing_data['reranking'] = time.perf_counter() - start

            final_ranked_results.extend(main_ranked_results)

            # === ДОПОЛНИТЕЛЬНЫЙ ПОИСК (ТОЛЬКО attached страница) ===
            if attached_link and attached_out_chunks > 0:
                start_attached = time.perf_counter()
                
                attached_documents = self.embedding_service.load_milvus_database_with_link_filter(
                    query=query,
                    embedder_name=embedder_name,
                    top_k=attached_out_chunks,
                    reference_project=reference_project,
                    attached_link=attached_link,
                    include_attached=True  # ТОЛЬКО attached страница
                )
                
                # Добавляем время поиска attached к embedding времени
                timing_data['embedding'] += time.perf_counter() - start_attached
                
                # Преобразуем attached результаты в нужный формат (без реранкера)
                for doc in attached_documents:
                    if doc['content']:  # Проверяем что документ не пустой
                        attached_result = (
                            doc['content'],
                            (doc['page_link'], doc['page_title']),
                            1.0  # Фиксированный score для attached результатов
                        )
                        final_ranked_results.append(attached_result)

            return final_ranked_results, timing_data

    def retrieve_table_context(
        self,
        query: str,
        reference_project: str,
        num_table_out: int,
    ) -> Tuple[str, float]:
        """Получение информации о похожих таблицах(xml схемы).

        Returns:
            - текстовая информация
            - информация о времени
        """
        start = time.perf_counter()
        table_context = self.table_search.search_table_info(
            user_question=query,
            reference_project=reference_project,
            num_table_out=num_table_out,
        )
        timing = time.perf_counter() - start

        return table_context, timing

    def retrieve_graph_context(
        self,
        query: str,
        reference_project: str,
        graph_out_chunks: int
    ) -> Tuple[str, float]:
        """Получение контекста из LightRAG.
        
        Args:
            query: Запрос пользователя
            reference_project: Название проекта
            page_name: Название страницы
            
        Returns:
            - контекст из LightRAG
            - время выполнения запроса
        """
        start = time.perf_counter()
        
        if not self.lightrag_url:
            logger.warning("URL для LightRAG не настроен")
            return "", 0.0
        
        try:
            response = requests.post(
                self.lightrag_url,
                json={
                    "query": query,
                    "project_id": reference_project,
                    "mode": "hybrid",
                    "top_k": graph_out_chunks
                }
            )
            
            if response.status_code == 200:
                lightrag_context = response.json().get("context", "")
                    
                return lightrag_context, time.perf_counter() - start
            else:
                logger.error(f"Ошибка при получении контекста LightRAG: {response.status_code}")
                return "", time.perf_counter() - start
        except Exception as e:
            logger.error(f"Исключение при получении контекста LightRAG: {str(e)}")
            return "", time.perf_counter() - start

    def retrieve_full_context(
        self,
        query: str,
        template: str,
        models_names: Dict[str, str],
        models_params: Dict[str, int],
        reference_project: str,
        attached_link: str,
        is_context: bool = True,
        is_table: bool = True,
        is_graph: bool = True,
        request_id: Optional[str] = None,
    ) -> Tuple[str, Dict[str, float], List[Tuple[str]]]:
        """Получение полного контекста из всех доступных источников."""
        timing_data = {'embedding': 0, 'reranking': 0.0, 'table_search': 0.0, 'graph_search': 0.0}
        context = ''

        names = {'krp_chunk_1_2': 'krp', 'krp_chunk_2_3': 'krp', 'krp_chunk_3_1': 'krp'}

        if is_context:
            full_query = query + template
            doc_results, doc_timing = self.retrieve_document_context(
                query=full_query,
                embedder_name=models_names.get('embedder', ''),
                reranker_name=models_names.get('reranker', ''),
                reference_project=names.get(reference_project, reference_project),
                attached_link=attached_link,
                embedder_out_chunks=models_params.get('embedder_out_chunks', 0),
                reranker_out_chunks=models_params.get('reranker_out_chunks', 0),
                attached_out_chunks=models_params.get('attached_out_chunks', 0),
            )
            context += '\n'.join([doc for doc, _, _ in doc_results])
            timing_data.update(doc_timing)

        if is_graph:
            graph_context, graph_timing = self.retrieve_graph_context(
                query=query,
                reference_project=reference_project,
                graph_out_chunks=models_params.get('graph_out_chunks', 0),
            )

            context += graph_context
            timing_data['graph_search'] = graph_timing

        if is_table:
            table_context, table_timing = self.retrieve_table_context(
                query=query,
                reference_project=names.get(reference_project, reference_project),
                num_table_out=models_params.get('num_table_out', 0),
            )

            context += table_context
            timing_data['table_search'] = table_timing

        if request_id:
            db_logger.log_generation_contexts(
                request_id=request_id,
                document_context=doc_results if doc_results and any(doc for doc, _, _ in doc_results) else [],
                table_context=table_context if is_table else '',
                graph_context=graph_context if graph_context else '',
            )

        top_page_infos = [page_info for _, page_info, _ in doc_results]
        return context, timing_data, top_page_infos