"""
CPU RAG Space — bge-small (fastembed) + FAISS + Qwen3.5-0.8B (llama.cpp),
served as an OpenAI-compatible API with a small web UI.

Everything runs on CPU and fits the Hugging Face free tier (2 vCPU / 16 GB).
Tuned for lowest CPU latency: small MoE model, 2K context, flash-attention,
short answers, relevance-gated RAG, and a startup warm-up so the first real
query is snappy. Needs llama-cpp-python >= 0.3.32 (qwen35 arch).
"""

import glob
import json
import os

import faiss
import numpy as np
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
from fastembed import TextEmbedding
from llama_cpp import Llama
from pydantic import BaseModel

# --------------------------------------------------------------------------- #
# Config (all overridable via Space "Variables")
# --------------------------------------------------------------------------- #
MODEL_DIR = os.environ.get("MODEL_DIR", "models")
LLM_FILE = os.environ.get("LLM_FILE", "Qwen3.5-0.8B-Q4_K_M.gguf")
LLM_PATH = os.path.join(MODEL_DIR, LLM_FILE)
EMBED_MODEL = os.environ.get("EMBED_MODEL", "BAAI/bge-small-en-v1.5")
FASTEMBED_CACHE = os.environ.get("FASTEMBED_CACHE")
# RAG prompts are ~500 tokens, so 2K context is plenty. A smaller KV cache means
# less RAM and faster per-token cache ops than the old 8192.
N_CTX = int(os.environ.get("N_CTX", "2048"))
N_THREADS = int(os.environ.get("N_THREADS", str(os.cpu_count() or 2)))
TOP_K = int(os.environ.get("TOP_K", "4"))
DOCS_DIR = os.environ.get("DOCS_DIR", "documents")
# Only inject retrieved context when the best hit is actually relevant (cosine
# similarity). Below this, treat it as normal chat instead of forcing doc QA —
# stops greetings/off-topic messages from dredging up irrelevant chunks.
RAG_MIN_SCORE = float(os.environ.get("RAG_MIN_SCORE", "0.4"))

CHUNK_SIZE = 800        # characters per chunk
CHUNK_OVERLAP = 120

# Plain assistant persona used when nothing relevant is retrieved.
CHAT_SYSTEM = "You are a friendly, concise assistant."

# Used when we DO have relevant context. No literal "[filename]" example — a
# tiny model will just echo it. Sources are returned separately in the JSON.
RAG_SYSTEM = (
    "You are a helpful assistant. Use the context below to answer the user's "
    "question. If the context does not contain the answer, say so briefly and "
    "answer from your own knowledge if you can. Keep answers concise.\n\n"
    "Context:\n{context}"
)

# --------------------------------------------------------------------------- #
# Lazily-initialised singletons
# --------------------------------------------------------------------------- #
app = FastAPI(title="CPU RAG Space")

_embedder = None
_llm = None
_index = None            # faiss.IndexFlatIP
_chunks = []             # list[{"text": str, "source": str}]


def embedder():
    global _embedder
    if _embedder is None:
        _embedder = TextEmbedding(EMBED_MODEL, cache_dir=FASTEMBED_CACHE)
    return _embedder


def embed(texts):
    vecs = np.array(list(embedder().embed(list(texts))), dtype="float32")
    faiss.normalize_L2(vecs)          # cosine similarity via inner product
    return vecs


def llm():
    global _llm
    if _llm is None:
        _llm = Llama(
            model_path=LLM_PATH,
            n_ctx=N_CTX,
            n_threads=N_THREADS,          # decode threads (= physical cores)
            n_threads_batch=N_THREADS,    # prefill threads
            n_batch=512,
            flash_attn=True,              # cheaper attention -> faster on CPU
            verbose=False,
        )
    return _llm


# --------------------------------------------------------------------------- #
# Indexing / retrieval
# --------------------------------------------------------------------------- #
def chunk_text(text, source):
    out, i, n = [], 0, len(text)
    step = max(CHUNK_SIZE - CHUNK_OVERLAP, 1)
    while i < n:
        piece = text[i:i + CHUNK_SIZE].strip()
        if piece:
            out.append({"text": piece, "source": source})
        i += step
    return out


def add_chunks(new_chunks):
    global _index, _chunks
    if not new_chunks:
        return 0
    vecs = embed([c["text"] for c in new_chunks])
    if _index is None:
        _index = faiss.IndexFlatIP(vecs.shape[1])
    _index.add(vecs)
    _chunks.extend(new_chunks)
    return len(new_chunks)


