main.py

import math
import random
import numpy as np
import time
import threading
import heapq
import json
import requests
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Tuple, Callable, Generator
from fastapi import FastAPI, Request, Response, UploadFile, File
from fastapi.responses import JSONResponse, FileResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
import os
import fitz  # PyMuPDF

# =====================================================================
#  CONSTANTS
# =====================================================================

DIMS = 16  # demo vectors dimension
SAVE_FILE = "doc_db_save.json"  # persistence file for DocumentDB

# =====================================================================
#  DATA TYPES
# =====================================================================

@dataclass
class VectorItem:
    id: int
    metadata: str
    category: str
    emb: List[float]


DistFn = Callable[[List[float], List[float]], float]


# =====================================================================
#  DISTANCE METRICS
# =====================================================================

def euclidean(a: List[float], b: List[float]) -> float:
    return float(np.linalg.norm(np.array(a) - np.array(b)))


def cosine(a: List[float], b: List[float]) -> float:
    a_ = np.array(a)
    b_ = np.array(b)
    na = np.linalg.norm(a_)
    nb = np.linalg.norm(b_)
    if na < 1e-9 or nb < 1e-9:
        return 1.0
    return float(1.0 - np.dot(a_, b_) / (na * nb))


def manhattan(a: List[float], b: List[float]) -> float:
    return float(np.sum(np.abs(np.array(a) - np.array(b))))


def get_dist_fn(m: str) -> DistFn:
    if m == "cosine":
        return cosine
    if m == "manhattan":
        return manhattan
    return euclidean


# =====================================================================
#  BRUTE FORCE
# =====================================================================

class BruteForce:
    def __init__(self):
        self.items: List[VectorItem] = []

    def insert(self, v: VectorItem):
        self.items.append(v)

    def knn(self, q: List[float], k: int, dist: DistFn) -> List[Tuple[float, int]]:
        r = [(dist(q, v.emb), v.id) for v in self.items]
        r.sort()
        return r[:k]

    def remove(self, id: int):
        self.items = [v for v in self.items if v.id != id]


# =====================================================================
#  KD-TREE
# =====================================================================

class KDNode:
    def __init__(self, item: VectorItem):
        self.item = item
        self.left: Optional['KDNode'] = None
        self.right: Optional['KDNode'] = None


class KDTree:
    def __init__(self, dims: int):
        self.root: Optional[KDNode] = None
        self.dims = dims

    def _ins(self, n: Optional[KDNode], v: VectorItem, d: int) -> KDNode:
        if n is None:
            return KDNode(v)
        ax = d % self.dims
        if v.emb[ax] < n.item.emb[ax]:
            n.left = self._ins(n.left, v, d + 1)
        else:
            n.right = self._ins(n.right, v, d + 1)
        return n

    def insert(self, v: VectorItem):
        self.root = self._ins(self.root, v, 0)

    def _knn(self, n: Optional[KDNode], q: List[float], k: int, d: int,
             dist: DistFn, heap: list):
        if n is None:
            return
        dn = dist(q, n.item.emb)
        # Max-heap: store negative distances
        if len(heap) < k or dn < -heap[0][0]:
            heapq.heappush(heap, (-dn, n.item.id))
            if len(heap) > k:
                heapq.heappop(heap)
        ax = d % self.dims
        diff = q[ax] - n.item.emb[ax]
        closer = n.left if diff < 0 else n.right
        farther = n.right if diff < 0 else n.left
        self._knn(closer, q, k, d + 1, dist, heap)
        if len(heap) < k or abs(diff) < -heap[0][0]:
            self._knn(farther, q, k, d + 1, dist, heap)

    def knn(self, q: List[float], k: int, dist: DistFn) -> List[Tuple[float, int]]:
        heap = []
        self._knn(self.root, q, k, 0, dist, heap)
        r = [(-neg_d, id_) for neg_d, id_ in heap]
        r.sort()
        return r

    def rebuild(self, items: List[VectorItem]):
        self.root = None
        for v in items:
            self.insert(v)


# =====================================================================
#  HNSW — Hierarchical Navigable Small World
# =====================================================================

@dataclass
class HNSWNode:
    item: VectorItem
    max_lyr: int
    nbrs: List[List[int]] = field(default_factory=list)


