import logging
from typing import Optional

from smart_bot_factory.rag import RagRouter, VectorStore
from typing import Literal

service_categories = VectorStore(
    bot_id="mdclinica",
    table_name="service_categories"
)

services = VectorStore(
    bot_id="mdclinica",
    table_name="services"
)

logger = logging.getLogger(__name__)

rag_router = RagRouter("mdclinica_rag")

def _format_results(results) -> str:
    text = ""
    for index, doc in enumerate(results):
        text += f"#{index + 1}\n {doc.page_content}\n\n"
    return text


@rag_router.tool
async def get_info_from_rag(query: str, doc_type: Literal["category_section", "service_section"], section: Optional[str] = None) -> str:
    """Запрос информации из RAG-системы.

    Args:
        query: Запрос к RAG-системе.
        doc_type: Тип документа (например, `category_section` или `service_section`, `indications`).
        section: Раздел внутри категории (например, `services`). Отправить можешь несколько через |. Пример - "services|prices|about_procedure|indications". Если не нужно фильтровать по разделу, то не отправляй этот параметр.
    Returns:
        Строка с подборкой категорий (один чанк на блок текста).
    """
    logger.info("Запрос категорий услуг: %s", query)
    
    metadata_filter = {
        "doc_type": doc_type
    }
    
    if section:
        metadata_filter["section"] = section.split("|")
    
    logger.info("Фильтры metadata для категорий: %s", metadata_filter)

    # Выбираем правильный vectorstore в зависимости от doc_type
    if doc_type == "category_section":
        vectorstore = service_categories
    elif doc_type == "service_section":
        vectorstore = services
    # Используем новый метод asimilarity_search с автоматической фильтрацией по score
    # score по умолчанию 0.6, можно изменить через параметр
    results = await vectorstore.asimilarity_search(
        query, k=10, filter=metadata_filter, score=0.55
    )
    
    text = _format_results(results)
    logger.info("Текст результатов: %s", text)
    
    return text



__all__ = ["get_info_from_rag", "rag_router"]