# @title Install dependencies
!pip install -q pyarrow numpy sentence-transformers datasets sqlite-vec tqdm psutil

# @title Configuration and core utilities
import gc
import hashlib
import json
import os
import pickle
import random
import re
import sqlite3
import sys
import time
from collections import Counter, defaultdict
from pathlib import Path
from typing import Optional

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
import psutil
import sqlite_vec
from sentence_transformers import SentenceTransformer
from tqdm import tqdm

random.seed(42)
np.random.seed(42)

EMBEDDING_MODEL = "all-MiniLM-L6-v2"  # 384-dim, fast, good enough
EMBEDDING_DIM = 384
INDEX_DIR = Path("./nanomem_index")
INDEX_DIR.mkdir(exist_ok=True)
CACHE_DIR = Path("./bench_cache")
CACHE_DIR.mkdir(exist_ok=True)

PROC = psutil.Process(os.getpid())


def mem_mb() -> float:
    """Current RSS in MB. Useful for the writeup section."""
    return PROC.memory_info().rss / 1_000_000


class MemoryBudget:
    """Track peak memory (RSS) across a block of code plus across the whole run.

    Use as a context manager around any code where memory matters.
    Required at Tier 93 (target 2000) and Tier 100 (target 1000).

        with MemoryBudget("Phase 1 build", limit_mb=2000) as mb:
            inv_idx.build(docs)
            mb.update()                # sample again after allocation-heavy steps
        # prints: [mem] Phase 1 build: peak 1340 MB / budget 2000 MB

    At the end of the run:
        print(f"Run peak: {MemoryBudget.run_peak():.0f} MB")

    The eval harness uses this automatically and reports peak_mb in
    every result dict, so your scoreboard table can quote it.

    Note: this reports `psutil` peak RSS, not the conceptual working set.
    At Tier 93/100, peak RSS will often read higher than 2 GB even when
    you only have one shard live at a time, because Python doesn't return
    freed pages to the OS. We assess live memory use by reading
    your notebook; this number is a soft signal. **Keep the `[mem]`
    prints visible in the submitted notebook.**
    """

    _run_peak: float = 0.0

    def __init__(self, label: str, limit_mb: float | None = None):
        self.label = label
        self.limit_mb = limit_mb
        self.start: float = 0.0
        self.peak: float = 0.0

    def __enter__(self):
        self.start = mem_mb()
        self.peak = self.start
        return self

    def update(self) -> float:
        cur = mem_mb()
        self.peak = max(self.peak, cur)
        MemoryBudget._run_peak = max(MemoryBudget._run_peak, self.peak)
        return cur

    def __exit__(self, *args):
        self.update()
        msg = f"[mem] {self.label}: peak {self.peak:.0f} MB"
        if self.limit_mb is not None:
            msg += f" / budget {self.limit_mb:.0f} MB"
            if self.peak > self.limit_mb:
                msg += " (RSS above target; inspect live working set)"
        print(msg)
        return False  # do not suppress exceptions

    @classmethod
    def run_peak(cls) -> float:
        """Highest RSS seen inside any MemoryBudget block so far."""
        return cls._run_peak

    @classmethod
    def reset(cls):
        cls._run_peak = 0.0


