"""Read exo's placement previews from standard input and make them legible.

Purpose: `GET /instance/previews` returns every way exo is willing to place a model
    across the cluster, and each entry carries four separate decisions: which nodes,
    pipeline or tensor sharding, which transport, and what it costs each node in
    memory. This turns that list into something you can read at a glance, and can
    also pick one entry out of it so a script does not have to parse JSON in shell.
Platform: all (spark, strix, mac, nvidia). Pure standard library; it only reads text.
Minimum memory: none of consequence; this is a filter, not a workload.
Assumes: the JSON body of GET /instance/previews on standard input. Placements whose
    `error` field is not null are shown but never selectable, because exo has already
    said why they will not work.

Usage:
    curl -fsS --get "$BASE/instance/previews" --data-urlencode "model_id=$MODEL" \
        | python3 exo-placement.py --list

    curl -fsS --get "$BASE/instance/previews" --data-urlencode "model_id=$MODEL" \
        | python3 exo-placement.py --select 0 > placement.json
"""

from __future__ import annotations

import argparse
import json
import sys


def parse_args():
    p = argparse.ArgumentParser(description="Read exo placement previews")
    mode = p.add_mutually_exclusive_group(required=True)
    mode.add_argument(
        "--list",
        dest="list_only",
        action="store_true",
        help="Print one line per placement, numbering the usable ones.",
    )
    mode.add_argument(
        "--select",
        type=int,
        metavar="N",
        help="Print the request body for usable placement N, counting from 0.",
    )
    return p.parse_args()


def usable_placements(payload) -> list:
    return [p for p in payload.get("previews", []) if p.get("error") is None]


def describe(index, preview) -> str:
    mark = f"[{index}]" if index is not None else "[--]"
    sharding = preview.get("sharding", "?")
    transport = preview.get("instance_meta", "?")
    lines = [f"{mark} sharding={sharding} transport={transport}"]
    for node, delta in (preview.get("memory_delta_by_node") or {}).items():
        lines.append(f"     {node}: {round(int(delta) / 1e9, 2)} GB")
    error = preview.get("error")
    if error is not None:
        lines.append(f"     not usable: {error}")
    return "\n".join(lines)


def main() -> int:
    args = parse_args()
    try:
        payload = json.load(sys.stdin)
    except ValueError as exc:
        print(f"Could not read JSON from standard input: {exc}", file=sys.stderr)
        return 1

    previews = payload.get("previews", [])
    if not previews:
        print("exo returned no placements at all for this model.", file=sys.stderr)
        print("Check the model id with: curl $BASE/models", file=sys.stderr)
        return 1

    if args.list_only:
        index = 0
        for preview in previews:
            if preview.get("error") is None:
                print(describe(index, preview))
                index += 1
            else:
                print(describe(None, preview))
        print("")
        print(f"{index} usable placement(s). Use --select with the number in [].")
        return 0

    usable = usable_placements(payload)
    if not usable:
        print("Every placement exo offered carries an error.", file=sys.stderr)
        return 1
    if args.select < 0 or args.select >= len(usable):
        print(
            f"Only {len(usable)} usable placement(s); {args.select} is out of range.",
            file=sys.stderr,
        )
        return 1
    print(json.dumps({"instance": usable[args.select]["instance"]}))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())