"""
Krok 2 pipeline'u RAG: zadaj pytanie, znajdź najlepiej pasujące fragmenty
OWU, i (jeśli jest ustawiony klucz API) poproś Claude o odpowiedź opartą
wyłącznie na tych fragmentach, z obowiązkowym cytowaniem paragrafu.

Uruchomienie:
    python src/query.py "Czy AC obejmuje szkodę na parkingu?"

Bez klucza ANTHROPIC_API_KEY w .env skrypt i tak pokaże, które fragmenty
zostały odnalezione -- to wystarczy, żeby ocenić jakość samego
wyszukiwania, zanim zapłacisz za pierwsze zapytanie do modelu.
"""

import os
import pickle
import sys
from pathlib import Path

from dotenv import load_dotenv
from sklearn.metrics.pairwise import cosine_similarity

INDEX_PATH = Path(__file__).parent.parent / ".index" / "index.pkl"
TOP_K = 5

SYSTEM_PROMPT = """Jesteś asystentem agenta ubezpieczeniowego. Odpowiadasz
WYŁĄCZNIE na podstawie dostarczonych fragmentów OWU -- nie korzystasz z
własnej wiedzy o produktach ubezpieczeniowych. Zasady:

1. Każde twierdzenie musi się powoływać na konkretny fragment (podaj etykietę
   paragrafu i nazwę pliku źródłowego dokładnie tak, jak podano we fragmentach).
2. Jeśli dostarczone fragmenty nie zawierają odpowiedzi, napisz to wprost --
   nie zgaduj i nie uzupełniaj wiedzą spoza fragmentów.
3. Zakończ odpowiedź zdaniem: "AI może się mylić -- zweryfikuj z oryginałem OWU."
4. Nie podejmuj decyzji za agenta (np. "wypłać odszkodowanie") -- tylko
   informuj, co wynika z dokumentu.
"""


def load_index():
    if not INDEX_PATH.exists():
        raise SystemExit("Brak indeksu. Uruchom najpierw: python src/ingest.py")
    with open(INDEX_PATH, "rb") as f:
        return pickle.load(f)


def retrieve(question: str, index: dict, top_k: int = TOP_K):
    vectorizer = index["vectorizer"]
    matrix = index["matrix"]
    chunks = index["chunks"]

    q_vec = vectorizer.transform([question])
    scores = cosine_similarity(q_vec, matrix).flatten()
    ranked = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:top_k]
    return [(chunks[i], float(scores[i])) for i in ranked if scores[i] > 0]


def build_context(results) -> str:
    parts = []
    for chunk, score in results:
        parts.append(
            f"[Źródło: {chunk.source_file}, {chunk.paragraph_label}, "
            f"trafność {score:.2f}]\n{chunk.text}"
        )
    return "\n\n---\n\n".join(parts)


def generate_answer(question: str, context: str) -> str:
    load_dotenv()
    api_key = os.getenv("ANTHROPIC_API_KEY")
    if not api_key:
        return (
            "[Brak ANTHROPIC_API_KEY w .env -- pokazuję tylko odnalezione "
            "fragmenty (patrz wyżej). Dodaj klucz w .env, żeby zobaczyć "
            "wygenerowaną odpowiedź.]"
        )

    from anthropic import Anthropic

    client = Anthropic(api_key=api_key)
    model = os.getenv("ANTHROPIC_MODEL", "claude-sonnet-4-5")

    response = client.messages.create(
        model=model,
        max_tokens=1024,
        system=SYSTEM_PROMPT,
        messages=[{
            "role": "user",
            "content": f"Fragmenty OWU:\n\n{context}\n\nPytanie agenta: {question}",
        }],
    )
    return response.content[0].text


def main() -> None:
    if len(sys.argv) < 2:
        raise SystemExit('Użycie: python src/query.py "treść pytania"')

    question = sys.argv[1]
    index = load_index()
    results = retrieve(question, index)

    print(f"Pytanie: {question}\n")
    if not results:
        print("Nie znaleziono żadnego pasującego fragmentu w bazie OWU.")
        return

    print(f"Znalezione fragmenty ({len(results)}):\n")
    for chunk, score in results:
        print(f"- {chunk.source_file} | {chunk.paragraph_label} | trafność {score:.2f}")
    print()

    context = build_context(results)
    answer = generate_answer(question, context)
    print("Odpowiedź:\n")
    print(answer)


if __name__ == "__main__":
    main()