import os
from fastapi import FastAPI
from langchain_community.utilities import SearxSearchWrapper
from langchain_community.vectorstores import Qdrant
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import LLaMAServer
from langchain.schema import Document

# Environment variables
SEARXNG_URL = os.environ["SEARXNG_URL"]
QDRANT_URL = os.environ["QDRANT_URL"]
LLAMA_SERVER_URL = os.environ["LLAMA_SERVER_URL"]

# Initialize FastAPI
app = FastAPI()

# Initialize components
searx = SearxSearchWrapper(searx_host=SEARXNG_URL)
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
qdrant = Qdrant(
    url=QDRANT_URL,
    collection_name="search_vectors",
    embeddings=embedding_model
)
llm = LLaMAServer(endpoint_url=LLAMA_SERVER_URL)

@app.post("/process")
async def process(query: str):
    # Step 1: Search with SearXNG
    raw_results = searx.run(query)

    # Step 2: Convert to LangChain Documents
    docs = [Document(page_content=result) for result in raw_results.split("\n") if result.strip()]

    # Step 3: Embed and store in Qdrant
    qdrant.add_documents(docs)

    # Step 4: Query Qdrant for similar vectors
    matches = qdrant.similarity_search(query, k=3)

    # Step 5: Combine raw + vector results
    combined = "\n".join([doc.page_content for doc in docs + matches])

    # Step 6: Send to LLaMA-Server
    response = llm.invoke(combined)

    return {"response": response}