def memory_breakdown(components: dict, label: str = "memory profile") -> dict:
    """Print a per-component memory breakdown table.

    Use this to figure out *where* your memory is going before you decide
    what to optimize at Tier 93 or Tier 100. Pass a dict of
    {label: object}; values can be:

      - numpy.ndarray:    uses .nbytes (exact)
      - InvertedIndex:    walks postings dict + doc_lengths
      - LSHIndex:         walks hyperplane arrays + bucket dicts
      - int / float:      treated as MB directly (use for measurements
                          you took elsewhere, e.g., measured native allocations)
      - anything else:    sys.getsizeof(obj) / 1_000_000 (approximate)

    Returns the {label: mb} dict so you can also persist it.

    **Keep the printed table visible in your notebook for Tier 93/100.**
    These are approximate component sizes, not proof of a peak-memory bound.
    Include embeddings and document text separately from their indexes.
    Account for shared state, native allocations, and temporary buffers.

    Measure components on your own data. Choose partition sizes from those
    measurements, leaving room for shared state and temporary buffers.
    """
    import sys
    sizes = {}
    for k, v in components.items():
        if isinstance(v, np.ndarray):
            sizes[k] = v.nbytes / 1_000_000
        elif isinstance(v, (int, float)):
            sizes[k] = float(v)
        elif hasattr(v, 'postings') and hasattr(v, 'doc_lengths'):
            n = sys.getsizeof(v.postings)
            for term, plist in v.postings.items():
                n += sys.getsizeof(term) + sys.getsizeof(plist) + len(plist) * 56
            n += sys.getsizeof(v.doc_lengths)
            sizes[k] = n / 1_000_000
        elif hasattr(v, 'bands'):
            n = 0
            for hp, buckets in v.bands:
                n += hp.nbytes if isinstance(hp, np.ndarray) else 0
                n += sys.getsizeof(buckets)
                for _bid, doc_ids in buckets.items():
                    n += sys.getsizeof(doc_ids) + len(doc_ids) * 28
            sizes[k] = n / 1_000_000
        else:
            sizes[k] = sys.getsizeof(v) / 1_000_000

    total = sum(sizes.values())
    print(f"\n=== {label} ===")
    print(f"{'COMPONENT':<28} {'MB':>10} {'%':>6}")
    print("-" * 46)
    for k in sorted(sizes, key=lambda x: -sizes[x]):
        mb = sizes[k]
        pct = (100 * mb / total) if total else 0
        print(f"{k:<28} {mb:>10.1f} {pct:>5.1f}%")
    print("-" * 46)
    print(f"{'TOTAL':<28} {total:>10.1f}  100.0%")
    print(f"process peak so far:        {MemoryBudget.run_peak():>7.1f} MB")
    return sizes

# @title Embedding helpers
# CPU works locally and in free Colab; change to "cuda" if your Colab has a GPU.
EMBED_DEVICE = "cpu"
import torch
torch.set_num_threads(min(4, os.cpu_count() or 1))
_embed_model = None


def get_embedder() -> SentenceTransformer:
    global _embed_model
    if _embed_model is None:
        print("Loading embedding model (one-time)...")
        _embed_model = SentenceTransformer(EMBEDDING_MODEL, device=EMBED_DEVICE)
    return _embed_model


def embed_text(text: str) -> np.ndarray:
    """Return a (EMBEDDING_DIM,) float32 unit vector."""
    v = get_embedder().encode([text], convert_to_numpy=True)[0].astype(np.float32)
    return v / (np.linalg.norm(v) + 1e-12)


def embed_batch(texts: list[str], batch_size: int = 32) -> np.ndarray:
    """Return an (N, EMBEDDING_DIM) float32 matrix of unit vectors."""
    M = get_embedder().encode(
        texts, batch_size=batch_size, show_progress_bar=True, convert_to_numpy=True
    ).astype(np.float32)
    norms = np.linalg.norm(M, axis=1, keepdims=True) + 1e-12
    return M / norms