class HNSW:
    def __init__(self, m: int = 16, ef_build: int = 200):
        self.M = m
        self.M0 = 2 * m
        self.ef_build = ef_build
        self.mL = 1.0 / math.log(float(m))
        self.top_layer = -1
        self.entry_pt = -1
        self.G: Dict[int, HNSWNode] = {}
        self._rng = random.Random(42)

    def _rand_level(self) -> int:
        u = self._rng.random()
        return int(math.floor(-math.log(u) * self.mL))

    def _search_layer(self, q: List[float], ep: int, ef: int, lyr: int,
                      dist: DistFn) -> List[Tuple[float, int]]:
        vis = {ep: True}
        # Min-heap for candidates
        d0 = dist(q, self.G[ep].item.emb)
        cands = [(d0, ep)]
        # Max-heap for found (store negatives)
        found = [(-d0, ep)]

        while cands:
            cd, cid = heapq.heappop(cands)
            best_found_dist = -found[0][0]
            if len(found) >= ef and cd > best_found_dist:
                break
            if lyr >= len(self.G[cid].nbrs):
                continue
            for nid in self.G[cid].nbrs[lyr]:
                if vis.get(nid) or nid not in self.G:
                    continue
                vis[nid] = True
                nd = dist(q, self.G[nid].item.emb)
                best_found_dist = -found[0][0]
                if len(found) < ef or nd < best_found_dist:
                    heapq.heappush(cands, (nd, nid))
                    heapq.heappush(found, (-nd, nid))
                    if len(found) > ef:
                        heapq.heappop(found)

        res = [(-neg_d, id_) for neg_d, id_ in found]
        res.sort()
        return res

    def _select_nbrs(self, cands: List[Tuple[float, int]], max_m: int) -> List[int]:
        return [cands[i][1] for i in range(min(len(cands), max_m))]

    def insert(self, item: VectorItem, dist: DistFn):
        id_ = item.id
        lvl = self._rand_level()
        node = HNSWNode(item=item, max_lyr=lvl, nbrs=[[] for _ in range(lvl + 1)])
        self.G[id_] = node

        if self.entry_pt == -1:
            self.entry_pt = id_
            self.top_layer = lvl
            return

        ep = self.entry_pt
        for lc in range(self.top_layer, lvl, -1):
            if lc < len(self.G[ep].nbrs):
                W = self._search_layer(item.emb, ep, 1, lc, dist)
                if W:
                    ep = W[0][1]

        for lc in range(min(self.top_layer, lvl), -1, -1):
            W = self._search_layer(item.emb, ep, self.ef_build, lc, dist)
            max_m = self.M0 if lc == 0 else self.M
            sel = self._select_nbrs(W, max_m)
            self.G[id_].nbrs[lc] = sel

            for nid in sel:
                if nid not in self.G:
                    continue
                if len(self.G[nid].nbrs) <= lc:
                    self.G[nid].nbrs.extend([] for _ in range(lc + 1 - len(self.G[nid].nbrs)))
                conn = self.G[nid].nbrs[lc]
                conn.append(id_)
                if len(conn) > max_m:
                    ds = []
                    for c in conn:
                        if c in self.G:
                            ds.append((dist(self.G[nid].item.emb, self.G[c].item.emb), c))
                    ds.sort()
                    self.G[nid].nbrs[lc] = [ds[i][1] for i in range(min(max_m, len(ds)))]

            if W:
                ep = W[0][1]

        if lvl > self.top_layer:
            self.top_layer = lvl
            self.entry_pt = id_

    def knn(self, q: List[float], k: int, ef: int, dist: DistFn) -> List[Tuple[float, int]]:
        if self.entry_pt == -1:
            return []
        ep = self.entry_pt
        for lc in range(self.top_layer, 0, -1):
            if lc < len(self.G[ep].nbrs):
                W = self._search_layer(q, ep, 1, lc, dist)
                if W:
                    ep = W[0][1]
        W = self._search_layer(q, ep, max(ef, k), 0, dist)
        return W[:k]

    def remove(self, id_: int):
        if id_ not in self.G:
            return
        for nid, nd in self.G.items():
            for layer in nd.nbrs:
                if id_ in layer:
                    layer.remove(id_)
        if self.entry_pt == id_:
            self.entry_pt = -1
            for nid in self.G:
                if nid != id_:
                    self.entry_pt = nid
                    break
        del self.G[id_]

    def get_info(self) -> dict:
        top_layer = self.top_layer
        node_count = len(self.G)
        max_l = max(top_layer + 1, 1)
        nodes_per_layer = [0] * max_l
        edges_per_layer = [0] * max_l
        nodes = []
        edges = []

        for id_, nd in self.G.items():
            nodes.append({
                "id": id_,
                "metadata": nd.item.metadata,
                "category": nd.item.category,
                "maxLyr": nd.max_lyr
            })
            for lc in range(min(nd.max_lyr + 1, max_l)):
                nodes_per_layer[lc] += 1
                if lc < len(nd.nbrs):
                    for nid in nd.nbrs[lc]:
                        if id_ < nid:
                            edges_per_layer[lc] += 1
                            edges.append({"src": id_, "dst": nid, "lyr": lc})

        return {
            "topLayer": top_layer,
            "nodeCount": node_count,
            "nodesPerLayer": nodes_per_layer,
            "edgesPerLayer": edges_per_layer,
            "nodes": nodes,
            "edges": edges
        }

    def size(self) -> int:
        return len(self.G)


# =====================================================================
#  VECTOR DATABASE  (demo 16D index)
# =====================================================================

@dataclass
class Hit:
    id: int
    meta: str
    cat: str
    emb: List[float]
    dist: float


@dataclass
class SearchOut:
    hits: List[Hit]
    us: int
    algo: str
    metric: str


@dataclass
class BenchOut:
    bf_us: int
    kd_us: int
    hnsw_us: int
    n: int


class VectorDB:
    def __init__(self, d: int):
        self.dims = d
        self.store: Dict[int, VectorItem] = {}
        self.bf = BruteForce()
        self.kdt = KDTree(d)
        self.hnsw = HNSW(16, 200)
        self.mu = threading.Lock()
        self.next_id = 1

    def insert(self, meta: str, cat: str, emb: List[float], dist: DistFn) -> int:
        with self.mu:
            v = VectorItem(id=self.next_id, metadata=meta, category=cat, emb=emb)
            self.next_id += 1
            self.store[v.id] = v
            self.bf.insert(v)
            self.kdt.insert(v)
            self.hnsw.insert(v, dist)
            return v.id

    def remove(self, id_: int) -> bool:
        with self.mu:
            if id_ not in self.store:
                return False
            del self.store[id_]
            self.bf.remove(id_)
            self.hnsw.remove(id_)
            rem = list(self.store.values())
            self.kdt.rebuild(rem)
            return True

    def search(self, q: List[float], k: int, metric: str, algo: str) -> SearchOut:
        with self.mu:
            dfn = get_dist_fn(metric)
            t0 = time.monotonic_ns()
            if algo == "bruteforce":
                raw = self.bf.knn(q, k, dfn)
            elif algo == "kdtree":
                raw = self.kdt.knn(q, k, dfn)
            else:
                raw = self.hnsw.knn(q, k, 50, dfn)
            us = (time.monotonic_ns() - t0) // 1000

            hits = []
            for d, id_ in raw:
                if id_ in self.store:
                    v = self.store[id_]
                    hits.append(Hit(id=id_, meta=v.metadata, cat=v.category, emb=v.emb, dist=d))
            return SearchOut(hits=hits, us=us, algo=algo, metric=metric)

    def benchmark(self, q: List[float], k: int, metric: str) -> BenchOut:
        with self.mu:
            dfn = get_dist_fn(metric)

            def time_fn(fn) -> int:
                t = time.monotonic_ns()
                fn()
                return (time.monotonic_ns() - t) // 1000

            bf_us = time_fn(lambda: self.bf.knn(q, k, dfn))
            kd_us = time_fn(lambda: self.kdt.knn(q, k, dfn))
            hnsw_us = time_fn(lambda: self.hnsw.knn(q, k, 50, dfn))
            return BenchOut(bf_us=bf_us, kd_us=kd_us, hnsw_us=hnsw_us, n=len(self.store))

    def all(self) -> List[VectorItem]:
        with self.mu:
            return list(self.store.values())

    def hnsw_info(self) -> dict:
        with self.mu:
            return self.hnsw.get_info()

    def size(self) -> int:
        with self.mu:
            return len(self.store)


