"""Fresh-process embedding-row experiment; never initialize Torch/MPS."""

import argparse
import hashlib
import json
import time
from pathlib import Path

import mlx.core as mx
import numpy as np

from h3_mlx.memory import FootprintMonitor
from h3_mlx.runtime import configure_memory, environment
from h3_mlx.weights import WeightStore


def main(a):
    configure_memory()
    store = WeightStore(a.encoder)
    key = "model.language_model.embed_tokens.weight"
    # Same number of padded/effective token rows as product, with duplicates,
    # padding, and vocabulary endpoints. No tokenizer/download dependency.
    ids = ([0, 1, 3, store.headers[key]["shape"][0] - 1, 3, 0, 100, 999] * 11)[:88]
    if max(ids) >= store.headers[key]["shape"][0]:
        raise ValueError("Unexpected vocabulary")
    tick = time.perf_counter()
    with FootprintMonitor(interval=0.001) as physical:
        if a.mode == "full":
            embedding = store.get(key)
            value = embedding[mx.array(ids, mx.int32)]
            mx.eval(value)
            del embedding
        else:
            value = store.get(key, rows=ids)
        mx.eval(value)
        elapsed = time.perf_counter() - tick
        bits = np.asarray(value.view(mx.uint16))
    row = {
        "mode": a.mode,
        "seconds": elapsed,
        "physical": physical.result(),
        "mlx_peak_bytes": mx.get_peak_memory(),
        "output_sha256": hashlib.sha256(bits.tobytes()).hexdigest(),
        "shape": list(value.shape),
        "environment": environment(),
        "scope": "Cold process, page-cache warm tensor read + gather only. Excludes tokenizer and full 50-layer encoder.",
    }
    dest = Path(a.output)
    dest.write_text(json.dumps(row, indent=2) + "\n")
    print(json.dumps(row), flush=True)


if __name__ == "__main__":
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--encoder", required=True)
    p.add_argument("--mode", choices=["full", "rows"], required=True)
    p.add_argument("--output", required=True)
    main(p.parse_args())