#!/usr/bin/env python3
"""Echo Space pilot client — the file we hand to pilot partners (reference client).

  export ECHOSPACE_API_KEY=es_pilot_...
  export ECHOSPACE_URL=https://api.mustafagalipayhan.com

  # 1. verdict from SAMPLES (drawn here with a seed; the corpus stays on your machine)
  python3 echospace_pilot.py forecast --target target/ --pool pool/ --label run1 \
      [--model-params 7000000000 --keep 0.25 --language en --regime pretraining] \
      [--table --out fc_run1.json]

  # 2. keep list computed LOCALLY over the whole pool from the returned scoring table
  python3 echospace_pilot.py select-local --forecast fc_run1.json --pool pool/ --keep 0.25 --out keep.json

  # 2'. or, for small pools, let the service compute the keep list (uploads the whole pool)
  python3 echospace_pilot.py select --target target/ --pool pool/ --keep 0.25 --out keep.json

  # 3. the receipt
  python3 echospace_pilot.py outcome --forecast fc_xxxxxxxx --selected 1.91 --baseline 2.05

  python3 echospace_pilot.py selftest      # checks this file's hashing against the service's vectors

A directory = one document per file (.txt/.md/.jsonl with a "text" field). A single file
is one document (.jsonl: one per line). Samples are sent over HTTPS gzip-compressed,
processed in memory and never stored by the service. `select-local` needs numpy.
"""
import argparse, gzip, hashlib, json, os, re, ssl, sys, urllib.request, zlib

try:  # python.org builds on macOS ship without root certificates; certifi fixes that
    import certifi
    _CTX = ssl.create_default_context(cafile=certifi.where())
except ImportError:
    _CTX = ssl.create_default_context()

CLIENT = "echospace-pilot-client/0.2"
HASHER = "es-h1"
NH = 16384
URL = os.environ.get("ECHOSPACE_URL", "http://127.0.0.1:8787").rstrip("/")
KEY = os.environ.get("ECHOSPACE_API_KEY", "")
_WORD = re.compile(r"\S+")


# ---------------- documents (streamed; nothing is held that need not be) ----------------

def iter_docs(path):
    if os.path.isdir(path):
        for root, _, files in os.walk(path):
            for f in sorted(files):
                yield from iter_docs(os.path.join(root, f))
        return
    if path.endswith(".jsonl"):
        with open(path, encoding="utf-8", errors="replace") as fh:
            for line in fh:
                if line.strip():
                    yield json.loads(line)["text"]
        return
    with open(path, encoding="utf-8", errors="replace") as fh:
        yield fh.read()


def count_docs(path):
    n = w = 0
    for d in iter_docs(path):
        n += 1; w += len(_WORD.findall(d))
    return n, w


def sample_docs(path, word_budget, seed):
    """Seeded random sample of whole documents whose total is about `word_budget` words.
    Two streaming passes: count, then keep the documents whose rank in a seeded
    permutation falls below the needed count. Returns (docs, total_documents)."""
    import numpy as np
    n, w = count_docs(path)
    if n == 0:
        sys.exit(f"no documents under {path}")
    if w <= word_budget:
        return list(iter_docs(path)), n
    k = max(1, int(round(n * word_budget / w)))
    rank = np.empty(n, dtype=np.int64)
    rank[np.random.default_rng(seed).permutation(n)] = np.arange(n)
    return [d for i, d in enumerate(iter_docs(path)) if rank[i] < k], n


# ---------------- HTTP ----------------

def call(ep, body=None, method=None):
    data = None
    headers = {"Authorization": f"Bearer {KEY}", "User-Agent": CLIENT}
    if body is not None:
        data = gzip.compress(json.dumps(body).encode(), compresslevel=6)
        headers.update({"Content-Type": "application/json", "Content-Encoding": "gzip"})
    req = urllib.request.Request(f"{URL}/v1/{ep}", method=method or ("POST" if body is not None else "GET"),
                                 headers=headers, data=data)
    try:
        with urllib.request.urlopen(req, timeout=900, context=_CTX) as r:
            return json.load(r)
    except urllib.error.HTTPError as e:
        sys.exit(f"HTTP {e.code}: {e.read().decode()}")
    except urllib.error.URLError as e:
        if "CERTIFICATE_VERIFY_FAILED" in str(e):
            sys.exit("TLS certificate store missing on this Python: run `pip install certifi` "
                     "(or macOS 'Install Certificates.command') and retry")
        raise


