"""OmniBrain API v3.2 — real routes over the exported brain.

Boot: pull parquet shards from the BRAIN_REPO dataset repo, load an fp32 matrix
in RAM, serve brute-force cosine. Space disk is ephemeral, so the dataset repo
is the durable store; restarting the Space just re-pulls.

v3.2: (1) matrix kept fp32 at load — v3.1 re-cast the whole fp16 matrix to
fp32 on EVERY query (~186MB alloc per request). (2) `scope` splits the personal
layer from bulk public reference rows, which otherwise outnumber it ~7:1 and
crowd real answers out of the top-k.
"""
import os, glob
from collections import Counter

import numpy as np
import pyarrow.parquet as pq
from fastapi import FastAPI
from pydantic import BaseModel
from huggingface_hub import snapshot_download
from sentence_transformers import SentenceTransformer

os.environ.setdefault("HF_HOME", "/tmp/hf")
REPO_ID = os.environ.get("BRAIN_REPO", "Fonedo/omnibrain-caleb-brain")

# Bulk public reference sets. Everything else (memories, vault, project docs,
# agents, skills) is Caleb's own knowledge and is what search should hit first.
PUBLIC_SOURCES = {
    "teknium/OpenHermes-2.5", "openbmb/UltraFeedback",
    "nvidia/OpenCodeReasoning", "nvidia/OpenCodeReasoning-2",
    "ise-uiuc/Magicoder-OSS-Instruct-75K", "princeton-nlp/SWE-bench",
    "zai-org/AgentInstruct", "markov-ai/gaming-500-hours",
    "HuggingFaceFW/fineweb-edu",
}

app = FastAPI(title="OmniBrain")
S = {"ready": False, "error": None}

@app.on_event("startup")
def load():
    try:
        d = snapshot_download(REPO_ID, repo_type="dataset",
                              token=os.environ.get("HF_TOKEN"))
        vecs, meta = [], []
        for f in sorted(glob.glob(f"{d}/*.parquet")):
            t = pq.read_table(f)
            vecs.append(np.asarray(t["vector"].combine_chunks().flatten(),
                                   dtype=np.float32).reshape(-1, 768))
            meta.extend(zip(t["project"].to_pylist(), t["source"].to_pylist(),
                            t["type"].to_pylist(), t["text"].to_pylist()))
        # fp32 once, not per query
        S["V"] = np.vstack(vecs) if vecs else np.zeros((0, 768), np.float32)
        S["M"] = meta
        S["IS_PUB"] = np.array([m[1] in PUBLIC_SOURCES for m in meta], dtype=bool)
        S["model"] = SentenceTransformer("nomic-ai/nomic-embed-text-v1.5",
                                         trust_remote_code=True, device="cpu")
        S["model"].max_seq_length = 512
        S["ready"] = True
        n_pub = int(S["IS_PUB"].sum())
        print(f"Loaded {len(meta):,} vectors "
              f"({len(meta)-n_pub:,} personal / {n_pub:,} public) from {REPO_ID}")
    except Exception as e:
        S["error"] = str(e)[:500]
        print("STARTUP FAILED:", S["error"])

@app.get("/health")
def health():
    return {"status": "online", "version": "3.2", "ready": S["ready"],
            "vectors": len(S.get("M", [])), "error": S["error"]}

@app.get("/stats")
def stats():
    c = Counter(m[0] for m in S.get("M", []))
    t = Counter(m[2] for m in S.get("M", []))
    pub = int(S["IS_PUB"].sum()) if S.get("ready") else 0
    return {"vectors": len(S.get("M", [])),
            "personal": len(S.get("M", [])) - pub, "public": pub,
            "by_project": dict(c), "by_type": dict(t), "dim": 768,
            "model": "nomic-ai/nomic-embed-text-v1.5", "brain_repo": REPO_ID}

class Q(BaseModel):
    query: str
    k: int = 8
    project: str | None = None
    agent: str | None = None          # kimi | claude | codex — biases type
    scope: str = "personal"           # personal | all | public

# each AI gets a different slice of the brain, per the OmniBrain design
AGENT_BIAS = {
    "kimi":   {"doc": 1.15, "vault": 1.15, "memory": 1.1},                # architecture
    "claude": {"memory": 1.15, "vault": 1.15, "agent": 1.1, "skill": 1.1},# systems/art
    "codex":  {"code": 1.2, "doc": 1.05},                                 # code
}

def _search(q: str, k: int, project: str | None, agent: str | None,
            scope: str = "personal"):
    if not len(S["M"]):
        return []
    # asymmetric model: docs were embedded with "search_document: "
    v = S["model"].encode(["search_query: " + q], normalize_embeddings=True)[0]
    sims = S["V"] @ v.astype(np.float32)

    if scope == "personal":
        sims = np.where(S["IS_PUB"], -1e9, sims)
    elif scope == "public":
        sims = np.where(S["IS_PUB"], sims, -1e9)

    if project:
        mask = np.array([m[0] == project or m[0] == "meta" for m in S["M"]])
        sims = np.where(mask, sims, -1e9)
    if agent and agent in AGENT_BIAS:
        b = AGENT_BIAS[agent]
        boost = np.array([b.get(m[2], 1.0) for m in S["M"]], dtype=np.float32)
        sims = np.where(sims > -1e8, sims * boost, sims)

    k = max(1, min(k, len(sims)))
    idx = np.argpartition(-sims, k - 1)[:k]
    idx = idx[np.argsort(-sims[idx])]
    return [{"score": round(float(sims[i]), 4), "project": S["M"][i][0],
             "source": S["M"][i][1], "type": S["M"][i][2],
             "text": S["M"][i][3]} for i in idx if sims[i] > -1e8]

@app.post("/search")
def search(q: Q):
    if not S["ready"]: return {"error": S["error"] or "still loading"}
    return {"results": _search(q.query, q.k, q.project, q.agent, q.scope)}

@app.post("/generate_context")
def generate_context(q: Q):
    if not S["ready"]: return {"error": S["error"] or "still loading"}
    hits = _search(q.query, q.k, q.project, q.agent, q.scope)
    ctx = "\n\n---\n".join(
        f"[{h['project']}/{h['type']}] {h['text']}" for h in hits)
    return {"context": ctx, "n": len(hits)}

@app.get("/")
def root():
    return {"message": "OmniBrain API",
            "routes": ["/health", "/stats", "/search", "/generate_context"],
            "scopes": ["personal (default)", "all", "public"]}