"""Prepare pinned image overlays offline. Writes only to an explicit empty output
directory; never deploys or edits the image snapshot. Run before any coordinator
window. The generated standard verifier accepts actual FP32 q, not log(q)/T.
"""
from __future__ import annotations

import argparse
import ast
import hashlib
import json
from pathlib import Path

UTILS = 'vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py'
RUNNER = 'vllm/v1/worker/gpu/model_runner.py'
PINS = {UTILS: 'ef1d40bf31140eb644a8f3b7a003373209d6a8e297f059e70ea66cf4972d8d90',
        RUNNER: 'af9be025ff8a17a5aabd452fc3e4b987abe64fb28d852c569e1372d6f1e4d3a6'}
ABI = 'pp3-dyadic24-standard-v1'


def replace_once(text: str, old: str, new: str) -> str:
    if text.count(old) != 1:
        raise ValueError(f'probability overlay source mismatch: {old[:80]!r}')
    return text.replace(old, new)


def utilities(text: str) -> str:
    text = 'import os\nREMOTE_DRAFT_PROBS_ABI = ' + repr(ABI) + '\n' + text
    old = '    if HAS_DRAFT_LOGITS:\n        # draft_logits is stored pre-temperature, so apply scale first.'
    new = '''    if HAS_DRAFT_LOGITS == 2:
        q = tl.load(draft_logits_ptr + req_state_idx * draft_logits_stride_0
                    + draft_step * draft_logits_stride_1 + token,
                    mask=mask, other=0.).to(tl.float32)
        draft_log_prob = tl.log(q)
        draft_lse = 0.
    elif HAS_DRAFT_LOGITS:
        # draft_logits is stored pre-temperature, so apply scale first.'''
    text = replace_once(text, old, new)
    text = replace_once(text, '        if HAS_DRAFT_LOGITS:\n            # Get local draft max',
                        '        if HAS_DRAFT_LOGITS == 1:\n            # Get local draft max')
    text = replace_once(text, '        draft_log_probs = draft_logits - draft_lse', '''        draft_log_probs = draft_logits - draft_lse
        if HAS_DRAFT_LOGITS == 2:
            q = tl.load(draft_logits_ptr + req_state_idx * draft_logits_stride_0
                        + resample_idx * draft_logits_stride_1 + block,
                        mask=mask, other=0.).to(tl.float32)
            draft_log_probs = tl.log(q)''')
    text = replace_once(text, '    # Compute the per-vocab-block logits stats,', '''    if os.environ.get("VLLM_DSPARK_SAMPLED_TOP32", "0") == "1":
        if use_block_verification or synthetic_conditional_rates is not None or watermark:
            raise ValueError("PP3 q ABI supports standard verification only, without watermarking")
        if draft_logits is None or draft_logits.dtype != torch.float32 or draft_logits.shape[-1] != target_logits.shape[-1]:
            raise ValueError("PP3 standard verification needs full-vocab FP32 probabilities")
        has_draft_logits = 2

    # Compute the per-vocab-block logits stats,''')
    return text


def runner(text: str) -> str:
    text = 'REMOTE_DRAFT_PROBS_ABI = ' + repr(ABI) + '\n' + text
    return replace_once(text, '        sampler_output: SamplerOutput | None\n', '''        # PP3 failed/missing-q rounds are target-only. Mask verification, never
        # reinterpret sampled IDs as point-mass proposals. Tensor values stay on GPU.
        pp3_valid = getattr(self.speculator, "sampled_valid", None)
        if pp3_valid is not None and input_batch.num_draft_tokens > 0:
            bad = ((input_batch.expanded_local_pos > 0)
                   & ~pp3_valid[input_batch.expanded_idx_mapping.long()])
            invalid_drafts = bad if invalid_drafts is None else invalid_drafts | bad

        sampler_output: SamplerOutput | None
''')


def prepare(source: Path, output: Path) -> dict[str, object]:
    if output.exists():
        raise ValueError('output must not exist; refuse overwriting any prepared overlay')
    texts = {}
    for name, pin in PINS.items():
        data = (source / name).read_bytes()
        if hashlib.sha256(data).hexdigest() != pin:
            raise ValueError(f'image snapshot differs from pinned PP3 source: {name}')
        texts[name] = utilities(data.decode()) if name == UTILS else runner(data.decode())
        ast.parse(texts[name])
        compile(texts[name], name, 'exec')
    result: dict[str, object] = {'abi': ABI, 'source_sha256': PINS,
                               'output_sha256': {name: hashlib.sha256(text.encode()).hexdigest()
                                                 for name, text in texts.items()}}
    output.mkdir(parents=True)
    for name, text in texts.items():
        path = output / name
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(text)
    (output / 'receipt.json').write_text(json.dumps(result, indent=2)+'\n')
    return result


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--source', type=Path, required=True)
    p.add_argument('--output', type=Path, required=True)
    a = p.parse_args()
    print(json.dumps(prepare(a.source, a.output), indent=2))


if __name__ == '__main__':
    main()