"""Thin FastAPI wrapper around langchain_postgres.PGVector, replacing the
vendored litellm-pgvector connector. See
docs/research/langchain-pgvector-vs-litellm-pgvector.md for the rationale
and docs/memory-knowledgebase.md for usage.

ponytail: one fixed collection (COLLECTION_NAME) — this stack only needs one
knowledgebase ("memory-and-notes"), not a multi-tenant store registry.
"""
import os

from fastapi import Depends, FastAPI, HTTPException
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from langchain_openai import OpenAIEmbeddings
from langchain_postgres import PGVector
from pydantic import BaseModel

DATABASE_URL = os.environ["DATABASE_URL"]
LITELLM_BASE_URL = os.environ["LITELLM_BASE_URL"]
LITELLM_API_KEY = os.environ["LITELLM_API_KEY"]
EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "local-embedding")
SERVER_API_KEY = os.environ["SERVER_API_KEY"]
COLLECTION_NAME = os.environ.get("COLLECTION_NAME", "memory-and-notes")

app = FastAPI(title="memory-retrieval", version="1.0.0")

embeddings = OpenAIEmbeddings(
    model=EMBEDDING_MODEL, base_url=LITELLM_BASE_URL, api_key=LITELLM_API_KEY
)
vector_store = PGVector(
    embeddings=embeddings,
    collection_name=COLLECTION_NAME,
    connection=DATABASE_URL,
    use_jsonb=True,
)

security = HTTPBearer()


def check_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
    if credentials.credentials != SERVER_API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key")


class Chunk(BaseModel):
    content: str
    metadata: dict = {}


class IngestRequest(BaseModel):
    chunks: list[Chunk]


class QueryRequest(BaseModel):
    query: str
    k: int = 5


@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/ingest", dependencies=[Depends(check_api_key)])
def ingest(req: IngestRequest):
    texts = [c.content for c in req.chunks]
    metadatas = [c.metadata for c in req.chunks]
    ids = vector_store.add_texts(texts=texts, metadatas=metadatas)
    return {"ingested": len(ids)}


@app.post("/query", dependencies=[Depends(check_api_key)])
def query(req: QueryRequest):
    results = vector_store.similarity_search_with_relevance_scores(req.query, k=req.k)
    return {
        "results": [
            {"content": doc.page_content, "metadata": doc.metadata, "score": score}
            for doc, score in results
        ]
    }