#!/usr/bin/env python3
"""Serve a verified text-only TT-native checkpoint without original HF weights."""
import argparse
import hashlib
import json
import os
import shlex
from pathlib import Path

p = argparse.ArgumentParser(description=__doc__)
p.add_argument('--checkpoint', type=Path, required=True)
p.add_argument('--cache-root', type=Path, required=True)
p.add_argument('--image', default='qwen35-tt-native-quant:mtp')
p.add_argument('--port', type=int, default=8001)
p.add_argument('--name', default='qwen9b-native-api')
p.add_argument('--max-model-len', type=int, default=8192)
p.add_argument('--mtp', choices=('on', 'off'), default='on',
               help='Enable MTP-1 (default), or use the verified non-speculative baseline')
p.add_argument('--device-ownership-confirmed', action='store_true')
p.add_argument('--print-command', action='store_true')
a = p.parse_args()
checkpoint, cache = a.checkpoint.resolve(), a.cache_root.resolve()
manifest_path = checkpoint / 'native_manifest.json'
manifest = json.loads(manifest_path.read_text())
mtp_enabled = a.mtp == 'on'
mtp_environment = '1' if mtp_enabled else '0'
proof_path = checkpoint / ('equivalence-mtp.json' if mtp_enabled else 'equivalence.json')
proof = json.loads(proof_path.read_text())
if (proof.get('exact_logits_equal') is not True
        or proof.get('tokens_compared', 0) < 1
        or proof.get('manifest_sha256') != hashlib.sha256(manifest_path.read_bytes()).hexdigest()
        or proof.get('precision') != manifest['precision']):
    p.error('Checkpoint lacks manifest-bound native reload equivalence evidence')
if proof.get('runtime_environment', {}).get('QWEN36_MTP') != mtp_environment:
    p.error('Checkpoint equivalence evidence does not match the selected MTP mode')
if manifest['scope'] != 'text-only-no-vision-with-mtp':
    p.error('Unsupported checkpoint scope')
mtp = manifest.get('mtp', {})
if (mtp.get('enabled') is not True or mtp.get('num_speculative_tokens') != 1
        or mtp.get('source_tensor_count') != 15 or mtp.get('storage') != 'lossless'):
    p.error('Checkpoint must preserve all 15 original MTP tensors losslessly for MTP-1')
if not 128 <= a.max_model_len <= 8192 or not 1 <= a.port <= 65535:
    p.error('Invalid context length or host port')
if cache == checkpoint or cache.is_relative_to(checkpoint) or checkpoint.is_relative_to(cache):
    p.error('Use independent checkpoint and writable cache directories')
command = ['docker', 'run', '--rm', '--name', a.name,
           '--publish', f'127.0.0.1:{a.port}:8000', '--device', '/dev/tenstorrent:/dev/tenstorrent',
           '-v', '/dev/hugepages:/dev/hugepages', '-v', '/dev/hugepages-1G:/dev/hugepages-1G',
           '-v', f'{checkpoint}:/model:ro', '-v', f'{cache}:/cache',
           '--workdir', '/home/container_app_user/tt-metal',
           '--entrypoint', '/home/container_app_user/tt-metal/python_env/bin/python']
environment = dict(proof['runtime_environment'])
environment['QWEN36_MTP'] = mtp_environment
cache_identity = hashlib.sha256(json.dumps(
    {'manifest': proof['manifest_sha256'], 'runtime_sources': proof['runtime_sources']},
    sort_keys=True).encode()).hexdigest()
environment.update({'HF_MODEL': '/model', 'MODEL_WEIGHTS_DIR': '/model', 'TT_CACHE_PATH': f'/cache/{cache_identity}',
                    'HF_HUB_OFFLINE': '1', 'TRANSFORMERS_OFFLINE': '1',
                    'VLLM_TARGET_DEVICE': 'tt', 'TORCHDYNAMO_DISABLE': '1',
                    'QWEN36_EVAL_PRECISION': json.dumps(manifest['precision'], separators=(',', ':'))})
for key, value in sorted(environment.items()):
    command.extend(['-e', f'{key}={value}'])
command.extend([a.image, '-m', 'vllm.entrypoints.openai.api_server', '--model', '/model',
                '--served-model-name', 'Qwen/Qwen3.5-9B', '--host', '0.0.0.0', '--port', '8000',
                '--max-model-len', str(a.max_model_len), '--max-num-seqs', '1', '--block-size', '64',
                '--max-num-batched-tokens', '16384', '--seed', '9472',
                '--enable-auto-tool-choice', '--tool-call-parser', 'qwen3_coder', '--reasoning-parser', 'qwen3',
                '--limit-mm-per-prompt', '{"image":0,"video":0}',
                '--additional-config', '{"tt":{"trace_region_size":1073741824,"l1_small_size":24576,"sample_on_device_mode":"decode_only"}}'])
if mtp_enabled:
    command.extend(['--speculative-config', '{"method":"mtp","num_speculative_tokens":1}'])
print(shlex.join(command), flush=True)
if not a.print_command:
    if not a.device_ownership_confirmed:
        p.error('Stop other P150 workloads and pass --device-ownership-confirmed; this launcher never stops them')
    cache.mkdir(parents=True, exist_ok=True)
    os.execvp(command[0], command)