# ---------------- local scoring (must match the service bit for bit; see selftest) --------

def word_hashes(text):
    return [zlib.crc32(t.encode("utf-8", "replace")) & 0xFFFFFFFF for t in _WORD.findall(text.lower())]


def buckets(text):
    hs = word_hashes(text)
    return [(hs[i] * 1000003 + hs[i + 1]) % NH for i in range(len(hs) - 1)]


def decode_table(tb):
    import numpy as np
    import base64
    if tb["hasher"] != HASHER:
        sys.exit(f"scoring table hasher {tb['hasher']} != this client's {HASHER}; update the client")
    raw = base64.b64decode(tb["values"])
    if hashlib.sha256(raw).hexdigest() != tb["sha256"]:
        sys.exit("scoring table checksum mismatch; download it again")
    t = np.frombuffer(raw, dtype="<f4").astype(np.float64)
    assert len(t) == tb["buckets"] == NH
    return t


def score_pool(table, docs, keep_fraction, seed=0):
    """Mean table weight over each document's bigrams, then Gumbel-perturbed top-k.
    `docs` is any iterable of strings; memory = two floats per document."""
    import numpy as np
    sums, cnts = [], []
    for d in docs:
        hs = np.array(word_hashes(d), dtype=np.int64)
        if len(hs) < 2:
            sums.append(0.0); cnts.append(0); continue
        b = (hs[:-1] * 1000003 + hs[1:]) % NH
        sums.append(float(table[b].sum())); cnts.append(len(b))
    sums, cnt = np.array(sums), np.array(cnts, dtype=np.float64)
    has = cnt > 0
    s_doc = np.full(len(cnt), -np.inf); s_doc[has] = sums[has] / cnt[has]
    eligible = np.flatnonzero(has)
    keep = max(1, int(round(keep_fraction * len(eligible))))
    rng = np.random.default_rng(9000 + seed)
    mean_len = float(cnt[has].mean())
    keys = s_doc[eligible] * mean_len + rng.gumbel(size=len(eligible))
    top = eligible[np.argpartition(-keys, keep - 1)[:keep]] if keep < len(eligible) else eligible
    return np.sort(top).tolist(), int(len(eligible)), int(len(cnt))


def selftest(verbose=True):
    v = call("hasher")
    if v["hasher"] != HASHER or v["buckets"] != NH:
        sys.exit(f"service hasher {v['hasher']}/{v['buckets']} != client {HASHER}/{NH}: update the client")
    for tv in v["test_vectors"]:
        assert word_hashes(tv["text"]) == tv["word_hashes"], tv
        assert buckets(tv["text"]) == tv["buckets"], tv
    if verbose:
        print(f"hasher {HASHER}: {len(v['test_vectors'])} vectors match the service")
    return True


