#!/usr/bin/env python3
"""RAG Query — search ChromaDB + answer via llama.cpp."""
import json, os, sys, urllib.request, time
import chromadb

LLAMA_URL = "http://localhost:8080/completion"

def llama_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=60) as resp:
        body = json.loads(resp.read())
    if isinstance(body, list):
        return [item["embedding"][0] for item in body]
    return body.get("embedding", [])


class LlamaEmbedFn:
    def __call__(self, input):
        all_embs = []
        for i in range(0, len(input), 32):
            batch = input[i:i+32]
            embs = llama_embed(batch)
            all_embs.extend(embs if isinstance(embs[0], list) else [embs])
        return all_embs


def search(collection, query, n=5):
    results = collection.query(query_texts=[query], n_results=n)
    docs = []
    for i in range(len(results["ids"][0])):
        docs.append({
            "id": results["ids"][0][i],
            "doc": results["documents"][0][i],
            "meta": results["metadatas"][0][i],
            "distance": results["distances"][0][i],
        })
    return docs


def ask_llama(prompt, max_tokens=256):
    data = json.dumps({
        "prompt": prompt, "temperature": 0.1,
        "n_predict": max_tokens,
        "stop": ["\n\n", "Documentation:", "Question:"],
    }).encode()
    req = urllib.request.Request(LLAMA_URL, data=data,
        headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=120) as resp:
        return json.loads(resp.read()).get("content", "")


def main():
    db_dir = os.path.expanduser("~/nubes/chroma_db")
    collection_name = "elsa_docs"
    
    ef = LlamaEmbedFn()
    
    client = chromadb.PersistentClient(path=db_dir)
    collection = client.get_collection(collection_name, embedding_function=ef)
    
    print(f"Collection '{collection_name}': {collection.count()} documents")
    
    while True:
        try:
            query = input("\n>>> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        
        if not query:
            continue
        if query.lower() in ("exit", "quit", "q"):
            break
        
        # 1. Поиск
        results = search(collection, query, n_results=5)
        
        # 2. Формируем контекст
        context_parts = []
        for r in results:
            title = r["meta"].get("title", "")[:120]
            ctx = f"[{title}] {r['doc'][:1000]}"
            context_parts.append(ctx)
        
        context = "\n\n---\n\n".join(context_parts)
        
        # 3. Промпт
        prompt = f"""You are an expert VAG (VW/Audi/Skoda/SEAT) automotive diagnostician. 
Answer the user's question using ONLY the provided repair documentation below.
If the answer is not in the documentation, say "Not found in documentation."
Answer in Russian.

DOCUMENTATION:
{context}

Question: {query}
Answer:"""
        
        print(f"\n[Found {len(results)} matches, sending to LLM...]")
        answer = ask_llama(prompt)
        print(f"\n{answer}")


if __name__ == "__main__":
    main()