#!/usr/bin/env python3
"""A thegrid `odin` operator backed by gridlm instead of a general LLM.

Same queue protocol as thegrid's tools/operator/grid_operator.py: the colony writes
<queue>/request.json (with "genome", a list of instruction names, and "request_id"); this
writes <queue>/proposal.json ({"request_id", "authored_at", "genome"}) atomically. The colony
validates every proposal with its own parser, so a proposal it can't use costs one birth (it
falls back to blind variation) and never the colony.

How a proposal is made: the generator samples K children for the parent; the screen scores
each; the best-scoring child that differs from the parent is proposed. Runs on CPU in
milliseconds, so there is no GPU queue to wait for.

Honest scope: trained on ONE colony (kannaka-grid's colony-one, an older thegrid ISA and a
starvation economy). Instruction names that colony never used are unknown to the model, and
anything it proposes that a destination colony can't run is rejected there by design.

usage:
  python gridlm_operator.py --model-dir . --queue /path/to/colony/operator [--queue ...]
  python gridlm_operator.py --model-dir . --once-genome "harvest harvest alloc copy ifnotdone jmpb fork scan move"
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from pathlib import Path

import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import train as g  # noqa: E402  (model classes, tokens, sampling)


def load(model_dir: Path):
    vocab = json.loads((model_dir / "vocab.json").read_text())
    screen = g.Screen(len(vocab))
    gen = g.Generator(len(vocab))
    for model, name in ((screen, "screen"), (gen, "generator")):
        st = model_dir / f"{name}.safetensors"
        if st.exists():                      # preferred: no pickle
            from safetensors.torch import load_file
            model.load_state_dict(load_file(str(st)))
        else:
            model.load_state_dict(torch.load(model_dir / f"{name}.pt", map_location="cpu", weights_only=True))
    screen.eval()
    gen.eval()
    return vocab, {t: i for i, t in enumerate(vocab)}, screen, gen


def propose(names, vocab, idx, screen, gen, k=8, temperature=0.9, seed=None):
    parent = [idx.get(n, g.UNK) for n in names][: g.MAXG]
    kids = g.sample_children(gen, [parent] * k, vocab, k, temperature=temperature,
                             rng_seed=seed if seed is not None else int(time.time() * 1000) % (2 ** 31))
    cands = [c for c in kids if 1 <= len(c) <= g.MAXG and c != list(names)]
    if not cands:
        return None, None
    rows = [{"parent": list(names)[: g.MAXG], "child": c, "label": 0, "child_id": ""} for c in cands]
    X, _ = g.screen_tensor(rows, idx)
    p = g.screen_scores(screen, X)
    best = int(p.argmax())
    return cands[best], float(p[best])


def serve(queue: Path, model, stats) -> bool:
    req_p, prop_p = queue / "request.json", queue / "proposal.json"
    if prop_p.exists() or not req_p.exists():
        return False
    try:
        req = json.loads(req_p.read_text())
    except (OSError, ValueError):
        return False
    child, score = propose(req.get("genome", []), *model)
    if child is None:
        stats["rejected"] += 1
        return True
    tmp = prop_p.with_suffix(".tmp")
    tmp.write_text(json.dumps({"request_id": req.get("request_id"), "authored_at": time.time(),
                               "genome": child, "gridlm_screen_p": score}))
    tmp.replace(prop_p)
    stats["served"] += 1
    print(f"{time.strftime('%H:%M:%S')} {queue}: {len(req.get('genome', []))} -> {len(child)} "
          f"(screen p={score:.3f}; served {stats['served']} rejected {stats['rejected']})", flush=True)
    return True


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model-dir", default=".")
    ap.add_argument("--queue", action="append", default=[])
    ap.add_argument("--once-genome", help="print one proposal for this genome (space-separated names) and exit")
    ap.add_argument("--poll", type=float, default=1.0)
    a = ap.parse_args()
    torch.set_num_threads(2)
    model = load(Path(a.model_dir))
    if a.once_genome:
        child, score = propose(a.once_genome.split(), *model)
        print(json.dumps({"parent": a.once_genome.split(), "child": child, "screen_p": score}))
        return
    stats = {"served": 0, "rejected": 0}
    queues = [Path(q) for q in a.queue]
    print("gridlm operator; queues:", ", ".join(map(str, queues)), flush=True)
    while True:
        worked = False
        for q in queues:
            if q.exists():
                worked = serve(q, model, stats) or worked
        if not worked:
            time.sleep(a.poll)


if __name__ == "__main__":
    main()