# Пользователь написал
question = "не могу законнектиться"

# Поиск нашёл документы про "коннект" и "проблемы соединения"
# Но пропустил документ "Настройка DATABASE_URL" — там слова "коннект" нет
# А именно там был ответ

from anthropic import Anthropic
import json

client = Anthropic()


def generate_queries(original_question: str, n: int = 4) -> list[str]:
    """Генерирует N перефразировок вопроса"""

    response = client.messages.create(
        model="claude-opus-4-5",
        max_tokens=512,
        messages=[{
            "role": "user",
            "content": f"""Перефразируй этот вопрос {n} разными способами для поиска по документации.
Каждая версия должна использовать разные слова и подходить с разного угла.
Отвечай только JSON-массивом строк, без пояснений.

Вопрос: {original_question}"""
        }]
    )

    queries = json.loads(response.content[0].text)
    return [original_question] + queries  # оригинал + перефразировки


# Пример
queries = generate_queries("не могу законнектиться")
print(queries)
# [
#   "не могу законнектиться",
#   "ошибка подключения к базе данных",
#   "как настроить соединение с PostgreSQL",
#   "DATABASE_URL не работает",
#   "Connection refused при запуске приложения"
# ]

from sentence_transformers import SentenceTransformer
import numpy as np
from typing import List, Dict


model = SentenceTransformer("all-MiniLM-L6-v2")


def multi_query_search(
    question: str,
    documents: List[Dict],
    top_k: int = 5,
    n_queries: int = 4
) -> List[Dict]:

    # Генерируем несколько версий вопроса
    queries = generate_queries(question, n=n_queries)
    print(f"Ищем по {len(queries)} запросам:")
    for q in queries:
        print(f"  - {q}")

    doc_embeddings = model.encode([d["content"] for d in documents])

    # Собираем результаты по всем запросам
    all_scores: Dict[int, float] = {}  # doc_index → max_score

    for query in queries:
        query_emb = model.encode([query])[0]

        similarities = [
            np.dot(query_emb, doc_emb) / (np.linalg.norm(query_emb) * np.linalg.norm(doc_emb))
            for doc_emb in doc_embeddings
        ]

        for doc_idx, score in enumerate(similarities):
            # Берём максимальный score по всем запросам (Reciprocal Rank Fusion — более сложный вариант)
            all_scores[doc_idx] = max(all_scores.get(doc_idx, 0), float(score))

    # Сортируем по лучшему score и берём top_k
    top_indices = sorted(all_scores, key=all_scores.get, reverse=True)[:top_k]
    return [documents[i] for i in top_indices]

def reciprocal_rank_fusion(
    ranked_lists: List[List[int]],
    k: int = 60
) -> Dict[int, float]:
    """
    ranked_lists: список списков индексов документов, отсортированных по релевантности
    k: константа сглаживания (обычно 60)
    """
    scores: Dict[int, float] = {}

    for ranked in ranked_lists:
        for rank, doc_idx in enumerate(ranked):
            scores[doc_idx] = scores.get(doc_idx, 0) + 1 / (k + rank + 1)

    return scores


def multi_query_search_rrf(question: str, documents: List[Dict], top_k: int = 5) -> List[Dict]:
    queries = generate_queries(question, n=4)
    doc_embeddings = model.encode([d["content"] for d in documents])

    ranked_lists = []
    for query in queries:
        query_emb = model.encode([query])[0]
        sims = [
            np.dot(query_emb, d) / (np.linalg.norm(query_emb) * np.linalg.norm(d))
            for d in doc_embeddings
        ]
        ranked_lists.append(np.argsort(sims)[::-1].tolist())

    rrf_scores = reciprocal_rank_fusion(ranked_lists)
    top_indices = sorted(rrf_scores, key=rrf_scores.get, reverse=True)[:top_k]
    return [documents[i] for i in top_indices]