def cached_embed_batch(doc_ids: list[str], texts: list[str], cache_path: Path) -> np.ndarray:
    """Compute embeddings once, save to Parquet, reload on subsequent calls.

    Use this to reuse embeddings across runs. Time a small batch on your
    machine before estimating the full run; document length affects runtime.
    The cache key is `cache_path`; pick a name that includes the dataset and
    tier (e.g., `bench_cache/codesearchnet_250k_embs.parquet`).

    On Colab specifically: Colab's local disk does NOT persist across session
    restarts. Mount Google Drive and point `cache_path` there:

        from google.colab import drive
        drive.mount('/content/drive')
        CACHE_DIR = Path("/content/drive/MyDrive/nanomem_cache")
        CACHE_DIR.mkdir(exist_ok=True)
        embs = cached_embed_batch(doc_ids, texts, CACHE_DIR / "codesearchnet_250k.parquet")

    For Tier 100, cache bounded batches separately so you can resume without
    holding the full embedding matrix in memory.

    The sidecar fingerprint includes document IDs, text, model configuration,
    and normalization. Changing any of these invalidates the cached embeddings.
    """
    if len(doc_ids) != len(texts):
        raise ValueError("Provide one text per document ID.")
    model = get_embedder()
    settings = {"model": EMBEDDING_MODEL, "dimension": EMBEDDING_DIM,
                "max_seq_length": model.max_seq_length,
                "prompts": getattr(model, "prompts", {}),
                "default_prompt": getattr(model, "default_prompt_name", None),
                "normalization": "l2-float32-v1"}
    key = hashlib.sha256()
    key.update(json.dumps(settings, sort_keys=True, default=str).encode())
    for doc_id, text in zip(doc_ids, texts):
        key.update(json.dumps([doc_id, text], ensure_ascii=False).encode("utf-8"))
        key.update(b"\n")
    key = key.hexdigest()
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    key_path = cache_path.with_suffix(cache_path.suffix + ".key")
    if cache_path.exists() and key_path.exists() and key_path.read_text().strip() == key:
        _, cached_embs = load_embeddings_parquet(cache_path)
        print(f"  [embed cache HIT] {cache_path.name} ({cached_embs.shape[0]} docs)")
        return cached_embs
    reason = "no cache" if not cache_path.exists() else "input or model changed"
    print(f"  [embed cache MISS: {reason}] computing embeddings for {len(texts)} docs...")
    embs = embed_batch(texts)
    save_embeddings_parquet(doc_ids, embs, cache_path)
    key_path.write_text(key)
    print(f"  [embed cache WROTE] {cache_path.name}")
    return embs

# @title Tokenizer
_TOKEN_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")


def tokenize(text: str) -> list[str]:
    """Lowercase identifier-style tokenizer.

    Splits on non-identifier characters; keeps `verify_jwt` as one token.
    Modify if you want, but you do not have to.
    """
    return [t.lower() for t in _TOKEN_RE.findall(text)]

# @title Storage helpers
def save_documents_json(docs: list[dict], path: Path) -> None:
    with open(path, "w") as f:
        for d in docs:
            f.write(json.dumps(d) + "\n")


def load_documents_json(path: Path) -> list[dict]:
    docs = []
    with open(path) as f:
        for line in f:
            docs.append(json.loads(line))
    return docs


def save_embeddings_parquet(doc_ids: list[str], embeddings: np.ndarray, path: Path) -> None:
    """One Parquet file. doc_id (string) + embedding (fixed-size list of float32)."""
    embeddings = np.ascontiguousarray(embeddings, dtype=np.float32)
    if embeddings.shape != (len(doc_ids), EMBEDDING_DIM):
        raise ValueError("Provide one embedding of the expected dimension per document.")
    table = pa.table({
        "doc_id": pa.array(doc_ids, type=pa.string()),
        "embedding": pa.FixedSizeListArray.from_arrays(
            pa.array(embeddings.reshape(-1), type=pa.float32()),
            EMBEDDING_DIM,
        ),
    })
    pq.write_table(table, path, compression="snappy")


def load_embeddings_parquet(path: Path) -> tuple[list[str], np.ndarray]:
    table = pq.read_table(path)
    doc_ids = table.column("doc_id").to_pylist()
    column = table.column("embedding").combine_chunks()
    embs = column.flatten().to_numpy(zero_copy_only=False).reshape(-1, EMBEDDING_DIM)
    return doc_ids, embs


def brute_force_cosine(query_vec: np.ndarray, embeddings: np.ndarray, top_k: int = 10) -> list[int]:
    """Return the top_k indices into `embeddings` by cosine similarity."""
    scores = embeddings @ query_vec
    return list(np.argsort(-scores)[:top_k])

# @title LSH math helpers
def random_projection_matrix(n_hyperplanes: int, dim: int = EMBEDDING_DIM) -> np.ndarray:
    """Returns (n_hyperplanes, dim) standard-normal hyperplanes."""
    return np.random.randn(n_hyperplanes, dim).astype(np.float32)


