"""Pinned MTPLX oracle for issue 107's real long two-image product request."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import subprocess
import sys

checkout, root, image, output = map(Path, sys.argv[1:5])
assert subprocess.check_output(["git", "-C", str(checkout), "rev-parse", "HEAD"], text=True).strip() == "9882703f3105363ddc37eca9f97aa09a1d387112"
sys.path.insert(0, str(checkout))
from mtplx.server.openai import _server_runtime_env_overrides
os.environ.update(_server_runtime_env_overrides(argparse.Namespace(model=str(root), generation_mode="mtp"), None))
os.environ["MTPLX_SUSTAINED_PREFILL"] = "1"
os.environ["MTPLX_PREFILL_CHUNK_SIZE"] = "2048"
import mlx.core as mx
import numpy as np
from mlx_lm.utils import load_model
from mtplx.models import qwen4_exp as qwen
from mtplx.runtime import _load_tokenizer_resilient, MTPLXRuntime
from mtplx.mtp_patch import MTPContract
from mtplx.sampling import SamplerConfig
from mtplx.generation import generate_mtpk
from mtplx.server.openai import ChatMessage, _encode_messages_uncached, _materialize_vision_splice

output.mkdir(parents=True, exist_ok=True)
print(json.dumps(dict(case="environment", overrides={k:v for k,v in os.environ.items() if k.startswith("MTPLX_")})), flush=True)
tokenizer = _load_tokenizer_resilient(root, json.loads((root/"config.json").read_text()))
model, _ = load_model(root, lazy=False, strict=True, get_model_classes=lambda **_: (qwen.Model,qwen.ModelArgs))
model.post_weight_load(root)
assert model.attach_mtp(root)
runtime = MTPLXRuntime(model, tokenizer, root, True, MTPContract())
text = "Describe these images" + "<|vision_start|><|image_pad|><|vision_end|>" * 2
ids = _encode_messages_uncached(tokenizer, [ChatMessage(role="user", content=text)], enable_thinking=True,
    reasoning_effort="low", preserve_reasoning_history=True, tools=None)
ids, splice = _materialize_vision_splice(argparse.Namespace(args=argparse.Namespace(model=str(root))), [image.read_bytes()]*2, ids)
assert len(ids) > 2048
(output/"input.json").write_text(json.dumps(ids))
out = generate_mtpk(runtime, ids, max_tokens=1, sampler=SamplerConfig(temperature=0., top_p=.95, top_k=20),
    speculative_depth=3, seed=1, stop_token_ids=set(tokenizer.eos_token_ids), mtp_history_policy="committed",
    verify_strategy="batched", vision_splice=splice, capture_final_state=True,
    prefill_callback=lambda event: print(json.dumps(dict(case="prefill", **event)), flush=True))
assert out.final_state.safe_to_commit
logits = np.asarray(out.final_state.final_logits.astype(mx.float32)).reshape(-1)
(output/"final.f32").write_bytes(logits.tobytes())
print(json.dumps(dict(case="long_vision", prompt_tokens=len(ids), tokens=out.tokens, text=out.text,
    final_logits_sha256=hashlib.sha256(logits.tobytes()).hexdigest())), flush=True)