# main.py
from qdrant_client.http import models as qmodels
import uuid
from fastapi import FastAPI
import requests, os, redis
from qdrant_client import QdrantClient
from langchain_community.vectorstores import Qdrant
from langchain_community.embeddings import HuggingFaceEmbeddings

app = FastAPI()

# Connectors
qdrant = QdrantClient(url=os.getenv("QDRANT_URL"))
r = redis.Redis.from_url(os.getenv("REDIS_URL"))
SEARXNG_URL = os.getenv("SEARXNG_URL")
LLAMA_SERVER_URL = os.getenv("LLAMA_SERVER_URL")
QDRANT_URL = os.getenv("QDRANT_URL")
REDIS_URL = os.getenv("REDIS_URL")

embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

@app.get("/ask")
def ask(query: str):
    # 1. Search via SearxNG
    resp = requests.get(f"{SEARXNG_URL}/search", params={"q": query, "format": "json"})
    snippets = [r["title"] + " " + r["content"] for r in resp.json()["results"][:3]]

    context = " ".join(snippets)

    # 2. Embed + store in Qdrant
    vectors = embeddings.embed_documents(snippets)
    points = [
        qmodels.PointStruct(
            id=str(uuid.uuid4()),          # unique ID for each snippet
            vector=vec,                    # the embedding vector
            payload={"text": text}         # optional metadata
        )
        for text, vec in zip(snippets, vectors)
    ]
    qdrant.upsert(collection_name="docs", points=points)

    # 3. Call llama-server
    llama_http = requests.post(
        f"{LLAMA_SERVER_URL}/completion",
        json={
            "prompt": f"Context:\n{context}\n\nQuestion: {query}\nAnswer:",
            "n_predict": 128,
            "temperature": 0.7,
            "stop": ["</s>"]
        }
    )
    llm_resp = llama_http.json()
    print("Llama raw response:", llm_resp)

    return {"answer": llm_resp.get("content", "")}
@app.get("/health")
def health():
    results = {}

    # Check Redis
    try:
        r = redis.Redis.from_url(REDIS_URL)
        r.ping()
        results["redis"] = "ok"
    except Exception as e:
        results["redis"] = f"error: {e}"

    # Check Qdrant
    try:
        q = requests.get(f"{QDRANT_URL}/readyz")
        results["qdrant"] = f"ok ({q.status_code})"
    except Exception as e:
        results["qdrant"] = f"error: {e}"

    # Check SearxNG
    try:
        s = requests.get(f"{SEARXNG_URL}/search",
                         params={"q": "ping", "format": "json"},
                         headers={"X-Forwarded-For": "127.0.0.1"})
        results["searxng"] = f"ok ({s.status_code})"
    except Exception as e:
        results["searxng"] = f"error: {e}"

    return results