"""Offline integrity, upstream parity, and real checkpoint inference checks."""
import hashlib
import json
from pathlib import Path
import struct
import inspect

import torch
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from transformers import Gemma4ForCausalLM, Gemma4TextConfig
from transformers.models.gemma4 import modeling_gemma4 as upstream_code


def main():
    path = Path(__file__).resolve().parent
    torch.set_num_threads(8)
    report = json.loads((path / 'conversion_report.json').read_text())
    with (path / 'model.safetensors').open('rb') as f:
        length = struct.unpack('<Q', f.read(8))[0]
        header = json.loads(f.read(length))
        for key, expected in report['tensor_sha256'].items():
            start, end = header[key]['data_offsets']
            f.seek(8 + length + start)
            digest = hashlib.sha256()
            left = end - start
            while left:
                chunk = f.read(min(left, 16 * 1024**2))
                assert chunk, key
                digest.update(chunk)
                left -= len(chunk)
            assert digest.hexdigest() == expected, key
    assert set(header) - {'__metadata__'} == set(report['tensor_sha256'])
    print(f'All {len(report["tensor_sha256"])} retained tensors match original bytes.', flush=True)

    config = AutoConfig.from_pretrained(path, trust_remote_code=True, local_files_only=True)
    assert not any(k.startswith('mtp_') or k in ['vision_config', 'audio_config'] for k in config.to_dict())
    small = json.loads((path / 'config.json').read_text())
    small['_name_or_path'] = str(path)
    small.update(vocab_size=128, vocab_size_per_layer_input=128, hidden_size=64,
                 hidden_size_per_layer_input=8, intermediate_size=96, num_hidden_layers=4,
                 num_attention_heads=2, num_key_value_heads=1, head_dim=16, global_head_dim=32,
                 num_kv_shared_layers=2, sliding_window=4,
                 layer_types=['sliding_attention', 'full_attention'] * 2,
                 eos_token_id=1, pad_token_id=0, dtype='float32')
    custom = AutoModelForCausalLM.from_config(type(config)(**small), trust_remote_code=True).eval()
    upstream = Gemma4ForCausalLM(Gemma4TextConfig(**small)).eval()
    upstream.load_state_dict(custom.state_dict(), strict=True)
    ids = torch.tensor([[3, 4, 5, 6, 7], [0, 0, 8, 9, 10]])
    mask = ids.ne(0).long()
    with torch.inference_mode():
        a = custom(ids, attention_mask=mask, use_cache=True)
        b = upstream(ids, attention_mask=mask, use_cache=True)
        torch.testing.assert_close(a.logits, b.logits, rtol=0, atol=0)
        extended_mask = torch.cat([mask, torch.ones(2, 1, dtype=torch.long)], dim=1)
        nxt = torch.tensor([[11], [12]])
        ac = custom(nxt, attention_mask=extended_mask, past_key_values=a.past_key_values, use_cache=True)
        bc = upstream(nxt, attention_mask=extended_mask, past_key_values=b.past_key_values, use_cache=True)
        torch.testing.assert_close(ac.logits, bc.logits, rtol=0, atol=0)
        full = custom(torch.cat([ids, nxt], dim=1), attention_mask=extended_mask, use_cache=False)
        torch.testing.assert_close(ac.logits[:, -1], full.logits[:, -1], rtol=1e-4, atol=1e-5)
    print('Upstream parity: exact prefill and cached decode; cached/full consistency passed.', flush=True)

    model, info = AutoModelForCausalLM.from_pretrained(
        path, trust_remote_code=True, local_files_only=True, dtype=torch.bfloat16,
        output_loading_info=True,
    )
    model.eval()
    for key in ['missing_keys', 'unexpected_keys', 'mismatched_keys', 'error_msgs']:
        assert not info.get(key), info
    info = {key: sorted(value) if isinstance(value, set) else value for key, value in info.items()}
    tied = model.lm_head.weight is model.model.embed_tokens.weight
    assert tied == config.tie_word_embeddings
    assert tied and "lm_head.weight" not in header
    assert sum(p.numel() for p in model.parameters()) == report["text_parameters"]
    assert not any(any(t in n for t in ['vision', 'audio', 'mtp']) for n, _ in model.named_parameters())
    tokenizer = AutoTokenizer.from_pretrained(path, local_files_only=True)
    prompt = tokenizer.apply_chat_template([{'role': 'user', 'content': '请用一句话介绍你自己。'}],
                                          tokenize=False, add_generation_prompt=True, enable_thinking=False)
    inputs = tokenizer(prompt, return_tensors='pt', add_special_tokens=False)
    with torch.inference_mode():
        model_inputs = {k: v for k, v in inputs.items() if k != 'token_type_ids'}
        logits = model(**model_inputs, logits_to_keep=1).logits
        assert torch.isfinite(logits).all()
        generated = model.generate(**model_inputs, max_new_tokens=32, do_sample=False)
    answer = tokenizer.decode(generated[0, inputs.input_ids.shape[1]:], skip_special_tokens=True)
    result = {'torch': torch.__version__, 'device': 'cpu', 'dtype': 'bfloat16',
              'tensor_integrity': f'{len(report["tensor_sha256"])}/{len(report["tensor_sha256"])} byte-identical to source',
              'upstream_small_model_prefill_and_decode': 'exact match',
              'cached_vs_full_forward': 'passed (rtol=1e-4, atol=1e-5)',
              'loading_info': info, 'tied_embeddings': tied, 'finite_logits': True,
              'generation': answer}
    (path / 'validation_report.json').write_text(json.dumps(result, indent=2, ensure_ascii=False) + '\n')
    print(json.dumps(result, indent=2, ensure_ascii=False))


if __name__ == '__main__':
    main()