def hash_vector(vec: np.ndarray, hyperplanes: np.ndarray) -> int:
    """Project vec onto each hyperplane; return bit signature as int."""
    bits = (hyperplanes @ vec > 0).astype(np.uint8)
    sig = 0
    for b in bits:
        sig = (sig << 1) | int(b)
    return sig


def hash_batch(M: np.ndarray, hyperplanes: np.ndarray) -> np.ndarray:
    """Vectorized hash of an (N, dim) matrix. Returns (N,) int64 sigs."""
    bits = (M @ hyperplanes.T > 0).astype(np.uint8)
    n_hp = bits.shape[1]
    weights = (1 << np.arange(n_hp - 1, -1, -1)).astype(np.int64)
    return (bits.astype(np.int64) * weights).sum(axis=1)

# @title SQLite reference helpers
def open_sqlite_db(path: Path) -> sqlite3.Connection:
    """Open a SQLite connection with the sqlite-vec extension loaded."""
    path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(path))
    try:
        conn.execute("CREATE VIRTUAL TABLE temp.fts_check USING fts5(text)")
        conn.execute("DROP TABLE temp.fts_check")
        conn.enable_load_extension(True)
        sqlite_vec.load(conn)
        conn.enable_load_extension(False)
    except (AttributeError, sqlite3.Error) as exc:
        conn.close()
        raise RuntimeError("SQLite needs FTS5 and loadable extensions. Use the supplied Colab setup or a Python build with both enabled.") from exc
    return conn


def initialize_sqlite_corpus(conn: sqlite3.Connection) -> None:
    """Create empty reference tables. Call once before inserting batches.

    This resets the reference corpus; do not call between batches.
    """
    with conn:
        conn.execute("DROP TABLE IF EXISTS docs_fts")
        conn.execute("DROP TABLE IF EXISTS docs_vec")
        conn.execute("DROP TABLE IF EXISTS doc_id_map")
        conn.execute("CREATE VIRTUAL TABLE docs_fts USING fts5(doc_id UNINDEXED, content)")
        conn.execute(f"CREATE VIRTUAL TABLE docs_vec USING vec0(doc_idx INTEGER PRIMARY KEY, embedding float[{EMBEDDING_DIM}])")
        conn.execute("CREATE TABLE doc_id_map(doc_idx INTEGER PRIMARY KEY, doc_id TEXT UNIQUE)")


def append_sqlite_batch(conn: sqlite3.Connection, docs: list[dict], embeddings: np.ndarray) -> None:
    """Append one bounded batch; retain earlier batches. IDs must be unique.

    Initialize once, then embed and append each batch. A failed batch rolls back.
    """
    embeddings = np.ascontiguousarray(embeddings, dtype=np.float32)
    if embeddings.shape != (len(docs), EMBEDDING_DIM) or not np.isfinite(embeddings).all():
        raise ValueError("Provide one finite embedding of the expected dimension per document.")
    if len({d["doc_id"] for d in docs}) != len(docs):
        raise ValueError("Document IDs must be unique.")
    start = conn.execute("SELECT COALESCE(MAX(doc_idx), -1) + 1 FROM doc_id_map").fetchone()[0]
    with conn:
        conn.executemany("INSERT INTO doc_id_map(doc_idx, doc_id) VALUES (?, ?)",
                         ((start+i, d["doc_id"]) for i,d in enumerate(docs)))
        conn.executemany("INSERT INTO docs_fts(doc_id, content) VALUES (?, ?)",
                         ((d["doc_id"], d["text"]) for d in docs))
        conn.executemany("INSERT INTO docs_vec(doc_idx, embedding) VALUES (?, ?)",
                         ((start+i, embeddings[i].tobytes()) for i in range(len(docs))))


def build_sqlite_corpus(conn: sqlite3.Connection, docs: list[dict], embeddings: np.ndarray) -> None:
    """Rebuild a small reference corpus. Larger runs use initialize + append."""
    initialize_sqlite_corpus(conn)
    append_sqlite_batch(conn, docs, embeddings)