# =====================================================================
#  DOCUMENT DATABASE  — HNSW over real Ollama embeddings
# =====================================================================

@dataclass
class DocItem:
    id: int
    title: str
    text: str
    emb: List[float]


class DocumentDB:
    def __init__(self):
        self.store: Dict[int, DocItem] = {}
        self.hnsw = HNSW(16, 200)
        self.bf = BruteForce()
        self.mu = threading.Lock()
        self.next_id = 1
        self.dims = 0

    def insert(self, title: str, text: str, emb: List[float]) -> int:
        with self.mu:
            if self.dims == 0:
                self.dims = len(emb)
            item = DocItem(id=self.next_id, title=title, text=text, emb=emb)
            self.next_id += 1
            self.store[item.id] = item
            vi = VectorItem(id=item.id, metadata=title, category="doc", emb=emb)
            self.hnsw.insert(vi, cosine)
            self.bf.insert(vi)
            return item.id

    def search(self, q: List[float], k: int, max_dist: float = 0.7) -> List[Tuple[float, 'DocItem']]:
        with self.mu:
            if not self.store:
                return []
            if len(self.store) < 10:
                raw = self.bf.knn(q, k, cosine)
            else:
                raw = self.hnsw.knn(q, k, 50, cosine)
            out = []
            for d, id_ in raw:
                if id_ in self.store and d <= max_dist:
                    out.append((d, self.store[id_]))
            return out

    def remove(self, id_: int) -> bool:
        with self.mu:
            if id_ not in self.store:
                return False
            del self.store[id_]
            self.hnsw.remove(id_)
            self.bf.remove(id_)
            return True

    def all(self) -> List['DocItem']:
        with self.mu:
            return list(self.store.values())

    def size(self) -> int:
        with self.mu:
            return len(self.store)

    def get_dims(self) -> int:
        return self.dims


# =====================================================================
#  OLLAMA CLIENT  — wraps local Ollama REST API
#  Install:  https://ollama.com
#  Models:   ollama pull nomic-embed-text
#            ollama pull llama3.2
# =====================================================================

class OllamaClient:
    def __init__(self, host: str = "127.0.0.1", port: int = 11434):
        self.host = host
        self.port = port
        self.embed_model = "nomic-embed-text"
        self.gen_model = "mistral"

    def _base_url(self) -> str:
        return f"http://{self.host}:{self.port}"

    def is_available(self) -> bool:
        try:
            r = requests.get(f"{self._base_url()}/api/tags", timeout=2)
            return r.status_code == 200
        except Exception:
            return False

    def embed(self, text: str) -> List[float]:
        try:
            r = requests.post(
                f"{self._base_url()}/api/embeddings",
                json={"model": self.embed_model, "prompt": text},
                timeout=30
            )
            if r.status_code != 200:
                return []
            data = r.json()
            return data.get("embedding", [])
        except Exception:
            return []

    def generate(self, prompt: str) -> str:
        try:
            r = requests.post(
                f"{self._base_url()}/api/generate",
                json={"model": self.gen_model, "prompt": prompt, "stream": False},
                timeout=180
            )
            if r.status_code != 200:
                return "ERROR: Ollama unavailable. Run: ollama serve"
            return r.json().get("response", "")
        except Exception:
            return "ERROR: Ollama unavailable. Run: ollama serve"

    def generate_stream(self, prompt: str) -> Generator[str, None, None]:
        """Yields one token at a time from Ollama streaming API."""
        try:
            with requests.post(
                f"{self._base_url()}/api/generate",
                json={"model": self.gen_model, "prompt": prompt, "stream": True},
                timeout=180,
                stream=True
            ) as r:
                if r.status_code != 200:
                    yield "__ERROR__: Ollama returned non-200. Run: ollama serve"
                    return
                for line in r.iter_lines():
                    if not line:
                        continue
                    try:
                        data = json.loads(line)
                    except json.JSONDecodeError:
                        continue
                    token = data.get("response", "")
                    if token:
                        yield token
                    # Ollama sets done=true on the final chunk
                    if data.get("done", False):
                        return
        except requests.exceptions.ConnectionError:
            yield "__ERROR__: Ollama went offline mid-stream. Run: ollama serve"
        except Exception as e:
            yield f"__ERROR__: {str(e)}"


# =====================================================================
#  TEXT CHUNKER
# =====================================================================

def chunk_text(text: str, chunk_words: int = 250, overlap_words: int = 30) -> List[str]:
    words = text.split()
    if not words:
        return []
    if len(words) <= chunk_words:
        return [text]
    chunks = []
    step = chunk_words - overlap_words
    i = 0
    while i < len(words):
        end = min(i + chunk_words, len(words))
        chunk = " ".join(words[i:end])
        chunks.append(chunk)
        if end == len(words):
            break
        i += step
    return chunks


# =====================================================================
#  DEMO DATA  (16D categorical vectors)
# =====================================================================