# ---------------- commands ----------------

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    for name in ("forecast", "select"):
        s = sub.add_parser(name); s.add_argument("--target", required=True); s.add_argument("--pool", required=True)
        s.add_argument("--label"); s.add_argument("--keep", type=float, default=0.25); s.add_argument("--out")
        s.add_argument("--seed", type=int, default=0, help="sampling seed (recorded on the receipt)")
        s.add_argument("--pool-sample-words", type=int, default=5_000_000)
        s.add_argument("--target-sample-words", type=int, default=2_000_000)
        s.add_argument("--model-params", type=int); s.add_argument("--language"); s.add_argument("--regime")
        s.add_argument("--table", action="store_true", help="forecast: ask for the scoring table on SELECT")
    sl = sub.add_parser("select-local"); sl.add_argument("--forecast", required=True, help="forecast JSON saved with --out")
    sl.add_argument("--pool", required=True); sl.add_argument("--keep", type=float, default=0.25)
    sl.add_argument("--out", required=True); sl.add_argument("--seed", type=int, default=0)
    o = sub.add_parser("outcome"); o.add_argument("--forecast", required=True); o.add_argument("--selected", type=float)
    o.add_argument("--baseline", type=float); o.add_argument("--metric", default="held-out loss")
    o.add_argument("--acted", default="both"); o.add_argument("--notes")
    sub.add_parser("me"); sub.add_parser("selftest")
    a = ap.parse_args()
    if not KEY:
        sys.exit("set ECHOSPACE_API_KEY")
    if a.cmd == "me":
        print(json.dumps(call("me"), indent=1)); return
    if a.cmd == "selftest":
        selftest(); return
    if a.cmd == "outcome":
        print(json.dumps(call("outcome", {"forecast_id": a.forecast, "acted": a.acted, "metric_name": a.metric,
                                          "selected_value": a.selected, "baseline_value": a.baseline, "notes": a.notes}), indent=1)); return
    if a.cmd == "select-local":
        fc = json.load(open(a.forecast))
        if not fc.get("scoring_table"):
            sys.exit(f"forecast {fc.get('forecast_id')} carries no scoring table (verdict {fc.get('verdict')}): "
                     "if the verdict is DO-NOT-SELECT, train on a random keep list of the same size")
        selftest(verbose=False)
        table = decode_table(fc["scoring_table"])
        keep, n_el, n_all = score_pool(table, iter_docs(a.pool), a.keep, a.seed)
        keep_sha = hashlib.sha256(json.dumps(keep).encode()).hexdigest()
        rep = call("selections", {"forecast_id": fc["forecast_id"], "keep_fraction": a.keep, "pool_total_documents": n_all,
                                  "kept": len(keep), "table_sha256": fc["scoring_table"]["sha256"], "keep_sha256": keep_sha,
                                  "client": CLIENT})
        json.dump({"forecast_id": fc["forecast_id"], "selection_id": rep["selection_id"], "mode": "local",
                   "keep_fraction": a.keep, "pool_documents": n_all, "eligible_documents": n_el, "kept": len(keep),
                   "keep_sha256": keep_sha, "hasher": HASHER, "keep_indices": keep}, open(a.out, "w"))
        print(json.dumps({**rep, "kept": len(keep), "pool_documents": n_all, "keep_indices": f"written to {a.out}"}, indent=1))
        return
    # forecast / select: samples in, verdict out
    if a.cmd == "select":
        pool, n_pool = list(iter_docs(a.pool)), None    # server-side keep list needs the whole pool
        n_pool = len(pool)
    else:
        pool, n_pool = sample_docs(a.pool, a.pool_sample_words, a.seed)
    target, n_target = sample_docs(a.target, a.target_sample_words, a.seed)
    ctx = {k: v for k, v in {"model_params": a.model_params, "keep_fraction": a.keep, "language": a.language,
                             "regime": a.regime}.items() if v is not None}
    body = {"target": target, "pool": pool, "label": a.label, "context": ctx or None,
            "provenance": {"client": CLIENT, "pool_total_documents": n_pool, "target_total_documents": n_target,
                           "sample_seed": a.seed}}
    if a.cmd == "select":
        body["keep_fraction"] = a.keep
    else:
        body["table"] = a.table
    r = call(a.cmd, body)
    shown = dict(r)
    if a.out:
        json.dump(r, open(a.out, "w"))
        for big in ("keep_indices", "scoring_table"):
            if big in shown:
                shown[big] = f"written to {a.out}"
    elif "scoring_table" in shown and shown["scoring_table"]:
        shown["scoring_table"] = "(use --out to save it for select-local)"
    print(json.dumps(shown, indent=1))


if __name__ == "__main__":
    main()