"""Dry-run SFT: load data, build model, resize embeddings, one forward pass."""
import sys, os, json, torch
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__))))

from sft import load_special_tokens, load_sft_dataset, get_sft_batch
from model import ModelConfig, Retriever500M
from tokenizers import Tokenizer

def main():
    load_special_tokens()
    tok = Tokenizer.from_file("tokenizer/tokenizer_agent.json")
    print("Loading dataset...")
    dataset = load_sft_dataset("data/sft_traces.jsonl", tok, max_seq_len=768)
    print(f"Dataset size: {len(dataset)}")

    config = ModelConfig(vocab_size=32009, d_model=1280, n_layers=23, n_heads=20,
                         d_ff=3456, max_seq_len=768, tie_embeddings=True)
    model = Retriever500M(config).cuda()
    print(f"Model: {model.count_parameters()/1e6:.1f}M params")

    # Load old checkpoint and resize embeddings
    ckpt = torch.load("checkpoints/latest.pt", map_location="cuda", weights_only=False)
    state = ckpt["model_state_dict"]
    old_w = state["token_embedding.weight"]
    new_w = torch.zeros(32009, 1280)
    new_w[:32000] = old_w
    torch.nn.init.normal_(new_w[32000:], mean=0.0, std=0.02)
    state["token_embedding.weight"] = new_w
    model.load_state_dict(state)
    print("Checkpoint loaded with resized embeddings")

    # One forward pass
    input_ids, targets, loss_mask = get_sft_batch(dataset, batch_size=2, seq_len=768, device=torch.device("cuda"))
    print(f"Batch: {input_ids.shape}")

    with torch.autocast("cuda", dtype=torch.bfloat16):
        out = model(input_ids, targets=targets)
        print(f"Forward pass OK: loss={out['loss'].item():.4f}")

    print("SFT DRY RUN PASSED")

if __name__ == "__main__":
    main()