def load_demo(db: VectorDB):
    dist = get_dist_fn("cosine")
    # Dims 0-3: CS | Dims 4-7: Math | Dims 8-11: Food | Dims 12-15: Sports
    db.insert("Linked List: nodes connected by pointers", "cs",
        [0.90,0.85,0.72,0.68,0.12,0.08,0.15,0.10,0.05,0.08,0.06,0.09,0.07,0.11,0.08,0.06], dist)
    db.insert("Binary Search Tree: O(log n) search and insert", "cs",
        [0.88,0.82,0.78,0.74,0.15,0.10,0.08,0.12,0.06,0.07,0.08,0.05,0.09,0.06,0.07,0.10], dist)
    db.insert("Dynamic Programming: memoization overlapping subproblems", "cs",
        [0.82,0.76,0.88,0.80,0.20,0.18,0.12,0.09,0.07,0.06,0.08,0.07,0.08,0.09,0.06,0.07], dist)
    db.insert("Graph BFS and DFS: breadth and depth first traversal", "cs",
        [0.85,0.80,0.75,0.82,0.18,0.14,0.10,0.08,0.06,0.09,0.07,0.06,0.10,0.08,0.09,0.07], dist)
    db.insert("Hash Table: O(1) lookup with collision chaining", "cs",
        [0.87,0.78,0.70,0.76,0.13,0.11,0.09,0.14,0.08,0.07,0.06,0.08,0.07,0.10,0.08,0.09], dist)
    db.insert("Calculus: derivatives integrals and limits", "math",
        [0.12,0.15,0.18,0.10,0.91,0.86,0.78,0.72,0.08,0.06,0.07,0.09,0.07,0.08,0.06,0.10], dist)
    db.insert("Linear Algebra: matrices eigenvalues eigenvectors", "math",
        [0.20,0.18,0.15,0.12,0.88,0.90,0.82,0.76,0.09,0.07,0.08,0.06,0.10,0.07,0.08,0.09], dist)
    db.insert("Probability: distributions random variables Bayes theorem", "math",
        [0.15,0.12,0.20,0.18,0.84,0.80,0.88,0.82,0.07,0.08,0.06,0.10,0.09,0.06,0.09,0.08], dist)
    db.insert("Number Theory: primes modular arithmetic RSA cryptography", "math",
        [0.22,0.16,0.14,0.20,0.80,0.85,0.76,0.90,0.08,0.09,0.07,0.06,0.08,0.10,0.07,0.06], dist)
    db.insert("Combinatorics: permutations combinations generating functions", "math",
        [0.18,0.20,0.16,0.14,0.86,0.78,0.84,0.80,0.06,0.07,0.09,0.08,0.06,0.09,0.10,0.07], dist)
    db.insert("Neapolitan Pizza: wood-fired dough San Marzano tomatoes", "food",
        [0.08,0.06,0.09,0.07,0.07,0.08,0.06,0.09,0.90,0.86,0.78,0.72,0.08,0.06,0.09,0.07], dist)
    db.insert("Sushi: vinegared rice raw fish and nori rolls", "food",
        [0.06,0.08,0.07,0.09,0.09,0.06,0.08,0.07,0.86,0.90,0.82,0.76,0.07,0.09,0.06,0.08], dist)
    db.insert("Ramen: noodle soup with chashu pork and soft-boiled eggs", "food",
        [0.09,0.07,0.06,0.08,0.08,0.09,0.07,0.06,0.82,0.78,0.90,0.84,0.09,0.07,0.08,0.06], dist)
    db.insert("Tacos: corn tortillas with carnitas salsa and cilantro", "food",
        [0.07,0.09,0.08,0.06,0.06,0.07,0.09,0.08,0.78,0.82,0.86,0.90,0.06,0.08,0.07,0.09], dist)
    db.insert("Croissant: laminated pastry with buttery flaky layers", "food",
        [0.06,0.07,0.10,0.09,0.10,0.06,0.07,0.10,0.85,0.80,0.76,0.82,0.09,0.07,0.10,0.06], dist)
    db.insert("Basketball: fast-paced shooting dribbling slam dunks", "sports",
        [0.09,0.07,0.08,0.10,0.08,0.09,0.07,0.06,0.08,0.07,0.09,0.06,0.91,0.85,0.78,0.72], dist)
    db.insert("Football: tackles touchdowns field goals and strategy", "sports",
        [0.07,0.09,0.06,0.08,0.09,0.07,0.10,0.08,0.07,0.09,0.08,0.07,0.87,0.89,0.82,0.76], dist)
    db.insert("Tennis: racket volleys groundstrokes and Wimbledon serves", "sports",
        [0.08,0.06,0.09,0.07,0.07,0.08,0.06,0.09,0.09,0.06,0.07,0.08,0.83,0.80,0.88,0.82], dist)
    db.insert("Chess: openings endgames tactics strategic board game", "sports",
        [0.25,0.20,0.22,0.18,0.22,0.18,0.20,0.15,0.06,0.08,0.07,0.09,0.80,0.84,0.78,0.90], dist)
    db.insert("Swimming: butterfly freestyle backstroke Olympic competition", "sports",
        [0.06,0.08,0.07,0.09,0.08,0.06,0.09,0.07,0.10,0.08,0.06,0.07,0.85,0.82,0.86,0.80], dist)

    # ── CS: 8 more ────────────────────────────────────────────────────
    db.insert("Sorting Algorithms: quicksort mergesort heapsort comparison", "cs",
        [0.86,0.80,0.74,0.70,0.14,0.09,0.11,0.13,0.07,0.06,0.08,0.07,0.08,0.07,0.09,0.06], dist)
    db.insert("Stack and Queue: LIFO FIFO push pop enqueue dequeue", "cs",
        [0.84,0.78,0.68,0.72,0.11,0.13,0.10,0.08,0.06,0.08,0.07,0.05,0.09,0.08,0.06,0.07], dist)
    db.insert("Recursion: base case call stack factorial Fibonacci", "cs",
        [0.80,0.74,0.84,0.78,0.17,0.15,0.11,0.10,0.08,0.07,0.06,0.08,0.07,0.10,0.08,0.06], dist)
    db.insert("Trie: prefix tree autocomplete dictionary word search", "cs",
        [0.83,0.77,0.71,0.75,0.12,0.10,0.08,0.13,0.07,0.06,0.09,0.07,0.08,0.06,0.10,0.08], dist)
    db.insert("Heap: priority queue min-heap max-heap Dijkstra", "cs",
        [0.85,0.79,0.73,0.77,0.16,0.12,0.09,0.11,0.06,0.09,0.07,0.08,0.10,0.07,0.08,0.09], dist)
    db.insert("Big O Notation: time complexity space complexity analysis", "cs",
        [0.81,0.75,0.86,0.76,0.19,0.14,0.13,0.08,0.07,0.06,0.08,0.06,0.09,0.08,0.07,0.06], dist)
    db.insert("Neural Network: layers weights backpropagation gradient descent", "cs",
        [0.79,0.83,0.76,0.80,0.21,0.17,0.14,0.11,0.08,0.07,0.06,0.09,0.07,0.09,0.08,0.07], dist)
    db.insert("Operating System: process thread scheduling memory management", "cs",
        [0.82,0.76,0.69,0.73,0.13,0.10,0.12,0.09,0.06,0.08,0.07,0.06,0.10,0.07,0.09,0.08], dist)

    # ── Math: 8 more ──────────────────────────────────────────────────
    db.insert("Differential Equations: ODEs PDEs initial value problems", "math",
        [0.14,0.17,0.19,0.11,0.89,0.83,0.75,0.70,0.07,0.06,0.08,0.09,0.08,0.07,0.06,0.09], dist)
    db.insert("Fourier Transform: frequency domain signal decomposition", "math",
        [0.19,0.14,0.16,0.13,0.85,0.88,0.80,0.74,0.08,0.07,0.06,0.08,0.09,0.06,0.08,0.07], dist)
    db.insert("Set Theory: unions intersections subsets cardinality", "math",
        [0.13,0.11,0.18,0.16,0.83,0.79,0.86,0.78,0.06,0.08,0.07,0.09,0.07,0.08,0.10,0.06], dist)
    db.insert("Topology: open sets continuity homeomorphism manifolds", "math",
        [0.17,0.13,0.15,0.19,0.81,0.86,0.74,0.88,0.09,0.07,0.08,0.06,0.06,0.09,0.07,0.08], dist)
    db.insert("Graph Theory: vertices edges paths cycles connectivity", "math",
        [0.21,0.19,0.17,0.15,0.87,0.76,0.82,0.79,0.07,0.08,0.06,0.07,0.08,0.07,0.09,0.10], dist)
    db.insert("Statistics: mean variance standard deviation hypothesis testing", "math",
        [0.16,0.13,0.21,0.17,0.82,0.77,0.85,0.81,0.08,0.06,0.09,0.07,0.07,0.10,0.06,0.08], dist)
    db.insert("Abstract Algebra: groups rings fields homomorphisms", "math",
        [0.12,0.18,0.14,0.20,0.78,0.83,0.77,0.86,0.06,0.09,0.08,0.07,0.09,0.06,0.08,0.07], dist)
    db.insert("Complex Analysis: analytic functions residues contour integration", "math",
        [0.20,0.15,0.12,0.17,0.84,0.81,0.73,0.89,0.07,0.08,0.06,0.08,0.08,0.07,0.09,0.06], dist)

    # ── Food: 7 more ──────────────────────────────────────────────────
    db.insert("Biryani: aromatic basmati rice slow-cooked spiced meat", "food",
        [0.07,0.09,0.08,0.06,0.08,0.07,0.09,0.06,0.88,0.83,0.77,0.80,0.07,0.08,0.06,0.09], dist)
    db.insert("Pasta Carbonara: eggs pecorino guanciale black pepper", "food",
        [0.08,0.06,0.07,0.09,0.07,0.09,0.06,0.08,0.84,0.88,0.80,0.74,0.06,0.09,0.08,0.07], dist)
    db.insert("Dim Sum: steamed dumplings har gow siu mai bamboo basket", "food",
        [0.06,0.08,0.09,0.07,0.09,0.06,0.08,0.07,0.80,0.76,0.88,0.86,0.08,0.07,0.09,0.06], dist)
    db.insert("Paella: saffron rice seafood chorizo Valencia Spain", "food",
        [0.09,0.07,0.06,0.08,0.06,0.08,0.07,0.09,0.76,0.80,0.84,0.88,0.07,0.06,0.08,0.09], dist)
    db.insert("Butter Chicken: tomato cream sauce tandoor marinated chicken", "food",
        [0.07,0.06,0.10,0.08,0.08,0.07,0.06,0.10,0.83,0.79,0.75,0.85,0.08,0.09,0.07,0.06], dist)
    db.insert("Tiramisu: espresso soaked ladyfingers mascarpone cocoa", "food",
        [0.06,0.09,0.07,0.08,0.09,0.06,0.10,0.07,0.81,0.85,0.79,0.73,0.09,0.07,0.06,0.08], dist)
    db.insert("Pho: Vietnamese beef broth rice noodles herbs anise", "food",
        [0.08,0.07,0.06,0.09,0.07,0.09,0.06,0.08,0.87,0.81,0.73,0.77,0.06,0.08,0.09,0.07], dist)

    # ── Sports: 7 more ────────────────────────────────────────────────
    db.insert("Cricket: batting bowling fielding wickets test match IPL", "sports",
        [0.08,0.07,0.09,0.06,0.09,0.08,0.07,0.06,0.07,0.09,0.08,0.06,0.88,0.83,0.77,0.74], dist)
    db.insert("Badminton: shuttlecock racket smash drop shot net play", "sports",
        [0.06,0.09,0.07,0.08,0.07,0.06,0.09,0.08,0.08,0.07,0.06,0.09,0.84,0.88,0.80,0.76], dist)
    db.insert("Cycling: Tour de France peloton sprint cadence endurance", "sports",
        [0.09,0.06,0.08,0.07,0.08,0.09,0.06,0.07,0.06,0.08,0.09,0.07,0.80,0.78,0.86,0.84], dist)
    db.insert("Volleyball: spike serve block rotation libero beach court", "sports",
        [0.07,0.08,0.06,0.09,0.06,0.07,0.08,0.09,0.09,0.06,0.07,0.08,0.82,0.86,0.84,0.78], dist)
    db.insert("Boxing: jab cross hook uppercut footwork ring knockouts", "sports",
        [0.08,0.06,0.09,0.07,0.07,0.08,0.09,0.06,0.07,0.08,0.06,0.09,0.86,0.80,0.82,0.88], dist)
    db.insert("Gymnastics: floor beam vault uneven bars artistic routines", "sports",
        [0.06,0.07,0.08,0.09,0.09,0.07,0.06,0.08,0.08,0.07,0.09,0.06,0.78,0.84,0.88,0.82], dist)
    db.insert("Golf: fairway iron wedge putt handicap birdie eagle swing", "sports",
        [0.09,0.08,0.07,0.06,0.06,0.09,0.07,0.08,0.06,0.09,0.08,0.07,0.83,0.79,0.85,0.87], dist)