def build_index():
    patterns = ("*.txt", "*.md")
    files = []
    for p in patterns:
        files += glob.glob(os.path.join(DOCS_DIR, "**", p), recursive=True)
    all_chunks = []
    for f in files:
        try:
            with open(f, encoding="utf-8") as fh:
                all_chunks += chunk_text(fh.read(), os.path.basename(f))
        except Exception as exc:
            print(f"[rag] skip {f}: {exc}")
    add_chunks(all_chunks)


def retrieve(query, k=TOP_K):
    if _index is None or _index.ntotal == 0:
        return []
    scores, ids = _index.search(embed([query]), min(k, _index.ntotal))
    hits = []
    for score, idx in zip(scores[0], ids[0]):
        if idx < 0:
            continue
        c = _chunks[idx]
        hits.append({"text": c["text"], "source": c["source"], "score": float(score)})
    return hits


# --------------------------------------------------------------------------- #
# Startup
# --------------------------------------------------------------------------- #
@app.on_event("startup")
def _startup():
    print("[rag] loading embedder + llm ...")
    embedder()
    llm()
    build_index()
    # Warm up the decode path (kernels, KV cache) so the FIRST real user query
    # doesn't pay the one-off jit/allocation cost.
    try:
        llm().create_chat_completion(
            messages=[{"role": "user", "content": "hi"}], max_tokens=1)
    except Exception as exc:
        print(f"[rag] warmup skipped: {exc}")
    print(f"[rag] ready. indexed_chunks={len(_chunks)}")


# --------------------------------------------------------------------------- #
# OpenAI-compatible chat endpoint (with RAG)
# --------------------------------------------------------------------------- #
class ChatMessage(BaseModel):
    role: str
    content: str


class ChatRequest(BaseModel):
    model: str = "cpu-rag"
    messages: list[ChatMessage]
    temperature: float = 0.3
    top_p: float = 0.9
    max_tokens: int = 256      # short RAG answers -> lower latency on CPU
    stream: bool = False       # API default off for OpenAI-client compat; UI streams
    use_rag: bool = True


def _augment(req: ChatRequest):
    msgs = [m.model_dump() for m in req.messages]
    users = [m for m in msgs if m["role"] == "user"]
    query = users[-1]["content"] if users else ""

    # Retrieve, then keep only chunks that clear the relevance bar.
    hits = retrieve(query) if req.use_rag else []
    ctxs = [c for c in hits if c["score"] >= RAG_MIN_SCORE]

    non_system = [m for m in msgs if m["role"] != "system"]
    if ctxs:
        context = "\n\n".join(f"[{c['source']}] {c['text']}" for c in ctxs)
        system = {"role": "system", "content": RAG_SYSTEM.format(context=context)}
    else:
        # Nothing relevant -> behave like a normal chat assistant, no forced
        # "answer only from context" (which made the model spit citations).
        system = {"role": "system", "content": CHAT_SYSTEM}
    return [system] + non_system, ctxs


@app.post("/v1/chat/completions")
def chat_completions(req: ChatRequest):
    messages, ctxs = _augment(req)
    params = dict(messages=messages, temperature=req.temperature, top_p=req.top_p,
                  max_tokens=req.max_tokens, stop=["<|im_end|>", "<|endoftext|>"])

    if req.stream:
        def gen():
            for chunk in llm().create_chat_completion(**params, stream=True):
                yield f"data: {json.dumps(chunk)}\n\n"
            yield "data: [DONE]\n\n"
        return StreamingResponse(gen(), media_type="text/event-stream")

    resp = llm().create_chat_completion(**params)
    resp["sources"] = [{"source": c["source"], "score": round(c["score"], 3)} for c in ctxs]
    return JSONResponse(resp)


# --------------------------------------------------------------------------- #
# Ingest / stats / UI
# --------------------------------------------------------------------------- #
@app.post("/ingest")
async def ingest(file: UploadFile = File(...)):
    text = (await file.read()).decode("utf-8", "ignore")
    added = add_chunks(chunk_text(text, file.filename))
    return {"file": file.filename, "added_chunks": added, "total_chunks": len(_chunks)}


@app.get("/stats")
def stats():
    return {"indexed_chunks": len(_chunks), "embed_model": EMBED_MODEL,
            "llm": LLM_FILE, "n_ctx": N_CTX, "threads": N_THREADS, "top_k": TOP_K}


@app.get("/", response_class=HTMLResponse)
def home():
    with open(os.path.join(os.path.dirname(__file__), "index.html"), encoding="utf-8") as f:
        return f.read()