"""Модуль для работы с эмбеддингами текстовых данных."""
import logging
import os
from typing import Dict, List

import requests
from pymilvus import model as milvus_model

from storage.milvus_client import MilvusConnectionManager

logger = logging.getLogger(__name__)


class EmbeddingService:
    """Сервис для генерации эмбеддингов и поиска векторов в базе данных."""

    def __init__(self, milvus_manager: MilvusConnectionManager) -> None:
        """Инициализация сервиса с использованием менеджера подключений к
        Milvus.

        MilvusConnectionManager
        """
        self.milvus_manager = milvus_manager

    def get_embeddings_from_api(
        self,
        chunks: List[str],
        url_embedder: str,
        model_name: str,
        batch_size: int = 32,
    ) -> List[float]:
        """Получение эмбеддингов для текстовых фрагментов через API.

        С помощью эмбеддера. Разбивает запросы на батчи для оптимизации
        производительности.
        """
        if not url_embedder:
            raise ValueError('URL_EMBEDDER environment variable not set')

        embeddings = []
        for index in range(0, len(chunks), batch_size):
            batch = chunks[index : index + batch_size]
            batch_data = {'input': batch, 'model': model_name}
            response = requests.post(url_embedder, headers={'Content-Type': 'application/json'}, json=batch_data)

            if response.status_code == requests.codes.ok:
                embeddings.extend([elem['embedding'] for elem in response.json()['data']])
            else:
                logger.error('API Error: %s, %s', response.status_code, response.text)
                raise Exception(f'API Error: {response.status_code}, {response.text}')

        return embeddings

    def load_milvus_database_with_link_filter(
            self,
            query: str,
            embedder_name: str,
            top_k: int,
            reference_project: str,
            attached_link: str,
            include_attached: bool = False
        ) -> List[Dict[str, str]]:
            """Поиск документов в базе Milvus с фильтрацией по attached_link.
            
            Args:
                query: Поисковый запрос
                embedder_name: Название модели эмбеддингов
                top_k: Количество результатов для возврата
                reference_project: Название проекта в базе
                attached_link: Ссылка для фильтрации
                include_attached: True - поиск ТОЛЬКО по attached странице, 
                                False - поиск БЕЗ attached страницы
            
            Returns:
                Список найденных документов с content, page_link и page_title
            """
            COLLECTION_NAME = f'project_{reference_project}'
            collection = self.milvus_manager.get_collection(COLLECTION_NAME)

            if not collection:
                return [{'content': '', 'page_link': '', 'page_title': ''}]

            url_embedder = os.getenv('URL_EMBEDDER')

            if url_embedder:
                query_vectors = self.get_embeddings_from_api([query], url_embedder, embedder_name, batch_size=1)
            else:
                embedding_model = milvus_model.dense.SentenceTransformerEmbeddingFunction(
                    model_name=embedder_name,
                    device='cpu',
                )
                query_vectors = embedding_model.encode_queries([query])
            try:
                # Формируем условие фильтрации
                if attached_link:
                    # Извлекаем pageId из attached_link
                    attached_page_id = attached_link.split('=')[-1]
                    
                    if include_attached:
                        # Поиск ТОЛЬКО по attached странице
                        expr = f'link like "%pageId={attached_page_id}%"'
                    else:
                        # Поиск БЕЗ attached страницы
                        expr = f'!(link like "%pageId={attached_page_id}%")'
                else:
                    # Если attached_link нет, ищем по всем документам
                    expr = 'id >= 0'

                search_results = collection.search(
                    data=query_vectors,
                    limit=top_k,
                    anns_field='embedding',
                    param={'metric_type': 'COSINE'},
                    expr=expr,
                    output_fields=['content', 'link', 'page_title'],
                )

                documents = []
                for hits in search_results:
                    for hit in hits:
                        document = {
                            'content': hit.entity.get('content'),
                            'page_link': hit.entity.get('link'),
                            'page_title': hit.entity.get('page_title'),
                        }
                        documents.append(document)

                self.milvus_manager.release_collection(collection)
                return documents
            except Exception as error:
                logger.error('Ошибка при получение контекста из milvus: %s', str(error))
                return [{'content': '', 'page_link': '', 'page_title': ''}]
