import os, json, torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from tokenizers import Tokenizer
from modeling_m31 import M31Model
from generation_m31 import generate

REPO_ID = os.environ.get('M31_REPO_ID', 'eshanized/M31Tesla')
CHECKPOINTS = {
    'pretrain': 'experiments/M31-Python-Agent-220M-v5/checkpoints/pretrain/step-00000733',
    'sft':      'experiments/M31-Python-Agent-220M-v5/checkpoints/sft/step-00000123',
    'agent':    'experiments/M31-Python-Agent-220M-v5/checkpoints/agent/step-00000123',
    'repair':   'experiments/M31-Python-Agent-220M-v5/checkpoints/repair/step-00000123',
}

def load_stage(stage='repair', device=None):
    if stage not in CHECKPOINTS:
        raise ValueError(f'Unknown stage {stage!r}: {sorted(CHECKPOINTS)}')
    device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
    prefix = CHECKPOINTS[stage]
    config_path = hf_hub_download(REPO_ID, f'{prefix}/config.json')
    weights_path = hf_hub_download(REPO_ID, f'{prefix}/model.safetensors')
    tokenizer_path = hf_hub_download(REPO_ID, 'tokenizer.json')
    with open(config_path, 'r', encoding='utf-8') as f:
        config = json.load(f)
    model = M31Model()
    state = load_file(weights_path, device='cpu')
    model.load_state_dict(state, strict=True)
    model.to(device).eval()
    tokenizer = Tokenizer.from_file(tokenizer_path)
    return model, tokenizer, config