# =====================================================================
#  HELPERS
# =====================================================================

def parse_vec(s: str) -> List[float]:
    v = []
    for t in s.split(","):
        try:
            v.append(float(t))
        except Exception:
            pass
    return v


# =====================================================================
#  EMBED AND STORE HELPER
# =====================================================================

async def embed_and_store(title: str, text: str) -> dict:
    """
    Chunk text, embed each chunk via Ollama, and store in doc_db.
    Returns a result dict on success, or an error dict on failure.
    Extracted to avoid duplicating this logic across /doc/insert,
    /doc/upload-pdf, and /doc/upload-txt.
    """
    chunks = chunk_text(text, 250, 30)
    ids = []

    for i, chunk in enumerate(chunks):
        emb = ollama.embed(chunk)
        if not emb:
            return {
                "error": (
                    "Ollama unavailable. "
                    "Install from https://ollama.com then run: "
                    "ollama pull nomic-embed-text && ollama pull llama3.2"
                )
            }
        chunk_title = (
            f"{title} [{i+1}/{len(chunks)}]" if len(chunks) > 1 else title
        )
        ids.append(doc_db.insert(chunk_title, chunk, emb))

    return {
        "ids": ids,
        "chunks": len(chunks),
        "dims": doc_db.get_dims()
    }


