"""
CPU RAG Space — bge-small (fastembed) + FAISS + Qwen2.5-1.5B (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).
"""

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", "qwen2.5-1.5b-instruct-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")
N_CTX = int(os.environ.get("N_CTX", "8192"))
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")

CHUNK_SIZE = 800        # characters per chunk
CHUNK_OVERLAP = 120

RAG_SYSTEM = (
    "You are a helpful assistant. Answer the user's question using ONLY the "
    "context below. If the answer is not in the context, say you don't know. "
    "Cite the source of each fact in square brackets like [filename].\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,
                     n_batch=512, 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()
    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 = 512
    stream: bool = False
    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 ""
    ctxs = retrieve(query) if req.use_rag else []
    if ctxs:
        context = "\n\n".join(f"[{c['source']}] {c['text']}" for c in ctxs)
        system = {"role": "system", "content": RAG_SYSTEM.format(context=context)}
        msgs = [system] + [m for m in msgs if m["role"] != "system"]
    return msgs, 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()