def fts5_query(conn: sqlite3.Connection, query: str, top_k: int = 10) -> list[tuple[str, float]]:
    """BM25 query via FTS5. Returns [(doc_id, score), ...] sorted by relevance."""
    tokens = re.findall(r"[A-Za-z0-9_]+", query)
    if not tokens:
        return []
    fts_query = " OR ".join(f'"{t}"' for t in tokens)
    try:
        rows = conn.execute(
            "SELECT doc_id, -bm25(docs_fts) FROM docs_fts WHERE docs_fts MATCH ? ORDER BY rank LIMIT ?",
            (fts_query, top_k)).fetchall()
        return [(r[0], r[1]) for r in rows]
    except sqlite3.OperationalError as exc:
        raise RuntimeError("FTS5 query failed; inspect the index and query instead of treating this as zero matches.") from exc


def vec_query(conn: sqlite3.Connection, query_vec: np.ndarray, top_k: int = 10) -> list[tuple[str, float]]:
    """Exact vector KNN via sqlite-vec. Returns [(doc_id, score), ...].

    Score is negative L2 distance: higher is better, but it is not cosine similarity.
    """
    rows = conn.execute(
        """SELECT m.doc_id, -v.distance FROM docs_vec v
           JOIN doc_id_map m ON v.doc_idx = m.doc_idx
           WHERE v.embedding MATCH ? AND k = ? ORDER BY distance LIMIT ?""",
        (query_vec.tobytes(), top_k, top_k)).fetchall()
    return [(r[0], r[1]) for r in rows]

# @title Corpus loaders
# These helpers load small corpora into memory. For Tiers 93/100, adapt loading to stay within the memory budget; fixed counts do not establish byte size.
TIER_SIZES = {"tiny": 1_000, "25k": 25_000, "250k": 250_000}


def load_codesearchnet(tier: str = "tiny") -> tuple[list[dict], list[dict]]:
    """Load CodeSearchNet at the chosen tier. Returns (docs, eval_set).

    eval_set is built from docstrings: query is the first nonempty line of the
    docstring; relevant_doc_ids is [the function whose docstring it was].
    """
    from datasets import load_dataset

    if tier not in TIER_SIZES:
        raise ValueError(f"unknown tier {tier!r}; choose from {list(TIER_SIZES)}")
    n = TIER_SIZES[tier]
    use_streaming = True

    ds = load_dataset(
        "code-search-net/code_search_net", "python", split="train",
        streaming=use_streaming,
    )
    docs, eval_pool = [], []
    for row in ds:
        if len(docs) >= n:
            break
        code = row["func_code_string"]
        if not code:
            continue
        doc_id = f"python/{row['repository_name']}/{row['func_name']}::{len(docs)}"
        docs.append({"doc_id": doc_id, "text": code})
        docstring = (row.get("func_documentation_string") or "").strip()
        if docstring:
            first = next((line.strip() for line in docstring.split("\n") if line.strip()), "")
            if first and 20 <= len(first) <= 200:
                eval_pool.append({"query": first, "relevant_doc_ids": [doc_id]})
    if len(docs) < n:
        raise ValueError(f"Requested {n:,} documents but source supplied {len(docs):,}. Use a smaller tier or extend the loader with documented sources.")
    rng = random.Random(42)
    eval_n = min(500, len(eval_pool))
    eval_set = rng.sample(eval_pool, eval_n)
    return docs, eval_set