# =====================================================================
#  PERSISTENCE  — save/load DocumentDB to JSON
# =====================================================================

def save_doc_db():
    """
    Write every DocItem in doc_db to SAVE_FILE as a JSON array.
    Called after every insert or delete so the file is always current.
    """
    data = []
    for item in doc_db.all():
        data.append({
            "id":    item.id,
            "title": item.title,
            "text":  item.text,
            "emb":   item.emb
        })
    try:
        with open(SAVE_FILE, "w", encoding="utf-8") as f:
            json.dump(data, f)
    except Exception as e:
        print(f"[persistence] WARNING: could not save doc_db — {e}")


def load_doc_db():
    """
    Re-hydrate doc_db from SAVE_FILE on startup.
    Silently skips if the file does not exist yet (first run).
    """
    if not os.path.exists(SAVE_FILE):
        return
    try:
        with open(SAVE_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
        for item in data:
            doc_db.insert(item["title"], item["text"], item["emb"])
        # Restore next_id so new inserts never reuse an existing ID
        if data:
            doc_db.next_id = max(item["id"] for item in data) + 1
        print(f"[persistence] Loaded {len(data)} document chunk(s) from {SAVE_FILE}")
    except Exception as e:
        print(f"[persistence] WARNING: could not load doc_db — {e}")


# =====================================================================
#  FASTAPI APPLICATION
# =====================================================================

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["GET", "POST", "DELETE", "OPTIONS"],
    allow_headers=["Content-Type"],
)

# Global state
db = VectorDB(DIMS)
doc_db = DocumentDB()
ollama = OllamaClient()

load_demo(db)
load_doc_db()

ollama_up = ollama.is_available()
print("=== VectorDB Engine ===")
print("http://localhost:8080")
print(f"{db.size()} demo vectors | {DIMS} dims | HNSW+KD-Tree+BruteForce")
print(f"Ollama: {'ONLINE' if ollama_up else 'OFFLINE (install from ollama.com)'}")
if ollama_up:
    print(f"  embed model: {ollama.embed_model}  gen model: {ollama.gen_model}")


# ── DEMO VECTOR ENDPOINTS ─────────────────────────────────────────

@app.get("/search")
def search(v: str = "", k: int = 5, metric: str = "cosine", algo: str = "hnsw", category: str = ""):
    q = parse_vec(v)
    if len(q) != DIMS:
        return JSONResponse({"error": f"need {DIMS}D vector"})
    out = db.search(q, k, metric, algo)
    results = []
    for h in out.hits:
        # ── Category filter: skip if category given and doesn't match ──
        if category and h.cat != category:
            continue
        results.append({
            "id": h.id,
            "metadata": h.meta,
            "category": h.cat,
            "distance": round(h.dist, 6),
            "embedding": [round(x, 4) for x in h.emb]
        })
    return {
        "results": results,
        "latencyUs": out.us,
        "algo": out.algo,
        "metric": out.metric,
        "categoryFilter": category if category else None
    }


@app.post("/insert")
async def insert(request: Request):
    body = await request.json()
    meta = body.get("metadata", "")
    cat = body.get("category", "")
    emb = body.get("embedding", [])
    if not meta or not emb or len(emb) != DIMS:
        return JSONResponse({"error": "invalid body"})
    id_ = db.insert(meta, cat, emb, get_dist_fn("cosine"))
    return {"id": id_}


@app.delete("/delete/{id_}")
def delete(id_: int):
    ok = db.remove(id_)
    return {"ok": ok}


@app.get("/items")
def items():
    all_items = db.all()
    return [
        {
            "id": v.id,
            "metadata": v.metadata,
            "category": v.category,
            "embedding": [round(x, 4) for x in v.emb]
        }
        for v in all_items
    ]


@app.get("/benchmark")
def benchmark(v: str = "", k: int = 5, metric: str = "cosine"):
    q = parse_vec(v)
    if len(q) != DIMS:
        return JSONResponse({"error": f"need {DIMS}D vector"})
    b = db.benchmark(q, k, metric)
    return {
        "bruteforceUs": b.bf_us,
        "kdtreeUs": b.kd_us,
        "hnswUs": b.hnsw_us,
        "itemCount": b.n
    }


@app.get("/hnsw-info")
def hnsw_info():
    gi = db.hnsw_info()
    return gi


# ── DOCUMENT + RAG ENDPOINTS ──────────────────────────────────────

@app.post("/doc/insert")
async def doc_insert(request: Request):
    body = await request.json()
    title = body.get("title", "")
    text = body.get("text", "")
    if not title or not text:
        return JSONResponse({"error": "need title and text"})

    result = await embed_and_store(title, text)
    if "error" in result:
        return JSONResponse(result)
    save_doc_db()
    return result


@app.post("/doc/upload-pdf")
async def doc_upload_pdf(file: UploadFile = File(...)):
    # ── Validate file type ────────────────────────────────────────────
    if not file.filename.lower().endswith(".pdf"):
        return JSONResponse({"error": "only .pdf files are accepted"})

    # ── Read raw bytes ────────────────────────────────────────────────
    raw = await file.read()
    if not raw:
        return JSONResponse({"error": "uploaded file is empty"})

    # ── Extract text with PyMuPDF ─────────────────────────────────────
    try:
        pdf = fitz.open(stream=raw, filetype="pdf")
    except Exception:
        return JSONResponse({"error": "corrupt or unreadable PDF"})

    text = ""
    for page in pdf:
        text += page.get_text()
    pdf.close()

    text = text.strip()
    if not text:
        return JSONResponse({"error": "PDF has no extractable text (may be scanned image)"})

    # ── Use filename (without .pdf) as the title ──────────────────────
    title = file.filename.removesuffix(".pdf")

    # ── Chunk → embed → store via shared helper ─────────────────────
    result = await embed_and_store(title, text)
    if "error" in result:
        return JSONResponse(result)
    save_doc_db()
    return result


