import gc
import hashlib
import json
import random
import re
import sys
import time
from pathlib import Path
import mlx.core as mx
from PIL import Image
from mlx_vlm import generate, load
from mlx_vlm.prompt_utils import apply_chat_template
from mlx_vlm.utils import prepare_inputs

model_path, cap_text = sys.argv[1:]
cap = int(cap_text)
assert cap in (224, 448, 896)
start = time.perf_counter()
model, processor = load(model_path)
mx.eval(model.parameters())
mx.synchronize()
load_s = time.perf_counter() - start
question = ('Read the code inside the outlined rectangle. '
            'Return only the six uppercase letters/digits, without explanation.')
prompt = apply_chat_template(processor, model.config, question, num_images=1)
cases = json.loads(Path('cases.json').read_text())
images = {c['id']: Image.open(c['path']).convert('RGB') for c in cases}

def measure(case, phase, repeat):
    gc.collect()
    mx.synchronize()
    mx.reset_peak_memory()
    start = time.perf_counter()
    inputs = prepare_inputs(
        processor, images=[images[case['id']]], prompts=prompt,
        min_pixels=56**2, max_pixels=cap**2,
    )
    mx.eval(inputs)
    mx.synchronize()
    prepare_s = time.perf_counter() - start
    grid = inputs['image_grid_thw'].tolist()[0]
    merge = processor.image_processor.merge_size
    patch = processor.image_processor.patch_size
    tokens = int(mx.sum(inputs['input_ids'] == processor.image_token_id).item())
    assert tokens == grid[0] * grid[1] * grid[2] // merge**2
    kwargs = {k: v for k, v in inputs.items()
              if k not in ('input_ids', 'pixel_values', 'attention_mask')}
    start_gen = time.perf_counter()
    result = generate(
        model, processor, prompt,
        input_ids=inputs['input_ids'], pixel_values=inputs['pixel_values'],
        mask=inputs.get('attention_mask'), **kwargs,
        max_tokens=32, temperature=0.0, verbose=False,
    )
    mx.synchronize()
    generation_s = time.perf_counter() - start_gen
    answer = result.text.strip()
    codes = re.findall(r'(?<![A-Z0-9])[A-Z0-9]{6}(?![A-Z0-9])', answer)
    codes = [c for c in codes if any(ch.isdigit() for ch in c)]
    row = dict(cap=cap, id=case['id'], phase=phase, repeat=repeat,
               source_sha256=case['sha256'], grid=grid,
               processed_hw=[grid[1]*patch, grid[2]*patch],
               vision_tokens=tokens, prompt_tokens=result.prompt_tokens,
               generation_tokens=result.generation_tokens,
               prepare_s=prepare_s, generation_s=generation_s,
               e2e_s=time.perf_counter()-start, load_s=load_s,
               peak_mlx_gib=mx.get_peak_memory()/2**30,
               answer=answer, expected=case['answer'],
               correct=codes == [case['answer']],
               strict_correct=answer == case['answer'],
               finish_reason=result.finish_reason,
               prompt_sha256=hashlib.sha256(prompt.encode()).hexdigest())
    print(json.dumps(row), flush=True)

measure(cases[0], 'first-generation', -1)
for repeat in range(3):
    order = cases.copy()
    random.Random(100 + repeat).shuffle(order)
    for case in order:
        measure(case, 'warm', repeat)