"""One bounded vector request over stdin; stdout contains only JSON."""
import argparse
import json
import math
import sys
import torch


def vector(value):
    if not isinstance(value, list) or not 1 <= len(value) <= 4096:
        raise ValueError("vector length must be 1..4096")
    if any(type(x) not in (int, float) or not math.isfinite(x) or abs(x) > 1e20 for x in value):
        raise ValueError("elements must be finite numbers in [-1e20, 1e20]")
    return value


def add(request, device):
    if not isinstance(request, dict) or set(request) != {"a", "b"}:
        raise ValueError("request must contain exactly a and b")
    a, b = vector(request["a"]), vector(request["b"])
    if len(a) != len(b):
        raise ValueError("vectors must have equal lengths")
    if device not in ("cpu", "cuda:0"):
        raise ValueError("device must be cpu or cuda:0")
    if device == "cuda:0" and not torch.cuda.is_available():
        raise ValueError("CUDA requested but unavailable; no silent CPU fallback")
    with torch.inference_mode():
        left = torch.tensor(a, dtype=torch.float32, device=device)
        right = torch.tensor(b, dtype=torch.float32, device=device)
        result = (left + right).cpu().tolist()
    return {"device": device, "dtype": "float32", "values": result}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", choices=["cpu", "cuda:0"], default="cpu")
    args = parser.parse_args()
    try:
        raw = sys.stdin.buffer.read(262145)
        if len(raw) > 262144:
            raise ValueError("request exceeds 256 KiB")
        result = add(json.loads(raw), args.device)
        print(json.dumps(result, allow_nan=False))
    except (ValueError, TypeError, OverflowError, RuntimeError) as error:
        print(str(error), file=sys.stderr)
        return 1
    return 0


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