@app.post("/doc/upload-txt")
async def doc_upload_txt(file: UploadFile = File(...)):
    # ── Validate file type ────────────────────────────────────────────
    if not file.filename.lower().endswith(".txt"):
        return JSONResponse({"error": "only .txt files are accepted"})

    # ── Read and decode ───────────────────────────────────────────────
    raw = await file.read()
    if not raw:
        return JSONResponse({"error": "uploaded file is empty"})

    try:
        text = raw.decode("utf-8").strip()
    except UnicodeDecodeError:
        return JSONResponse({"error": "could not decode file — make sure it is UTF-8 encoded"})

    if not text:
        return JSONResponse({"error": "text file has no content"})

    # ── Use filename (without .txt) as the title ──────────────────────
    title = file.filename.removesuffix(".txt")

    # ── Chunk → embed → store via shared helper ─────────────────────
    result = await embed_and_store(title, text)
    if "error" in result:
        return JSONResponse(result)
    save_doc_db()
    return result


@app.delete("/doc/delete/{id_}")
def doc_delete(id_: int):
    ok = doc_db.remove(id_)
    if ok:
        save_doc_db()
    return {"ok": ok}


@app.get("/doc/list")
def doc_list():
    docs = doc_db.all()
    result = []
    for d in docs:
        preview = d.text[:120] + ("…" if len(d.text) > 120 else "")
        result.append({
            "id": d.id,
            "title": d.title,
            "preview": preview,
            "words": len(d.text.split())
        })
    return result


@app.post("/doc/search")
async def doc_search(request: Request):
    body = await request.json()
    question = body.get("question", "")
    k = body.get("k", 3)
    max_distance = body.get("max_distance", 0.7)  # ── similarity threshold
    if not question:
        return JSONResponse({"error": "need question"})

    q_emb = ollama.embed(question)
    if not q_emb:
        return JSONResponse({"error": "Ollama unavailable"})

    hits = doc_db.search(q_emb, k, max_distance)
    contexts = [
        {
            "id": item.id,
            "title": item.title,
            "distance": round(d, 4)
        }
        for d, item in hits
    ]
    return {"contexts": contexts}


@app.post("/doc/ask")
async def doc_ask(request: Request):

    body = await request.json()
    question = body.get("question", "")
    k = body.get("k", 3)

    t = time.time()
    q_emb = ollama.embed(question)
    print("EMBED:", round(time.time() - t, 2), "sec")
    
    max_distance = body.get("max_distance", 0.7)  # ── similarity threshold
    if not question:
        return JSONResponse({"error": "need question"})

    # Step 1: embed the question
    q_emb = ollama.embed(question)
    if not q_emb:
        return JSONResponse({"error": "Ollama unavailable"})

    # Step 2: retrieve top-k relevant chunks
    # hits = doc_db.search(q_emb, k, max_distance)
    t = time.time()
    hits = doc_db.search(q_emb, k, max_distance)
    print("SEARCH:", round(time.time() - t, 2), "sec")

    # Step 3: build prompt
    ctx = ""
    t = time.time()
    for i, (d, item) in enumerate(hits):
        ctx += f"[{i+1}] {item.title}:\n{item.text}\n\n"
    print("BUILD PROMPT:", round(time.time() - t, 2), "sec")

    prompt = (
        "You are a helpful assistant. Answer the user's question directly. "
        "Use the provided context if it contains relevant information. "
        "If it doesn't, just use your own general knowledge. "
        "IMPORTANT: Do NOT mention the 'context', 'provided text', or say things like "
        "'the context doesn't mention'. Just answer the question naturally.\n\n"
        f"Context:\n{ctx}"
        f"Question: {question}\n\n"
        "Answer:"
    )

    # Step 4: build context payload to send as first SSE event
    contexts = [
        {
            "id": item.id,
            "title": item.title,
            "text": item.text,
            "distance": round(d, 4)
        }
        for d, item in hits
    ]

    def event_stream() -> Generator[str, None, None]:
        # ── Event 1: context chunks ───────────────────────────────────
        ctx_payload = json.dumps({
            "type": "context",
            "contexts": contexts,
            "docCount": doc_db.size(),
            "model": ollama.gen_model
        })
        yield f"data: {ctx_payload}\n\n"

        # ── Event 2..N: stream each token ─────────────────────────────
        for token in ollama.generate_stream(prompt):
            if token.startswith("__ERROR__"):
                err_payload = json.dumps({"type": "error", "message": token})
                yield f"data: {err_payload}\n\n"
                return
            token_payload = json.dumps({"type": "token", "value": token})
            yield f"data: {token_payload}\n\n"

        # ── Final event: done ─────────────────────────────────────────
        done_payload = json.dumps({
            "type": "done",
            "model": ollama.gen_model,
            "docCount": doc_db.size()
        })
        yield f"data: {done_payload}\n\n"

    return StreamingResponse(event_stream(), media_type="text/event-stream")


@app.post("/doc/ask-sync")
async def doc_ask_sync(request: Request):
    """Original non-streaming endpoint kept for backward compatibility."""
    body = await request.json()
    question = body.get("question", "")
    k = body.get("k", 3)
    if not question:
        return JSONResponse({"error": "need question"})

    # Step 1: embed the question
    q_emb = ollama.embed(question)
    if not q_emb:
        return JSONResponse({"error": "Ollama unavailable"})

    # Step 2: retrieve top-k relevant chunks
    hits = doc_db.search(q_emb, k)

    # Step 3: build prompt
    ctx = ""
    for i, (d, item) in enumerate(hits):
        ctx += f"[{i+1}] {item.title}:\n{item.text}\n\n"

    prompt = (
        "You are a helpful assistant. Answer the user's question directly. "
        "Use the provided context if it contains relevant information. "
        "If it doesn't, just use your own general knowledge. "
        "IMPORTANT: Do NOT mention the 'context', 'provided text', or say things like "
        "'the context doesn't mention'. Just answer the question naturally.\n\n"
        f"Context:\n{ctx}"
        f"Question: {question}\n\n"
        "Answer:"
    )

    # Step 4: generate answer
    answer = ollama.generate(prompt)

    # Step 5: return everything
    contexts = [
        {
            "id": item.id,
            "title": item.title,
            "text": item.text,
            "distance": round(d, 4)
        }
        for d, item in hits
    ]
    return {
        "answer": answer,
        "model": ollama.gen_model,
        "contexts": contexts,
        "docCount": doc_db.size()
    }



