"""Opt-in on-host three-resident-model CPU/RSS benchmark.

Run from service root with model snapshots downloaded and API_KEY configured:
BENCH_LONG=1 python scripts/bench_local.py
Reports are checkpointed after every case to survive failures/OOM.
"""
from __future__ import annotations

import json
import os
import threading
import time
from datetime import datetime, timezone
from pathlib import Path

from app.embedding import MODEL_ID as E5_ID
from app.models import ModelRegistry, QWEN_ID, COLBERT_ID, memory_mb

REPORT = Path(os.getenv("BENCH_REPORT", "bench-result.json"))
LONG = os.getenv("BENCH_LONG", "0") == "1"
rows = []
report = {"time": datetime.now(timezone.utc).isoformat(),
          "long_context": LONG, "startup": None, "cases": rows}


def save():
    REPORT.write_text(json.dumps(report, indent=2, ensure_ascii=False),
                      encoding="utf-8")


def repeated_text(tokenizer, max_tokens: int, prefix: str = "") -> str:
    """Find a repetition count whose true tokenizer length is <= target.

    Model-specific input overhead (query instructions/markers) is counted.
    A result slightly shorter than target is reported, not mislabeled.
    """
    base = "Planet earth history. "
    def length(n):
        return len(tokenizer(prefix + base * n, add_special_tokens=True,
                             truncation=False)["input_ids"])
    lo, hi = 0, max_tokens
    while lo < hi:
        mid = (lo + hi + 1) // 2
        if length(mid) <= max_tokens:
            lo = mid
        else:
            hi = mid - 1
    if lo == 0:
        raise RuntimeError("Model prefix exceeds the test target")
    return base * lo


def measure(name, runner, token_count=None):
    readings = []
    finish = threading.Event()

    def monitor():
        while not finish.wait(0.05):
            value = memory_mb()["rss_mb"]
            if value is not None:
                readings.append(value)

    sampler = threading.Thread(target=monitor, daemon=True)
    sampler.start()
    started = time.perf_counter()
    result = {"case": name, "requested_tokens": token_count}
    try:
        payload = runner()
        result["status"] = "pass"
        result["output_tokens"] = payload.tokens_count
        if hasattr(payload, "embedding"):
            result["dimensions"] = len(payload.embedding)
        else:
            result["dimensions"] = len(payload.token_embeddings[0])
            result["vector_count"] = len(payload.token_embeddings)
    except Exception as error:
        result["status"] = "fail"
        result["error"] = f"{type(error).__name__}: {error}"
    finally:
        finish.set()
        sampler.join()
        result["elapsed_seconds"] = round(time.perf_counter() - started, 3)
        result["rss_mb"] = memory_mb()["rss_mb"]
        result["peak_sampled_rss_mb"] = max(readings, default=None)
        result["peak_process_rss_mb"] = memory_mb()["peak_rss_mb"]
        rows.append(result)
        save()
        print(result, flush=True)


def main():
    started = time.perf_counter()
    registry = ModelRegistry.load()  # all 3 weights resident, once
    report["startup"] = {
        "seconds": round(time.perf_counter() - started, 3),
        "per_model_rss": registry.startup_memory,
        "post_load_ram": memory_mb(),
    }
    save()

    for kind in ("query", "passage"):
        for language, text in (
            ("en", "A long voyage through the forest."),
            ("ru", "Долгое путешествие сквозь лес."),
            ("he", "מסע ארוך דרך היער."),
        ):
            measure(f"e5-{kind}-{language}",
                    lambda t=text, k=kind: registry.e5.encode(t, k))
            measure(f"qwen-{kind}-{language}",
                    lambda t=text, k=kind: registry.qwen.encode(t, k))
            measure(f"colbert-{kind}-{language}",
                    lambda t=text, k=kind: registry.colbert.encode(t, k))

    if LONG:
        for target in (8192, 32768):
            text = repeated_text(registry.qwen.tokenizer, target)
            actual = len(registry.qwen.tokenizer(
                text, add_special_tokens=True, truncation=False)["input_ids"])
            measure(f"qwen-passage-{target}", lambda t=text: registry.qwen.encode(t, "passage"), actual)

        model = registry.colbert.model
        prefix = model.document_prefix
        text = repeated_text(model.tokenizer, 8192, prefix)
        count = len(model.tokenizer(
            prefix + text, add_special_tokens=True, truncation=False)["input_ids"])
        measure("colbert-passage-8192",
                lambda: registry.colbert.encode(text, "passage"), count)


if __name__ == "__main__":
    main()