from fastapi import FastAPI
from pydantic import BaseModel
from langchain_community.document_loaders import TextLoader
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
import os
import logging
import sys

logging.basicConfig(level=logging.INFO, stream=sys.stdout)
logger = logging.getLogger(__name__)

app = FastAPI()
vector_db = None

# Voyage-2 embeddings via OpenRouter API
embeddings = OpenAIEmbeddings(
    model="openai/text-embedding-3-small",
    openai_api_base="https://openrouter.ai/api/v1",
    openai_api_key=os.getenv("OPENROUTER_API_KEY")
)

@app.on_event("startup")
async def startup_event():
    global vector_db
    data_path = "./data/hobbies.md"
    if os.path.exists(data_path):
        try:
            loader = TextLoader(data_path)
            documents = loader.load()
            text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
            chunks = text_splitter.split_documents(documents)
            vector_db = Chroma.from_documents(documents=chunks, embedding=embeddings, persist_directory="./chroma_db")
            logger.info("Librarian: ChromaDB is loaded with openAi embeddings.")
        except Exception as e:
            logger.error(f"Librarian: DB error: {str(e)}")
    else:
        logger.warning(f"Librarian: Missing data file at {data_path}")

@app.get("/health")
async def health():
    return {"status": "ready", "vectors_loaded": vector_db is not None}

class QueryRequest(BaseModel):
    question: str

@app.post("/query")
async def query_knowledge(request: QueryRequest):
    if not vector_db: return {"context": ""}
    results = vector_db.similarity_search(request.question, k=2)
    return {"context": "\n".join([res.page_content for res in results])}