@app.post("/doc/evaluate")
async def doc_evaluate(request: Request):
    """
    RAG Evaluation endpoint — no external libraries needed.

    Body:
        {
            "pairs": [
                {"question": "...", "expected_answer": "..."},
                ...
            ],
            "k": 3,
            "max_distance": 0.7
        }

    Per-question metrics:
        context_hit        — 1 if retrieved chunks contain keywords from expected_answer
        keyword_hit_rate   — fraction of expected keywords found in retrieved chunks
        answer_length      — word count of generated answer
        retrieval_distance — average cosine distance of retrieved chunks
    """
    body         = await request.json()
    pairs        = body.get("pairs", [])
    k            = body.get("k", 3)
    max_distance = body.get("max_distance", 0.7)

    if not pairs:
        return JSONResponse({"error": "need at least one question/expected_answer pair"})
    if not isinstance(pairs, list):
        return JSONResponse({"error": "pairs must be a list"})

    if not ollama.is_available():
        return JSONResponse({
            "error": (
                "Ollama unavailable. "
                "Run: ollama serve && ollama pull nomic-embed-text && ollama pull llama3.2"
            )
        })

    results = []

    for pair in pairs:
        question        = pair.get("question", "").strip()
        expected_answer = pair.get("expected_answer", "").strip()

        if not question or not expected_answer:
            results.append({
                "question": question,
                "error": "missing question or expected_answer"
            })
            continue

        # Step 1 — embed the question
        q_emb = ollama.embed(question)
        if not q_emb:
            results.append({"question": question, "error": "Ollama embed failed"})
            continue

        # Step 2 — retrieve context chunks
        hits = doc_db.search(q_emb, k, max_distance)

        # Step 3 — build prompt and generate answer
        ctx = ""
        for i, (d, item) in enumerate(hits):
            ctx += "[" + str(i + 1) + "] " + item.title + ":\n" + item.text + "\n\n"

        prompt = (
            "You are a helpful assistant. Answer the user question directly. "
            "Use the provided context if it contains relevant information. "
            "If it does not, just use your own general knowledge. "
            "Do NOT mention the context or say things like the context does not mention. "
            "Just answer the question naturally.\n\n"
            "Context:\n" + ctx +
            "Question: " + question + "\n\nAnswer:"
        )

        generated_answer = ollama.generate(prompt)

        # Step 4 — score

        # Metric 1: context_hit and keyword_hit_rate
        expected_keywords = set()
        for w in expected_answer.split():
            cleaned = w.lower().strip(".,!?;:")
            if len(cleaned) > 4:
                expected_keywords.add(cleaned)

        retrieved_text = " ".join(item.text for _, item in hits).lower()

        if expected_keywords:
            keyword_hits     = sum(1 for kw in expected_keywords if kw in retrieved_text)
            context_hit      = 1 if keyword_hits > 0 else 0
            keyword_hit_rate = round(keyword_hits / len(expected_keywords), 3)
        else:
            context_hit      = 0
            keyword_hit_rate = 0.0

        # Metric 2: answer_length
        answer_length = len(generated_answer.split())

        # Metric 3: retrieval_distance
        if hits:
            retrieval_distance = round(sum(d for d, _ in hits) / len(hits), 4)
        else:
            retrieval_distance = 1.0

        results.append({
            "question":         question,
            "expected_answer":  expected_answer,
            "generated_answer": generated_answer,
            "chunks_retrieved": len(hits),
            "scores": {
                "context_hit":        context_hit,
                "keyword_hit_rate":   keyword_hit_rate,
                "answer_length":      answer_length,
                "retrieval_distance": retrieval_distance,
            },
            "retrieved_chunks": [
                {"title": item.title, "distance": round(d, 4)}
                for d, item in hits
            ]
        })

    # Aggregate averages
    scored = [r for r in results if "scores" in r]
    if scored:
        avg_context_hit        = round(sum(r["scores"]["context_hit"]        for r in scored) / len(scored), 3)
        avg_keyword_hit_rate   = round(sum(r["scores"]["keyword_hit_rate"]   for r in scored) / len(scored), 3)
        avg_answer_length      = round(sum(r["scores"]["answer_length"]      for r in scored) / len(scored), 1)
        avg_retrieval_distance = round(sum(r["scores"]["retrieval_distance"] for r in scored) / len(scored), 4)
    else:
        avg_context_hit = avg_keyword_hit_rate = avg_answer_length = avg_retrieval_distance = 0.0

    return {
        "summary": {
            "total_pairs":             len(pairs),
            "evaluated":               len(scored),
            "avg_context_hit":         avg_context_hit,
            "avg_keyword_hit_rate":    avg_keyword_hit_rate,
            "avg_answer_length_words": avg_answer_length,
            "avg_retrieval_distance":  avg_retrieval_distance,
            "doc_count":               doc_db.size(),
            "model":                   ollama.gen_model,
        },
        "results": results
    }


@app.get("/status")
def status():
    up = ollama.is_available()
    return {
        "ollamaAvailable": up,
        "embedModel": ollama.embed_model,
        "genModel": ollama.gen_model,
        "docCount": doc_db.size(),
        "docDims": doc_db.get_dims(),
        "demoDims": DIMS,
        "demoCount": db.size()
    }


@app.get("/stats")
def stats():
    return {
        "count": db.size(),
        "dims": DIMS,
        "algorithms": ["bruteforce", "kdtree", "hnsw"],
        "metrics": ["euclidean", "cosine", "manhattan"]
    }


@app.get("/")
def index():
    return FileResponse("index.html")


# =====================================================================
#  ENTRY POINT
# =====================================================================

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8080)