import hashlib
import json
import os
from pathlib import Path
import shutil
import sys

import torch
from peft import PeftModel
from safetensors import safe_open
from transformers import AutoModelForCausalLM

base, checkpoint, target = (Path(x) for x in sys.argv[1:])
base_revision = '2c95908929198d6e69af8638f0dbbd9bc6b93f9e'
pinned = Path.home()/'.cache/huggingface/hub/models--ministral--Ministral-3b-instruct/snapshots'/base_revision
stage = target.with_name(target.name + '_staging_' + os.environ['SLURM_JOB_ID'])

def digest(path):
    with path.open('rb') as handle:
        return hashlib.file_digest(handle, 'sha256').hexdigest()

if target.exists():
    raise SystemExit('Merged output already exists')
index = 'model.safetensors.index.json'
metadata = ('config.json','generation_config.json','special_tokens_map.json','tokenizer_config.json','tokenizer.json')
original_keys = json.loads((base/index).read_text())['weight_map']
for name in (*metadata, index, *sorted(set(original_keys.values()))):
    if digest(base/name) != digest(pinned/name):
        raise RuntimeError(f'Baseline file differs from the pinned training snapshot: {name}')
print('Baseline model matches the pinned training snapshot', base_revision, flush=True)

torch.set_num_threads(8)
model = AutoModelForCausalLM.from_pretrained(str(base), dtype=torch.float32, local_files_only=True, trust_remote_code=False)
lora = PeftModel.from_pretrained(model, str(checkpoint), is_trainable=False)
with torch.no_grad():
    merged = lora.merge_and_unload(safe_merge=True).to(torch.bfloat16)
stage.mkdir()
merged.save_pretrained(str(stage), safe_serialization=True, max_shard_size='3GB')
for name in metadata:
    shutil.copy2(base/name, stage/name)

saved_keys = json.loads((stage/index).read_text())['weight_map']
if set(saved_keys) != set(original_keys) or any('lora_' in key for key in saved_keys):
    raise RuntimeError('Unexpected merged parameter set')
for key, expected_to_change in (
    ('model.layers.0.self_attn.q_proj.weight', True),
    ('model.embed_tokens.weight', False),
):
    with safe_open(str(base/original_keys[key]), framework='pt', device='cpu') as left:
        with safe_open(str(stage/saved_keys[key]), framework='pt', device='cpu') as right:
            original, merged_tensor = left.get_tensor(key), right.get_tensor(key)
    changed = not torch.equal(original, merged_tensor)
    if merged_tensor.dtype != torch.bfloat16 or changed != expected_to_change:
        raise RuntimeError(f'Unexpected merged tensor: {key}')
(stage/'merge_provenance.json').write_text(json.dumps({
    'checkpoint': str(checkpoint), 'adapter_sha256': digest(checkpoint/'adapter_model.safetensors'),
    'base_revision': base_revision, 'output_dtype': 'bfloat16',
    'tokenizer_and_config': {name: digest(stage/name) for name in metadata},
}, indent=2) + chr(10))
stage.rename(target)
print('MERGED_MODEL_READY', target, flush=True)