#!/usr/bin/env python3
"""RAG test - embed query, search, answer via Qwen."""
import json, os, chromadb, urllib.request

def embed(texts):
    data = json.dumps({"content": texts}).encode()
    req = urllib.request.Request("http://localhost:8081/embedding", data=data,
        headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=30) as resp:
        body = json.loads(resp.read())
    if isinstance(body, list):
        return [item["embedding"][0] for item in body]
    return body.get("embedding", [])

os.chdir(os.path.expanduser("~/nubes"))
client = chromadb.PersistentClient(path=os.path.expanduser("~/nubes/chroma_db"))
collection = client.get_collection("elsa_docs")

# 1. Embed query
query = "What is the tightening torque for fuel filter on VW Phaeton 3.2 VR6 2004?"
print(f"Query: {query}")
q_emb = embed([query])[0]

# 2. Search
results = collection.query(query_embeddings=[q_emb], n_results=3)
print(f"Matches: {len(results['ids'][0])}")
for i in range(len(results['ids'][0])):
    meta = results['metadatas'][0][i] or {}
    dist = results['distances'][0][i]
    print(f"\n[{i}] dist={dist:.3f}")
    print(f"  source: {meta.get('source','')}")
    print(f"  title: {meta.get('title','')[:120]}")
    print(f"  text: {results['documents'][0][i][:300]}...")

# 3. Build prompt for Qwen
context = "\n\n".join([f"[{m.get('title','')[:80] if m else ''}]\n{d[:1500]}" 
    for d, m in zip(results['documents'][0], results['metadatas'][0])])

prompt = f"""You are a VAG automotive diagnostician. Answer using ONLY the documentation below.
DOCUMENTATION:
{context}

Question: {query}
Answer:"""

data = json.dumps({"prompt": prompt, "temperature": 0.1, "n_predict": 300,
    "stop": ["\n\n", "Answer:", "\nThe answer", "\nDocumentation"]}).encode()
req = urllib.request.Request("http://localhost:8080/completion", data=data,
    headers={"Content-Type": "application/json"})
with urllib.request.urlopen(req, timeout=120) as resp:
    answer = json.loads(resp.read()).get("content", "")

print(f"\n\n=== LLM Answer ===\n{answer}")