# cspell:ignore healthz
"""PLaMo Embedding 1B の HTTP wrapper.

`POST /embed` に {"text": "...", "mode": "query"|"document"} を投げると、
2048 次元の float 配列を返す。`POST /chunks` は同じtokenizerを使い、Markdownを
学習時のcontext長に合わせて分割する。tidb-embedder (backfill) と blog-api (検索) が
同じServiceを利用する。
"""
import logging
import os

import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field, model_validator
from transformers import AutoModel, AutoTokenizer

from chunking import CHUNKING_VERSION, chunk_document

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

MODEL_ID = os.environ.get("MODEL_ID", "pfnet/plamo-embedding-1b")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

log.info("loading model=%s device=%s", MODEL_ID, DEVICE)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True).to(DEVICE).eval()
log.info("model loaded")

app = FastAPI()


class EmbedRequest(BaseModel):
    text: str
    mode: str  # "query" | "document"


class EmbedResponse(BaseModel):
    vector: list[float]
    dim: int


class ChunksRequest(BaseModel):
    title: str
    description: str = ""
    content: str
    max_tokens: int = Field(default=1024, ge=64, le=4096)
    overlap_tokens: int = Field(default=128, ge=0)

    @model_validator(mode="after")
    def validate_overlap(self) -> "ChunksRequest":
        if self.overlap_tokens >= self.max_tokens:
            raise ValueError("overlap_tokens must be less than max_tokens")
        return self


class ChunkResponseItem(BaseModel):
    index: int
    heading: str | None
    content: str
    embedding_text: str
    token_count: int


class ChunksResponse(BaseModel):
    version: str
    max_tokens: int
    overlap_tokens: int
    chunks: list[ChunkResponseItem]


@app.post("/embed", response_model=EmbedResponse)
def embed(req: EmbedRequest) -> EmbedResponse:
    if req.mode not in ("query", "document"):
        raise HTTPException(status_code=400, detail="mode must be 'query' or 'document'")
    with torch.inference_mode():
        if req.mode == "query":
            vec = model.encode_query(req.text, tokenizer)
        else:
            vec = model.encode_document(req.text, tokenizer)
    # encode_* は (1, hidden_size) を返す。squeeze して 1 次元にする
    vec_list = vec.squeeze(0).cpu().float().tolist()
    return EmbedResponse(vector=vec_list, dim=len(vec_list))


@app.post("/chunks", response_model=ChunksResponse)
def chunks(req: ChunksRequest) -> ChunksResponse:
    items = chunk_document(
        tokenizer,
        title=req.title,
        description=req.description,
        content=req.content,
        max_tokens=req.max_tokens,
        overlap_tokens=req.overlap_tokens,
    )
    return ChunksResponse(
        version=CHUNKING_VERSION,
        max_tokens=req.max_tokens,
        overlap_tokens=req.overlap_tokens,
        chunks=[ChunkResponseItem(**item.__dict__) for item in items],
    )


@app.get("/healthz")
def healthz() -> dict[str, str]:
    return {"status": "ok"}

"""PLaMO tokenizer に合わせた Markdown document chunking."""

from dataclasses import dataclass
import re
from typing import Protocol


CHUNKING_VERSION = "plamo-markdown-1024-v1"
METADATA_TOKEN_BUDGET = 256


class Tokenizer(Protocol):
    def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: ...

    def decode(self, token_ids: list[int], skip_special_tokens: bool = True) -> str: ...


@dataclass(frozen=True)
class MarkdownSection:
    heading: str | None
    content: str


@dataclass(frozen=True)
class DocumentChunk:
    index: int
    heading: str | None
    content: str
    embedding_text: str
    token_count: int


_HEADING_PATTERN = re.compile(r"^(#{1,6})[ \t]+(.+?)[ \t]*$")
_CLOSING_HASHES_PATTERN = re.compile(r"[ \t]+#+[ \t]*$")
_FENCE_PATTERN = re.compile(r"^[ \t]*(`{3,}|~{3,})")