def load_msmarco(tier: str = "tiny") -> tuple[list[dict], list[dict]]:
    """Load MS MARCO at the chosen tier. Returns (docs, eval_set).

    eval_set is real Bing/SQ queries from the BeIR/msmarco-qrels training split,
    filtered to those whose relevant docs are in the loaded subset.
    """
    from datasets import load_dataset

    if tier not in TIER_SIZES:
        raise ValueError(f"unknown tier {tier!r}; choose from {list(TIER_SIZES)}")
    n = TIER_SIZES[tier]

    corpus_ds = load_dataset("BeIR/msmarco", "corpus", split="corpus", streaming=True)
    queries_ds = load_dataset("BeIR/msmarco", "queries", split="queries", streaming=True)
    qrels_ds = load_dataset("BeIR/msmarco-qrels", split="train", streaming=True)

    docs = []
    docs_in_subset = set()
    for i, row in enumerate(corpus_ds):
        if i >= n:
            break
        doc_id = str(row["_id"])
        text = ((row.get("title") or "") + " " + (row.get("text") or "")).strip()
        docs.append({"doc_id": doc_id, "text": text})
        docs_in_subset.add(doc_id)

    qid_to_text = {str(row["_id"]): row["text"] for row in queries_ds}
    qid_to_relevant = {}
    for row in qrels_ds:
        qid = str(row["query-id"])
        did = str(row["corpus-id"])
        if did in docs_in_subset and int(row["score"]) > 0:
            qid_to_relevant.setdefault(qid, []).append(did)
    eligible = [q for q, r in qid_to_relevant.items() if r and q in qid_to_text]
    if len(docs) < n:
        raise ValueError(f"Requested {n:,} documents but source supplied {len(docs):,}. Use a smaller tier or extend the loader with documented sources.")
    rng = random.Random(42)
    eval_n = min(500, len(eligible))
    eval_set = [{"query": qid_to_text[q], "relevant_doc_ids": qid_to_relevant[q]}
                for q in rng.sample(eligible, eval_n)]
    return docs, eval_set

# @title Evaluation helpers
def recall_at_k(retrieved: list[str], relevant: list[str], k: int = 5) -> float:
    """Fraction of distinct relevant documents found in the first k results."""
    relevant_ids = set(relevant)
    return len(set(retrieved[:k]) & relevant_ids) / len(relevant_ids) if relevant_ids else 0.0


def mean_reciprocal_rank(retrieved: list[str], relevant: list[str]) -> float:
    rel_set = set(relevant)
    for i, doc in enumerate(retrieved, start=1):
        if doc in rel_set:
            return 1.0 / i
    return 0.0


def run_eval(query_fn, eval_set: list[dict], top_k: int = 10,
             memory_limit_mb: float | None = None) -> dict:
    """Run query_fn against an eval set. Reports R@1, R@5, R@10, MRR@10,
    p50/p95 latency, AND peak RSS during the run.

    Pass `memory_limit_mb=2000` for Tier 93 or `memory_limit_mb=1000`
    for Tier 100 to label the RSS report with the tier target. This
    does not enforce or measure the live working-set cap. Report live
    components separately with memory_breakdown and explain how your implementation stays within the memory budget.
    """
    r1, r5, r10, mrrs, lat = [], [], [], [], []
    with MemoryBudget("eval run", limit_mb=memory_limit_mb) as mb:
        for ex in eval_set:
            t0 = time.perf_counter()
            results = query_fn(ex["query"])
            lat.append((time.perf_counter() - t0) * 1000)
            mb.update()
            retrieved = [d for d, _ in results] if (results and isinstance(results[0], tuple)) else (results or [])
            r1.append(recall_at_k(retrieved, ex["relevant_doc_ids"], 1))
            r5.append(recall_at_k(retrieved, ex["relevant_doc_ids"], 5))
            r10.append(recall_at_k(retrieved, ex["relevant_doc_ids"], 10))
            mrrs.append(mean_reciprocal_rank(retrieved[:10], ex["relevant_doc_ids"]))
    return {
        "n": len(eval_set),
        "recall@1": float(np.mean(r1)),
        "recall@5": float(np.mean(r5)),
        "recall@10": float(np.mean(r10)),
        "mrr@10": float(np.mean(mrrs)),
        "p50_ms": float(np.percentile(lat, 50)),
        "p95_ms": float(np.percentile(lat, 95)),
        "peak_mb": mb.peak,
    }

