#!/usr/bin/env python3
"""Offline native MLX text completion for OpenSML-Base-175M."""
from pathlib import Path
import argparse
import hashlib
import json
import math
import sys
import tempfile
import time
import zipfile
ROOT = Path(__file__).resolve().parent
BUNDLE_SHA256 = "14f85f0ff677d35a6fb34f118f2dd21a3fdd504fff27c163860ab919b21e0e98"

def sha256(path):
    with Path(path).open("rb") as file:
        return hashlib.file_digest(file, "sha256").hexdigest()

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    inputs = parser.add_mutually_exclusive_group(required=True)
    inputs.add_argument("--prompt", help="Raw text prefix; no chat template is applied")
    inputs.add_argument("--prompt-file", type=Path, help="UTF-8 text prefix file")
    parser.add_argument("--weights", type=Path, default=ROOT/"model.safetensors")
    parser.add_argument("--max-tokens", type=int, default=128)
    parser.add_argument("--temperature", type=float, default=0)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--json", action="store_true", help="Print completion and generation metadata as JSON")
    args = parser.parse_args()
    if not 1 <= args.max_tokens <= 2047:
        parser.error("--max-tokens must be between 1 and 2047")
    if not math.isfinite(args.temperature) or args.temperature < 0:
        parser.error("--temperature must be finite and nonnegative")
    prompt = args.prompt if args.prompt is not None else args.prompt_file.read_text(encoding="utf-8")
    config = json.loads((ROOT/"config.json").read_text())
    if not args.weights.is_file():
        parser.error("Missing model.safetensors; download the complete repository")
    if sha256(args.weights) != config["weight_sha256"]:
        parser.error("Weights do not match the published base model identity")
    archive = ROOT/"native_bundle.zip"
    if sha256(archive) != BUNDLE_SHA256:
        parser.error("Native source bundle failed its integrity check")
    with tempfile.TemporaryDirectory(prefix="opensml-base-runtime-") as directory:
        runtime = Path(directory)
        with zipfile.ZipFile(archive) as bundle:
            for member in bundle.infolist():
                p = Path(member.filename)
                if p.is_absolute() or ".." in p.parts or (member.external_attr >> 16) & 0o170000 == 0o120000:
                    parser.error("Unsafe source-bundle entry")
            bundle.extractall(runtime)
        sys.path.insert(0, str(runtime/"pretraining"))
        from sml_v1.tokenization import Tokenizer
        from sml_v1.model import TransformerConfig, TransformerLM, count_parameters
        import mlx.core as mx
        from mlx.utils import tree_map
        tokenizer = Tokenizer(runtime/"tokenizer")
        ids = tokenizer.encode(prompt)
        if not ids:
            parser.error("Prompt must contain at least one token")
        if len(ids)+args.max_tokens > config["context_limit"]:
            parser.error("Prompt plus generation budget exceeds 2048 tokens; no truncation is performed")
        mx.set_default_device(mx.gpu)
        mx.set_cache_limit(4*1024**3)
        mx.random.seed(args.seed)
        model = TransformerLM(TransformerConfig(**config["model"]))
        model.update(tree_map(lambda value:value.astype(mx.bfloat16),model.parameters()))
        model.load_weights(str(args.weights))
        mx.eval(model.parameters())
        model.eval()
        if count_parameters(model) != config["parameters"]:
            raise ValueError("Unexpected parameter count")
        inputs = mx.array([ids],dtype=mx.int32)
        cache = None
        generated = []
        stop = "token_limit"
        start = time.perf_counter()
        try:
            for _ in range(args.max_tokens):
                logits, cache = model.logits(inputs, cache)
                final = logits[0,-1]
                if args.temperature == 0:
                    token = int(mx.argmax(final).item())
                else:
                    token = int(mx.random.categorical(final.astype(mx.float32)/args.temperature).item())
                if token == tokenizer.eos:
                    stop = "document_end"; break
                if token in (0,2,3):
                    stop = "structural_token"; break
                generated.append(token)
                inputs = mx.array([[token]],dtype=mx.int32)
            completion = tokenizer.decode(generated)
            record = {"model":config["model_name"],"prompt":prompt,"completion":completion,
                      "prompt_tokens":len(ids),"output_tokens":len(generated),
                      "stop_reason":stop,"seconds":time.perf_counter()-start,
                      "protocol":{"serialization":"raw text prefix","temperature":args.temperature,
                                  "seed":args.seed,"max_tokens":args.max_tokens,"context_limit":2048,
                                  "truncation":False,"compute_dtype":"bfloat16"}}
            print(json.dumps(record,indent=2,ensure_ascii=False) if args.json else completion)
        finally:
            mx.clear_cache()

if __name__ == "__main__":
    main()