"""Generate once from a model split across the cluster, and record what it cost.

Purpose: the shortest path from a configured cluster to a number worth writing down.
    It loads one model across the distributed group, generates a completion, and
    records prefill speed, decode speed, token counts and peak memory per rank in the
    lab notebook. The --pipeline flag selects pipelining instead of tensor
    parallelism, which is the one decision this script makes differently from run to
    run and the one the lab asks you to compare.
Platform: mac (Track M). Runs on one machine too, as a group of one or as two ranks
    launched with `mlx.launch -n 2`, which is the single-machine path.
Minimum memory: 32 GB per Mac for a model larger than one machine; 24 GB for the
    single-Mac baseline with a smaller model.
Assumes: mlx-lm installed in the python mlx.launch starts, this file present at the
    SAME absolute path on every machine, and the model repository readable by every
    rank. With --pipeline each rank downloads only its own shard; without it every
    rank downloads the whole repository, so check the disk on both machines first.

Usage:
    mlx.launch --backend jaccl --hostfile hosts.json -- \
        /path/to/python sharded-generate.py --model mlx-community/Qwen3-32B-8bit \
        --label two-macs-jaccl-tensor

    mlx.launch --backend ring --hostfile hosts.json -- \
        /path/to/python sharded-generate.py --model mlx-community/Qwen3-32B-8bit \
        --pipeline --label two-macs-ring-pipeline

    python3 sharded-generate.py --model mlx-community/Qwen3-8B-4bit --label one-mac
"""

from __future__ import annotations

import argparse
import json
import platform
from datetime import date, datetime, timezone
from pathlib import Path

import mlx.core as mx
from mlx_lm import stream_generate
from mlx_lm.utils import sharded_load

HERE = Path(__file__).resolve().parent

# Long enough that prefill is measurable rather than noise, and boring enough that
# the answer is not the point. Repeat it to reach the prompt length you want.
DEFAULT_PROMPT = (
    "Explain, for an engineer who has just built a two-machine cluster, why the "
    "prompt-processing phase and the token-generation phase of a language model "
    "have different bottlenecks, what each phase asks of the link between the two "
    "machines, and which of the two a faster cable actually helps. Answer in no "
    "more than six sentences."
)


def parse_args():
    p = argparse.ArgumentParser(description="Distributed generation with mlx-lm")
    p.add_argument("--model", required=True, help="MLX repository or local path.")
    p.add_argument("--prompt", default=DEFAULT_PROMPT, help="The prompt text.")
    p.add_argument(
        "--prompt-file",
        default=None,
        help="Read the prompt from this file instead of --prompt.",
    )
    p.add_argument(
        "--prompt-repeat",
        type=int,
        default=1,
        help="Repeat the prompt this many times, to lengthen prefill.",
    )
    p.add_argument("--max-tokens", type=int, default=256)
    p.add_argument(
        "--pipeline",
        action="store_true",
        help="Use pipelining instead of tensor parallelism.",
    )
    p.add_argument("--label", default="", help="A name for this run, for the record.")
    p.add_argument("--labbook", default=str(HERE / "labbook.md"))
    p.add_argument(
        "--quiet",
        action="store_true",
        help="Do not stream the answer; report only the numbers.",
    )
    return p.parse_args()


def build_prompt(args) -> str:
    if args.prompt_file:
        text = Path(args.prompt_file).read_text(encoding="utf-8")
    else:
        text = args.prompt
    return ("\n\n".join([text] * max(1, args.prompt_repeat))).strip()


def main() -> int:
    args = parse_args()

    group = mx.distributed.init()
    rank = group.rank()
    size = group.size()

    def rprint(*a, **kw):
        if rank == 0:
            print(*a, **kw)

    pipeline_group = group if args.pipeline else None
    tensor_group = None if args.pipeline else group

    rprint(f"loading {args.model} across {size} rank(s)")
    rprint("split: " + ("pipeline" if args.pipeline else "tensor parallel"))
    model, tokenizer = sharded_load(args.model, pipeline_group, tensor_group)

    messages = [{"role": "user", "content": build_prompt(args)}]
    prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)

    response = None
    for response in stream_generate(
        model, tokenizer, prompt, max_tokens=args.max_tokens
    ):
        if not args.quiet:
            rprint(response.text, end="", flush=True)
    if not args.quiet:
        rprint()

    if response is None:
        rprint("No tokens were generated. Check --max-tokens and the prompt.")
        return 1

    # Every rank knows its own peak memory; only rank 0 writes the notebook line,
    # so the peaks are gathered rather than reported separately per machine.
    peaks = mx.distributed.all_gather(mx.array([response.peak_memory], mx.float32))
    mx.eval(peaks)
    peak_by_rank = [round(float(v), 3) for v in peaks.tolist()]

    if rank != 0:
        return 0

    record = {
        "lab": "part-21/two-mac-cluster-over-thunderbolt-5",
        "record": "sharded-generate",
        "date": date.today().isoformat(),
        "recorded_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat(),
        "label": args.label,
        "model": args.model,
        "group_size": size,
        "split": "pipeline" if args.pipeline else "tensor",
        "machine": platform.node().split(".")[0],
        "prompt_tokens": response.prompt_tokens,
        "prompt_tokens_per_second": round(response.prompt_tps, 3),
        "generation_tokens": response.generation_tokens,
        "generation_tokens_per_second": round(response.generation_tps, 3),
        "peak_memory_gb_by_rank": peak_by_rank,
    }
    text = json.dumps(record, sort_keys=True)
    with open(args.labbook, "a", encoding="utf-8") as fh:
        fh.write(text + "\n")
    print(text)
    print(f"Appended to {args.labbook}")
    return 0


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