"""Hidden-state feature of the hidden-state arms (EXPERIMENT_LOG.md 2026-10-19, tasks 3 and 4a) and its HTTP server.
Runs in .train-venv.

Feature: Qwen3.5-9B's last-layer hidden state (after the final norm, i.e. the lm_head input; 4096 dims) at the last
prompt token of A's grounding question (ctrl_policy.ground(): the frame, GROUND.format(name), thinking off, the answer
started as START, rendered by the chat template with continue_final_message as vLLM renders it), prefill only. Batches
are right-padded and each row's last real token is gathered. Values travel as base64 of the bf16 bits (encode/decode;
hid_policy.decode is the numpy twin), so a logged feature is exactly the one the controller used.

Server: one model replica per GPU (threads), a shared queue, each replica takes up to MAX_BATCH waiting requests per
forward. POST /feat {"frame": base64 PNG, "name": target name} -> {"feat": ...} and, with --base, also "base": the
feature with the adapter disabled (the 4a KL reference input). The port opens once every replica is loaded.
Usage: CUDA_VISIBLE_DEVICES=2 .train-venv/bin/python train/hid_server.py --port 9210 [--gpus 0] [--adapter DIR --base]
"""
import argparse
import base64
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import io
import json
from pathlib import Path
import queue
import sys
import threading
import time

import torch
from PIL import Image
from transformers import AutoProcessor

from bt_train import MODULES, load_model  # noqa: F401  (MODULES: the LoRA targets, re-exported for hid_grpo)
from ctrl_policy import START

ROOT = Path(__file__).resolve().parents[1]
sys.path.append(str(ROOT / "bench"))
from diag_perception_direction import GROUND  # noqa: E402

MAX_BATCH = 32


def processor():
    return AutoProcessor.from_pretrained(ROOT / "models/Qwen3.5-9B", local_files_only=True)


def prompt(proc, name):
    return proc.apply_chat_template(
        [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": GROUND.format(name)}]},
         {"role": "assistant", "content": START}],
        tokenize=False, add_generation_prompt=False, continue_final_message=True, enable_thinking=False)


def batch_inputs(proc, images, names):
    return dict(proc(text=[prompt(proc, n) for n in names], images=images, padding=True, padding_side="right",
                     return_tensors="pt"))


def features(model, inputs, device):
    """(B, 4096) bf16 last-layer hidden states at each row's last prompt token; model: the Qwen3.5 model or its
    PeftModel (grad follows the caller's mode)."""
    body = (model.get_base_model() if hasattr(model, "get_base_model") else model).model
    x = {k: v.to(device) for k, v in inputs.items()}
    with torch.autocast("cuda", dtype=torch.bfloat16):
        h = body(**x, use_cache=False).last_hidden_state
    return h[torch.arange(len(h), device=device), x["attention_mask"].sum(1) - 1]


def encode(f):
    """bf16 vector -> base64 of its bits."""
    return base64.b64encode(f.to(torch.bfloat16).view(torch.int16).cpu().numpy().tobytes()).decode()


def decode(s):
    return torch.frombuffer(bytearray(base64.b64decode(s)), dtype=torch.int16).view(torch.bfloat16)


def png(b64):
    return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, required=True)
    ap.add_argument("--gpus", default="0", help="local cuda indices, one replica each")
    ap.add_argument("--adapter", type=Path)
    ap.add_argument("--base", action="store_true", help="also return the feature with the adapter disabled")
    args = ap.parse_args()
    proc = processor()
    jobs = queue.Queue()

    def replica(model, device):
        while True:
            items = [jobs.get()]
            while len(items) < MAX_BATCH:
                try:
                    items.append(jobs.get_nowait())
                except queue.Empty:
                    break
            try:
                inputs = batch_inputs(proc, [png(i["frame"]) for i in items], [i["name"] for i in items])
                with torch.no_grad():
                    out = {"feat": features(model, inputs, device)}
                    if args.base:
                        with model.disable_adapter():
                            out["base"] = features(model, inputs, device)
                for k, item in enumerate(items):
                    item["out"] = {key: encode(f[k]) for key, f in out.items()}
            except Exception as e:  # answer every waiting request rather than leave it hanging
                for item in items:
                    item["out"] = {"error": f"{type(e).__name__}: {e}"}
            for item in items:
                item["done"].set()

    gpus = [int(g) for g in args.gpus.split(",")]
    for g in gpus:
        model = load_model(g)
        if args.adapter:
            from peft import PeftModel
            model = PeftModel.from_pretrained(model, args.adapter)
        model.eval()
        threading.Thread(target=replica, args=(model, torch.device("cuda", g)), daemon=True).start()

    class Handler(BaseHTTPRequestHandler):
        def log_message(self, *a):
            pass

        def reply(self, code, body):
            data = json.dumps(body).encode()
            self.send_response(code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(data)))
            self.end_headers()
            self.wfile.write(data)

        def do_GET(self):
            self.reply(200, {"gpus": gpus})

        def do_POST(self):
            item = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
            item["done"] = threading.Event()
            jobs.put(item)
            item["done"].wait()
            self.reply(500 if "error" in item["out"] else 200, item["out"])

    server = ThreadingHTTPServer(("0.0.0.0", args.port), Handler)
    server.daemon_threads = True
    print(f"{time.strftime('%Y-%m-%d %H:%M:%S')} feature server on port {args.port}, gpus {gpus}, adapter "
          f"{args.adapter}, base {args.base}", flush=True)
    server.serve_forever()


if __name__ == "__main__":
    main()