"""Kleines lokales RAG ohne Framework.

Vorher: ollama pull embeddinggemma
         ollama pull gemma3
Dann:    python rag.py index
         python rag.py ask "Wie lange dauert der Versand?"
"""

import json
import math
import sys
import urllib.request
from pathlib import Path

DOCS = Path("wissen")
INDEX = Path("index.json")
OLLAMA = "http://localhost:11434"
EMBED_MODEL = "embeddinggemma"
CHAT_MODEL = "gemma3"


def post(path, data):
    """Sendet JSON an Ollama und gibt dessen JSON-Antwort zurück."""
    request = urllib.request.Request(
        OLLAMA + path,
        data=json.dumps(data).encode(),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(request) as response:
        return json.load(response)


def embed(text):
    """Macht aus Text eine Zahlenreihe, die wir vergleichen können."""
    return post("/api/embed", {"model": EMBED_MODEL, "input": text})["embeddings"][0]


def cosine(a, b):
    """Gibt zurück, wie ähnlich zwei Zahlenreihen ausgerichtet sind."""
    dot = sum(x * y for x, y in zip(a, b))
    length_a = math.sqrt(sum(x * x for x in a))
    length_b = math.sqrt(sum(y * y for y in b))
    return dot / (length_a * length_b) if length_a and length_b else 0


def chunks():
    """Liest Markdown-Dateien; jeder Absatz wird ein Chunk."""
    for file in DOCS.glob("*.md"):
        for paragraph in file.read_text(encoding="utf-8").split("\n\n"):
            text = paragraph.strip()
            if text:
                yield {"source": file.name, "text": text}


def build_index():
    """Speichert jeden Chunk zusammen mit seinem Embedding in index.json."""
    entries = []
    for chunk in chunks():
        chunk["embedding"] = embed(chunk["text"])
        entries.append(chunk)
        print("indexiert:", chunk["source"])
    INDEX.write_text(json.dumps(entries, ensure_ascii=False), encoding="utf-8")
    print(f"Fertig: {len(entries)} Chunks in {INDEX}")


def search(question, limit=3):
    """Findet die drei Chunks, deren Embedding der Frage am ähnlichsten ist."""
    question_embedding = embed(question)
    entries = json.loads(INDEX.read_text(encoding="utf-8"))
    return sorted(
        entries,
        key=lambda entry: cosine(question_embedding, entry["embedding"]),
        reverse=True,
    )[:limit]


def answer(question):
    """Gibt die Fundstellen mit der Frage an das Chat-Modell."""
    found = search(question)
    context = "\n\n".join(f"[{entry['source']}]\n{entry['text']}" for entry in found)
    prompt = f"""Beantworte die Frage nur mit dem folgenden Kontext.
Wenn die Antwort nicht im Kontext steht, schreibe das deutlich.
Nenne die Quelldatei in Klammern.

KONTEXT:
{context}

FRAGE: {question}"""
    result = post("/api/chat", {
        "model": CHAT_MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "stream": False,
    })
    print(result["message"]["content"])


if __name__ == "__main__":
    if sys.argv[1:] == ["index"]:
        build_index()
    elif len(sys.argv) > 2 and sys.argv[1] == "ask":
        answer(" ".join(sys.argv[2:]))
    else:
        print('Nutze: python rag.py index | python rag.py ask "Frage"')