# @title Tier 93 data helper
# Supplied utility: defines a loader; downloads nothing until you call it.
def prepare_tier93_data(data_dir="nanomem_scale"):
    import hashlib, importlib.util, urllib.request
    from pathlib import Path
    root = Path(data_dir)
    root.mkdir(parents=True, exist_ok=True)
    helper_path = root / "nanomem_benchmark_v1.py"
    payload = urllib.request.urlopen(
        "https://cs145.web.app/projects/nanomem/bonus/benchmark_v1.py"
    ).read()
    if hashlib.sha256(payload).hexdigest() != "c3acbaf9bfb77dd4bb471df1b67b8a3d371df09b4e9afc9ba4af96c9baaed1c1":
        raise ValueError("Helper version mismatch")
    helper_path.write_bytes(payload)
    spec = importlib.util.spec_from_file_location("nanomem_benchmark", helper_path)
    utility = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(utility)
    data = utility.load_bonus_benchmark(root, config_sha256="75744dcb478a27f4ef75972361bcf015d626e88633cda85110da122d2ce2456f")
    utility.prepare_documents(data)
    return utility, data

# @title Tiny test data
TINY_CORPUS = [
    {"doc_id": "auth/jwt.py:1", "text": "def verify_jwt(token, secret): return jwt.decode(token, secret, algorithms=['HS256'])"},
    {"doc_id": "auth/jwt.py:2", "text": "def issue_jwt(user_id, secret, ttl=3600): return jwt.encode({'sub': user_id, 'exp': time.time() + ttl}, secret)"},
    {"doc_id": "auth/scopes.py:1", "text": "def verify_jwt_with_scope(token, secret, required_scope): payload = verify_jwt(token, secret); assert required_scope in payload['scopes']"},
    {"doc_id": "config/parser.py:1", "text": "def parse_config(path): with open(path) as f: return json.load(f)"},
    {"doc_id": "config/parser.py:2", "text": "def merge_configs(base, override): merged = dict(base); merged.update(override); return merged"},
    {"doc_id": "utils/load_json.py:1", "text": "def load_json(path): with open(path, 'r') as fh: return json.load(fh)"},
    {"doc_id": "utils/load_json.py:2", "text": "def save_json(obj, path): with open(path, 'w') as fh: json.dump(obj, fh, indent=2)"},
    {"doc_id": "db/connection.py:1", "text": "def connect_postgres(dsn): conn = psycopg2.connect(dsn); conn.autocommit = False; return conn"},
    {"doc_id": "db/connection.py:2", "text": "def close_connection(conn): conn.commit(); conn.close()"},
    {"doc_id": "db/transactions.py:1", "text": "def with_transaction(conn, fn): try: result = fn(conn); conn.commit(); return result; except: conn.rollback(); raise"},
    {"doc_id": "search/tokenize.py:1", "text": "def tokenize(text): return [w.lower() for w in re.split(r'\\W+', text) if w]"},
    {"doc_id": "search/bm25.py:1", "text": "def bm25_score(tf, df, n_docs, doc_len, avg_len, k1=1.2, b=0.75): idf = math.log((n_docs - df + 0.5) / (df + 0.5))"},
    {"doc_id": "cache/redis_client.py:1", "text": "def cache_get(key): return redis.get(key)"},
    {"doc_id": "cache/redis_client.py:2", "text": "def cache_set(key, value, ttl=300): redis.setex(key, ttl, value)"},
    {"doc_id": "models/user.py:1", "text": "class User: def __init__(self, id, email): self.id = id; self.email = email"},
    {"doc_id": "models/user.py:2", "text": "def get_user_by_email(session, email): return session.query(User).filter_by(email=email).first()"},
    {"doc_id": "api/routes.py:1", "text": "@app.route('/login', methods=['POST']) def login(): user = authenticate(request.json); return issue_jwt(user.id, SECRET)"},
    {"doc_id": "api/routes.py:2", "text": "@app.route('/logout', methods=['POST']) def logout(): revoke_token(request.headers['Authorization'])"},
    {"doc_id": "tests/test_auth.py:1", "text": "def test_jwt_roundtrip(): token = issue_jwt('user123', 'secret'); assert verify_jwt(token, 'secret')['sub'] == 'user123'"},
    {"doc_id": "tests/test_config.py:1", "text": "def test_parse_config(): cfg = parse_config('test.json'); assert cfg['db'] == 'postgres'"},
]