def split_markdown_sections(content: str) -> list[MarkdownSection]:
    """見出し階層を保ちながら Markdown を section に分割する。"""

    sections: list[MarkdownSection] = []
    headings: list[str] = []
    lines: list[str] = []
    current_heading: str | None = None
    fence_marker: str | None = None

    def flush() -> None:
        nonlocal lines
        body = "\n".join(lines).strip()
        if body:
            sections.append(MarkdownSection(heading=current_heading, content=body))
        lines = []

    for line in content.replace("\r\n", "\n").replace("\r", "\n").split("\n"):
        fence = _FENCE_PATTERN.match(line)
        if fence:
            marker = fence.group(1)[0]
            if fence_marker is None:
                fence_marker = marker
            elif fence_marker == marker:
                fence_marker = None
            lines.append(line)
            continue

        heading = _HEADING_PATTERN.match(line) if fence_marker is None else None
        if heading:
            flush()
            level = len(heading.group(1))
            title = _CLOSING_HASHES_PATTERN.sub("", heading.group(2)).strip()
            headings = headings[: level - 1]
            while len(headings) < level - 1:
                headings.append("")
            headings.append(title)
            current_heading = " > ".join(value for value in headings if value)
            continue

        lines.append(line)

    flush()
    return sections


def _token_ids(tokenizer: Tokenizer, text: str) -> list[int]:
    return list(tokenizer.encode(text, add_special_tokens=False))


def _token_count(tokenizer: Tokenizer, text: str) -> int:
    return len(_token_ids(tokenizer, text))


