"""Shared helpers for the pgvector.app benchmarks.

Every number published on pgvector.app comes from these scripts running against
PostgreSQL 18 + pgvector 0.8.7 (PGDG packages) on a 4 vCPU / 15 GB cloud VM.
"""
import json, os, statistics, time
import numpy as np
import psycopg
from pgvector.psycopg import register_vector
from pgvector import HalfVector, Bit

NP = os.environ.get("NP_DIR", os.path.join(os.path.dirname(__file__), "..", "np"))
OUT = os.environ.get("OUT_DIR", os.path.join(os.path.dirname(__file__), "..", "results"))
os.makedirs(OUT, exist_ok=True)


def connect(autocommit=True):
    conn = psycopg.connect(host="/var/run/postgresql", port=5433, user="postgres",
                           dbname="bench", autocommit=autocommit)
    register_vector(conn)
    return conn


def corpus(n, dim=1536):
    v = np.load(f"{NP}/corpus{dim}.npy", mmap_mode="r")
    return np.ascontiguousarray(v[:n])


def queries():
    return np.load(f"{NP}/queries1536.npy")


def ground_truth(base, q, k=100, metric="ip"):
    """Exact top-k ids (row numbers) by cosine distance (== inner product on unit vectors)."""
    gt = np.empty((len(q), k), dtype=np.int64)
    for s in range(0, len(q), 100):
        sims = q[s:s + 100] @ base.T
        idx = np.argpartition(-sims, k, axis=1)[:, :k]
        order = np.take_along_axis(sims, idx, axis=1).argsort(axis=1)[:, ::-1]
        gt[s:s + 100] = np.take_along_axis(idx, order, axis=1)
    return gt


def recall(found, truth, k=10):
    return len(set(found[:k]) & set(truth[:k])) / k


def load_table(conn, name, vecs, coltype="vector", dim=1536, extra_cols="", start_id=0):
    """Create table and binary-COPY vectors in; id = row number."""
    conn.execute(f"DROP TABLE IF EXISTS {name}")
    conn.execute(f"CREATE TABLE {name} (id int PRIMARY KEY, embedding {coltype}({dim}){extra_cols})")
    t = time.perf_counter()
    with conn.cursor().copy(f"COPY {name} (id, embedding) FROM STDIN WITH (FORMAT BINARY)") as cp:
        cp.set_types(["int4", coltype])
        for i, v in enumerate(vecs):
            cp.write_row([start_id + i, HalfVector(v) if coltype == "halfvec" else v])
    elapsed = time.perf_counter() - t
    conn.execute(f"VACUUM ANALYZE {name}")
    return elapsed


def rel_size(conn, rel):
    return conn.execute("SELECT pg_relation_size(%s::regclass)", (rel,)).fetchone()[0]


def table_size(conn, rel):
    """Heap + TOAST (vectors > ~2KB are TOASTed), excluding indexes."""
    return conn.execute("SELECT pg_table_size(%s::regclass)", (rel,)).fetchone()[0]


def timed(conn, sql, params=None):
    t = time.perf_counter()
    conn.execute(sql, params)
    return time.perf_counter() - t


def run_queries(conn, sql, qvecs, truth=None, k=10, warmup=50, wrap=lambda v: v):
    """Run each query once (after warmup); return recall@k, latency percentiles, rows returned."""
    cur = conn.cursor()
    for v in qvecs[:warmup]:
        cur.execute(sql, (wrap(v),), prepare=True).fetchall()
    lat, rec, nret = [], [], []
    for i, v in enumerate(qvecs):
        t = time.perf_counter()
        rows = cur.execute(sql, (wrap(v),), prepare=True).fetchall()
        lat.append((time.perf_counter() - t) * 1000)
        ids = [r[0] for r in rows]
        nret.append(len(ids))
        if truth is not None:
            rec.append(recall(ids, truth[i], k))
    lat.sort()
    res = {
        "p50_ms": round(lat[len(lat) // 2], 3),
        "p95_ms": round(lat[int(len(lat) * 0.95)], 3),
        "p99_ms": round(lat[int(len(lat) * 0.99)], 3),
        "mean_ms": round(statistics.fmean(lat), 3),
        "qps_1client": round(1000 / statistics.fmean(lat), 1),
        "avg_rows": round(statistics.fmean(nret), 2),
    }
    if rec:
        res["recall"] = round(statistics.fmean(rec), 4)
    return res


def save(name, obj):
    obj = dict(obj)
    obj.setdefault("env", env_info())
    with open(f"{OUT}/{name}.json", "w") as fh:
        json.dump(obj, fh, indent=1, default=lambda o: o.decode() if isinstance(o, bytes) else str(o))
    print("saved", name)


_ENV = None


def env_info():
    global _ENV
    if _ENV is None:
        with connect() as c:
            _ENV = {
                "postgres": c.execute("SHOW server_version").fetchone()[0],
                "pgvector": c.execute("SELECT extversion FROM pg_extension WHERE extname='vector'").fetchone()[0],
                "shared_buffers": c.execute("SHOW shared_buffers").fetchone()[0],
                "maintenance_work_mem": c.execute("SHOW maintenance_work_mem").fetchone()[0],
                "max_parallel_maintenance_workers": c.execute("SHOW max_parallel_maintenance_workers").fetchone()[0],
                "cpu": "4 vCPU Intel Xeon @ 2.10GHz (cloud VM)",
                "ram": "15 GB",
                "dataset": "DBpedia entities, OpenAI text-embedding-3-large (Qdrant/dbpedia-entities-openai3-text-embedding-3-large-1536-1M on Hugging Face)",
            }
    return _ENV


def bge_small():
    """BAAI/bge-small-en-v1.5 via fastembed, using the fp32 ONNX export (Xenova/bge-small-en-v1.5).
    Same embeddings as fastembed's default quantized build (cosine > 0.99999) but ~6x faster on our CPU."""
    from fastembed import TextEmbedding
    from fastembed.common.model_description import PoolingType, ModelSource
    name = "bge-small-en-v1.5-fp32"
    if name not in [m["model"] for m in TextEmbedding.list_supported_models()]:
        TextEmbedding.add_custom_model(model=name, pooling=PoolingType.CLS, normalization=True,
                                       sources=ModelSource(hf="Xenova/bge-small-en-v1.5"), dim=384,
                                       model_file="onnx/model.onnx")
    return TextEmbedding(name)