import argparse
import hashlib
import json
from pathlib import Path
from transformers import AutoTokenizer

p = argparse.ArgumentParser()
p.add_argument('--model', required=True)
p.add_argument('--input', type=Path, required=True)
p.add_argument('--output', type=Path, required=True)
a = p.parse_args()
tok = AutoTokenizer.from_pretrained(a.model, local_files_only=True)
template = tok.chat_template
start = '{%- elif message.role == "assistant" %}'
stop = '{%- elif message.role == "tool" %}'
if template.count(start) != 1 or template.count(stop) != 1:
    raise ValueError('Pinned template assistant boundaries differ; cannot infer masks')
tracked = template.replace(start, start + '\n{%- generation %}').replace(stop, '{%- endgeneration %}\n' + stop)
header = tok.encode('<|im_start|>assistant\n', add_special_tokens=False)
eos = tok.convert_tokens_to_ids('<|im_end|>')
records = {}
for row in map(json.loads, a.input.read_text().splitlines()):
    encoded = tok.apply_chat_template(row['messages'], tools=row.get('tools'), chat_template=tracked, tokenize=True, return_dict=True, return_assistant_tokens_mask=True, add_generation_prompt=False, enable_thinking=False)
    ids, mask = encoded['input_ids'], encoded['assistant_masks']
    if ids != row['token_ids']:
        raise ValueError(f'Mask instrumentation changed actual tokens: {row["id"]}')
    index = 0
    assistant_messages = 0
    while index < len(mask):
        if not mask[index]:
            index += 1
            continue
        begin = index
        while index < len(mask) and mask[index]:
            index += 1
        end = index
        if ids[begin:begin+len(header)] != header:
            raise ValueError(f'Unexpected assistant header tokenization: {row["id"]}')
        assistant_messages += 1
        mask[begin:begin+len(header)] = [0] * len(header)
        eos_positions = [j for j in range(begin+len(header), end) if ids[j] == eos]
        if len(eos_positions) != 1:
            raise ValueError(f'Ambiguous assistant termination: {row["id"]}')
        mask[eos_positions[0]+1:end] = [0] * (end-eos_positions[0]-1)
    if assistant_messages != sum(m['role'] == 'assistant' for m in row['messages']):
        raise ValueError(f'Assistant span count mismatch: {row["id"]}')
    if not any(mask[1:]):
        raise ValueError(f'No assistant targets: {row["id"]}')
    records[row['id']] = {'token_ids': ids, 'target_mask': mask[1:], 'domain': row['domain'], 'split': row['split']}
result = {'method': 'Instrument pinned Jinja assistant branch with generation tags; verify exact original token equality; exclude supplied role headers and post-EOS newline; include assistant content/tool-call syntax and EOS', 'input_sha256': hashlib.sha256(a.input.read_bytes()).hexdigest(), 'template_sha256': hashlib.sha256(template.encode()).hexdigest(), 'records': records}
a.output.write_text(json.dumps(result, ensure_ascii=False) + '\n')
print(json.dumps({'records': len(records), 'assistant_targets': sum(sum(r['target_mask']) for r in records.values()), 'all_targets': sum(len(r['target_mask']) for r in records.values())}))