def _metadata_prefix(
    tokenizer: Tokenizer,
    title: str,
    description: str,
    heading: str | None,
    max_tokens: int,
) -> str:
    fields = [f"タイトル: {title.strip()}"]
    if heading:
        fields.append(f"見出し: {heading.strip()}")
    if description.strip():
        fields.append(f"概要: {description.strip()}")

    header = "\n".join(fields)
    suffix = "\n本文:\n"
    # metadata が長くても本文用の token budget を確保する。タイトル、見出し、概要の
    # 順なので、切り詰めが必要な場合も検索に重要な情報を優先できる。
    budget = min(METADATA_TOKEN_BUDGET, max_tokens // 2)
    suffix_tokens = _token_ids(tokenizer, suffix)
    header_tokens = _token_ids(tokenizer, header)
    header_budget = max(1, budget - len(suffix_tokens))
    if len(header_tokens) > header_budget:
        header = tokenizer.decode(
            header_tokens[:header_budget], skip_special_tokens=True
        ).strip()
    return f"{header}{suffix}"


def chunk_document(
    tokenizer: Tokenizer,
    *,
    title: str,
    description: str,
    content: str,
    max_tokens: int = 1024,
    overlap_tokens: int = 128,
) -> list[DocumentChunk]:
    """Markdown section を PLaMO の token 数で window 分割する。"""

    if not 64 <= max_tokens <= 4096:
        raise ValueError("max_tokens must be between 64 and 4096")
    if not 0 <= overlap_tokens < max_tokens:
        raise ValueError("overlap_tokens must be between 0 and max_tokens - 1")

    sections = split_markdown_sections(content)
    if not sections:
        sections = [MarkdownSection(heading=None, content="")]

    chunks: list[DocumentChunk] = []
    for section in sections:
        prefix = _metadata_prefix(
            tokenizer, title, description, section.heading, max_tokens
        )
        prefix_count = _token_count(tokenizer, prefix)
        body_budget = max_tokens - prefix_count
        if body_budget <= 0:
            raise ValueError("metadata leaves no token budget for content")

        body_tokens = _token_ids(tokenizer, section.content)
        if not body_tokens:
            embedding_text = prefix.rstrip()
            chunks.append(
                DocumentChunk(
                    index=len(chunks),
                    heading=section.heading,
                    content="",
                    embedding_text=embedding_text,
                    token_count=_token_count(tokenizer, embedding_text),
                )
            )
            continue

        start = 0
        while start < len(body_tokens):
            end = min(start + body_budget, len(body_tokens))
            window = body_tokens[start:end]
            body = tokenizer.decode(window, skip_special_tokens=True).strip()
            embedding_text = f"{prefix}{body}"
            token_count = _token_count(tokenizer, embedding_text)

            # SentencePiece は prefix と本文を結合した境界で再tokenize結果がわずかに
            # 変わり得る。最終入力を数え直し、1024 tokenを確実に超えないよう縮める。
            while token_count > max_tokens and len(window) > 1:
                overflow = token_count - max_tokens
                window = window[: max(1, len(window) - overflow - 1)]
                end = start + len(window)
                body = tokenizer.decode(window, skip_special_tokens=True).strip()
                embedding_text = f"{prefix}{body}"
                token_count = _token_count(tokenizer, embedding_text)

            chunks.append(
                DocumentChunk(
                    index=len(chunks),
                    heading=section.heading,
                    content=body,
                    embedding_text=embedding_text,
                    token_count=token_count,
                )
            )
            if end >= len(body_tokens):
                break
            effective_overlap = min(overlap_tokens, len(window) - 1)
            start = end - effective_overlap

    return chunks

# syntax=docker/dockerfile:1
FROM python:3.12-slim

RUN apt-get update && apt-get install -y --no-install-recommends \
      ca-certificates \
      libjemalloc2 \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app

# CPU版PyTorch(クラスタにGPUなし)
RUN pip install --no-cache-dir \
      --index-url https://download.pytorch.org/whl/cpu \
      torch==2.5.1 \
    && pip install --no-cache-dir \
      transformers==4.46.0 \
      sentencepiece==0.2.0 \
      fastapi==0.115.4 \
      uvicorn==0.32.0 \
      pydantic==2.9.2

COPY server.py chunking.py /app/

# モデル(約2.3GB)はビルド時に取得してイメージへ焼き込み、実行時はHF cacheからロードする
ARG MODEL_ID=pfnet/plamo-embedding-1b
ENV MODEL_ID=${MODEL_ID}
RUN python -c "from huggingface_hub import snapshot_download; \
  snapshot_download(repo_id='${MODEL_ID}')"

# メモリアロケータをjemallocへ差し替える(経緯は後述)
ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2
ENV MALLOC_CONF=background_thread:true,dirty_decay_ms:10000

EXPOSE 8080
CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8080"]

apiVersion: v1
kind: Namespace
metadata:
  name: plamo-embedding
---
apiVersion: apps/v1
kind: Deployment
metadata:
  name: plamo-embedding
  namespace: plamo-embedding
spec:
  # node1 は TiFlash の vector index build 用に空け、node2 / node3 へ 1 Pod ずつ配置する。
  replicas: 2
  # maxSurge=0 で更新中にも追加 Pod を作らず、モデル分のメモリ増加を防ぐ。
  strategy:
    type: RollingUpdate
    rollingUpdate:
      maxSurge: 0
      maxUnavailable: 1
  selector:
    matchLabels:
      app: plamo-embedding
  template:
    metadata:
      labels:
        app: plamo-embedding
    spec:
      affinity:
        nodeAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
            nodeSelectorTerms:
              - matchExpressions:
                  - key: kubernetes.io/hostname
                    operator: In
                    values:
                      - node2
                      - node3
        podAntiAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
            - labelSelector:
                matchLabels:
                  app: plamo-embedding
              topologyKey: kubernetes.io/hostname
      containers:
        - name: server
          image: ghcr.io/shuntaka9576/plamo-embedding:latest
          imagePullPolicy: Always
          ports:
            - containerPort: 8080
          readinessProbe:
            httpGet:
              path: /healthz
              port: 8080
            periodSeconds: 10
            timeoutSeconds: 3
          startupProbe:
            httpGet:
              path: /healthz
              port: 8080
            # node ごとの初回 image pull と model load を最大 10 分待つ。
            periodSeconds: 10
            timeoutSeconds: 3
            failureThreshold: 60
          livenessProbe:
            httpGet:
              path: /healthz
              port: 8080
            periodSeconds: 30
            timeoutSeconds: 5
            failureThreshold: 3
          resources:
            requests:
              cpu: '500m'
              memory: '4Gi'
            limits:
              # 常駐約5.5GiB + 同時リクエスト時のアクティベーションで6Giでは
              # OOMKillされた(2026-07-17)ため8Giへ引き上げ
              memory: '8Gi'
---
apiVersion: v1
kind: Service
metadata:
  name: plamo-embedding
  namespace: plamo-embedding
spec:
  type: ClusterIP
  selector:
    app: plamo-embedding
  ports:
    - port: 80
      targetPort: 8080

# /// script
# requires-python = ">=3.11"
# dependencies = ["pymysql"]
# ///
"""日本語Wikipedia(cirrus_search_indexダンプ)をチャンク化してTiDBへ投入する。

シャード(.json.bz2)をストリームで読み、PLaMoサービスの /chunks でチャンク化、
/embed (mode=document) で2048次元ベクトルを生成して bench_wiki.wiki_embedding_chunks
へINSERTする。--target-chunks に達したら終了する。

再開可能: 投入済みのpage_idはスキップし、既存チャンク数をターゲットに算入する。
INSERTはページ単位のトランザクションなので、中断してもページ単位で整合する。

Usage:
  export TAILNET=$(tailscale status --json | jq -r '.MagicDNSSuffix')
  uv run ingest_wiki.py jawiki_content-20260712-00000.json.bz2 \
    --target-chunks 10000 --concurrency 4
  uv run ingest_wiki.py <shard> --dry-run   # チャンク数の見積もりだけ(embed/INSERTなし)
"""

import argparse
import bz2
import gzip
import json
import os
import sys
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor

import pymysql


def log(msg: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


def http_post_json(url: str, payload: dict, timeout: int, retries: int = 20) -> dict:
    body = json.dumps(payload).encode()
    last_err: Exception | None = None
    for attempt in range(1, retries + 1):
        try:
            req = urllib.request.Request(
                url, data=body, headers={"Content-Type": "application/json"}
            )
            with urllib.request.urlopen(req, timeout=timeout) as resp:
                return json.load(resp)
        except Exception as e:  # noqa: BLE001
            last_err = e
            log(f"HTTP retry {attempt}/{retries} {url}: {e}")
            # Pod再起動時のモデルロード(最大10分)を跨げるよう、計16分程度まで粘る
            time.sleep(min(5 * attempt, 120))
    raise RuntimeError(f"HTTP failed after {retries} retries: {url}: {last_err}")


def open_dump(path: str):
    if path.endswith(".bz2"):
        return bz2.open(path, "rt", encoding="utf-8")
    if path.endswith(".gz"):
        return gzip.open(path, "rt", encoding="utf-8")
    return open(path, encoding="utf-8")


def iter_pages(paths: list[str], min_chars: int):
    """ダンプのNDJSONからドキュメント行だけを取り出す。

    形式: {"index": {"_id": N}} 行と {"page_id": N, "title": ..., "text": ...} 行の2行1組。
    ドキュメント行はpage_idを含むため、index行は読み飛ばすだけでよい。
    """
    for path in paths:
        log(f"reading {path}")
        with open_dump(path) as f:
            for line in f:
                try:
                    doc = json.loads(line)
                except json.JSONDecodeError:
                    continue
                if "index" in doc and "page_id" not in doc:
                    continue
                page_id = doc.get("page_id")
                title = doc.get("title")
                text = doc.get("text")
                if not page_id or not title or not text:
                    continue
                if doc.get("namespace", 0) != 0:
                    continue
                if len(text) < min_chars:
                    continue
                yield {"page_id": page_id, "title": title, "text": text}


def main() -> int:
    tailnet = os.environ.get("TAILNET", "")
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("files", nargs="+", help="ダンプシャード (.json.bz2 / .json.gz)")
    ap.add_argument(
        "--embed-endpoint",
        default=os.environ.get(
            "PLAMO_EMBED_ENDPOINT",
            f"http://plamo-embedding.{tailnet}" if tailnet else "",
        ),
    )
    ap.add_argument("--tidb-host", default=f"tidb.{tailnet}" if tailnet else "")
    ap.add_argument("--tidb-port", type=int, default=4000)
    ap.add_argument("--tidb-user", default="root")
    ap.add_argument("--database", default="bench_wiki")
    ap.add_argument("--table", default="wiki_embedding_chunks")
    ap.add_argument("--target-chunks", type=int, default=10_000)
    ap.add_argument("--concurrency", type=int, default=4)
    ap.add_argument("--max-tokens", type=int, default=1024)
    ap.add_argument("--overlap-tokens", type=int, default=128)
    ap.add_argument("--min-chars", type=int, default=200)
    ap.add_argument("--dry-run", action="store_true", help="embed/INSERTせず見積もりのみ")
    args = ap.parse_args()

    if not args.embed_endpoint:
        ap.error("--embed-endpoint (または TAILNET / PLAMO_EMBED_ENDPOINT) が必要")

    conn = None
    done_pages: set[int] = set()
    inserted = 0
    if not args.dry_run:
        if not args.tidb_host:
            ap.error("--tidb-host (または TAILNET) が必要")
        conn = pymysql.connect(
            host=args.tidb_host,
            port=args.tidb_port,
            user=args.tidb_user,
            database=args.database,
            charset="utf8mb4",
            autocommit=False,
        )
        with conn.cursor() as cur:
            cur.execute(f"SELECT DISTINCT page_id FROM {args.table}")
            done_pages = {row[0] for row in cur.fetchall()}
            cur.execute(f"SELECT COUNT(*) FROM {args.table}")
            inserted = cur.fetchone()[0]
        log(f"resume: {len(done_pages)} pages / {inserted} chunks already in DB")

    executor = ThreadPoolExecutor(max_workers=args.concurrency)

    def embed(text: str) -> list[float]:
        resp = http_post_json(
            f"{args.embed_endpoint}/embed",
            {"text": text, "mode": "document"},
            timeout=120,
        )
        return resp["vector"]

    def flush(pages: list[dict]) -> int:
        """保留ページのチャンクを並列embedし、ページ単位のトランザクションでINSERTする。"""
        jobs = [c for p in pages for c in p["chunks"]]
        vectors = list(executor.map(lambda c: embed(c["embedding_text"]), jobs))
        for c, v in zip(jobs, vectors):
            c["vector"] = v
        # 長時間実行中にTCPが切れていた場合はここで張り直す
        conn.ping(reconnect=True)
        n = 0
        for p in pages:
            try:
                with conn.cursor() as cur:
                    cur.executemany(
                        f"INSERT INTO {args.table}"
                        " (page_id, title, chunk_index, content, token_count, embedding)"
                        " VALUES (%s, %s, %s, %s, %s, %s)",
                        [
                            (
                                p["page_id"],
                                p["title"][:512],
                                c["index"],
                                c["content"],
                                c["token_count"],
                                json.dumps(c["vector"], separators=(",", ":")),
                            )
                            for c in p["chunks"]
                        ],
                    )
                conn.commit()
            except pymysql.err.IntegrityError:
                # コミット成功後に応答だけ失われた再実行ケース。UNIQUE制約で検出し投入済みとして扱う
                conn.rollback()
            else:
                n += len(p["chunks"])
        return n

    t0 = time.time()
    session_start = inserted  # 今回セッション開始時点の投入済み件数(rate計算用)
    planned = inserted  # 投入済み + 保留分
    pending: list[dict] = []
    pending_chunks = 0
    total_pages = 0

    try:
        for page in iter_pages(args.files, args.min_chars):
            if planned >= args.target_chunks:
                break
            if page["page_id"] in done_pages:
                continue

            chunk_resp = http_post_json(
                f"{args.embed_endpoint}/chunks",
                {
                    "title": page["title"],
                    "description": "",
                    "content": page["text"],
                    "max_tokens": args.max_tokens,
                    "overlap_tokens": args.overlap_tokens,
                },
                timeout=300,
            )
            chunks = chunk_resp["chunks"]
            if not chunks:
                continue

            total_pages += 1
            planned += len(chunks)

            if args.dry_run:
                inserted = planned
                if total_pages % 200 == 0:
                    log(f"dry-run: {total_pages} pages -> {planned} chunks")
                continue

            page["chunks"] = chunks
            pending.append(page)
            pending_chunks += len(chunks)

            if pending_chunks >= args.concurrency * 2:
                inserted += flush(pending)
                pending, pending_chunks = [], 0
                rate = (inserted - session_start) / max(time.time() - t0, 1)
                remain = max(args.target_chunks - inserted, 0)
                eta_h = remain / rate / 3600 if rate > 0 else float("inf")
                log(
                    f"{inserted}/{args.target_chunks} chunks"
                    f" ({total_pages} pages, {rate:.2f} chunks/s, ETA {eta_h:.1f}h)"
                )

        if pending and not args.dry_run:
            inserted += flush(pending)
    except KeyboardInterrupt:
        log("interrupted: 投入済みページまでは保存されています(再実行で再開)")
        return 130
    finally:
        executor.shutdown(wait=False, cancel_futures=True)
        if conn:
            conn.close()

    if args.dry_run:
        log(f"dry-run done: {total_pages} pages -> {planned} chunks")
    else:
        log(f"done: {inserted} chunks ({total_pages} pages this run)")
    return 0


if __name__ == "__main__":
    sys.exit(main())

# /// script
# requires-python = ">=3.11"
# dependencies = ["pymysql>=1.1"]
# ///
"""TiDBベクトル検索の3経路レイテンシとrecall@10を計測する。

同一クエリを繰り返すとTiKV経路がTiDBのcoprocessor cacheに乗って
実測にならないため、異なる20クエリを各1回ずつ流して分布を取る。

usage: TAILNET=... uv run bench_search.py
"""
import json
import math
import os
import time
import urllib.request

import pymysql

TAILNET = os.environ["TAILNET"]
EMBED_URL = f"http://plamo-embedding.{TAILNET}/embed"
QUERIES = [
    "日本の城の石垣の構造",
    "相対性理論における時間の遅れ",
    "江戸時代の参勤交代の制度",
    "ワクチンで免疫がつく仕組み",
    "新幹線が開業するまでの歴史",
    "火山の噴火が起きる仕組み",
    "平安時代の貴族の暮らし",
    "サッカーのオフサイドのルール",
    "株式会社と合同会社の違い",
    "太陽系の惑星の並び順",
    "発酵食品と微生物の関係",
    "CPUが計算を実行する仕組み",
    "織田信長と本能寺の変",
    "地震の震度とマグニチュードの違い",
    "世界遺産に登録される条件",
    "オリンピックの開催地の決め方",
    "日本酒の醸造工程",
    "遺伝子の突然変異と進化",
    "気候変動が農業に与える影響",
    "印象派の画家と代表作",
]
# (hint, ORDER BY式)。tiflash_exactの「+ 0」はANNインデックス回避(本文参照)
PATTERNS = {
    "tikv_full":     ("READ_FROM_STORAGE(TIKV[c])",    "VEC_COSINE_DISTANCE(embedding, %s)"),
    "tiflash_exact": ("READ_FROM_STORAGE(TIFLASH[c])", "VEC_COSINE_DISTANCE(embedding, %s) + 0"),
    "tiflash_ann":   ("READ_FROM_STORAGE(TIFLASH[c])", "VEC_COSINE_DISTANCE(embedding, %s)"),
}


def embed_query(text: str) -> tuple[str, float]:
    body = json.dumps({"text": text, "mode": "query"}).encode()
    req = urllib.request.Request(EMBED_URL, data=body,
                                 headers={"Content-Type": "application/json"})
    t0 = time.perf_counter()
    vec = json.load(urllib.request.urlopen(req, timeout=60))["vector"]
    return json.dumps(vec), time.perf_counter() - t0


def search(cur, pattern: str, qvec: str) -> tuple[tuple, float]:
    hint, order_by = PATTERNS[pattern]
    sql = (f"SELECT /*+ {hint} */ chunk_id, title, "
           f"VEC_COSINE_DISTANCE(embedding, %s) AS distance "
           f"FROM wiki_embedding_chunks AS c "
           f"ORDER BY {order_by} LIMIT 10")
    t0 = time.perf_counter()
    cur.execute(sql, [qvec] * sql.count("%s"))
    return cur.fetchall(), time.perf_counter() - t0


def stats(label: str, samples: list[float]) -> None:
    s = sorted(samples)
    n = len(s)
    p50 = s[math.ceil(0.50 * n) - 1] * 1000
    p95 = s[math.ceil(0.95 * n) - 1] * 1000
    print(f"{label:14s} min={s[0] * 1000:6.1f}ms p50={p50:6.1f}ms "
          f"p95={p95:6.1f}ms max={s[-1] * 1000:6.1f}ms")


def main() -> None:
    conn = pymysql.connect(host=f"tidb.{TAILNET}", port=4000, user="root",
                           database="bench_wiki")
    cur = conn.cursor()

    # ウォームアップ(計測対象外のクエリで各経路を一度ずつ温める)
    warm, _ = embed_query("ウォームアップ用のクエリ")
    for pattern in PATTERNS:
        search(cur, pattern, warm)

    embed_times: list[float] = []
    times: dict[str, list[float]] = {p: [] for p in PATTERNS}
    recalls: list[float] = []
    for q in QUERIES:
        qvec, embed_sec = embed_query(q)
        embed_times.append(embed_sec)
        rows = {}
        for pattern in PATTERNS:
            rows[pattern], sec = search(cur, pattern, qvec)
            times[pattern].append(sec)
        exact = {r[0] for r in rows["tiflash_exact"]}
        ann = {r[0] for r in rows["tiflash_ann"]}
        recalls.append(len(exact & ann) / len(exact))
        print(f"recall@10={recalls[-1]:.2f} {q}")

    print(f"mean recall@10={sum(recalls) / len(recalls):.2f} (n={len(QUERIES)})")
    stats("query_embed", embed_times)
    for pattern in PATTERNS:
        stats(pattern, times[pattern])


if __name__ == "__main__":
    main()