"""Tiny Llama pretrain demo (CPU/4GB friendly). For 1B+ use cloud + DeepSpeed."""
import json
from pathlib import Path
import torch
from transformers import LlamaConfig, LlamaForCausalLM, Trainer, TrainingArguments
from torch.utils.data import Dataset

class JsonlDS(Dataset):
    def __init__(self, path, tok_len=128):
        self.rows = [json.loads(l) for l in open(path, encoding="utf-8")]
        self.tok_len = tok_len
    def __len__(self): return len(self.rows)
    def __getitem__(self, i):
        import random
        r = self.rows[i]
        # char-level pseudo-tokenization for demo (real run uses train_tokenizer.py + proper encode)
        s = (r.get("prompt", "") + "\n" + r.get("think", "") + "\n" + r.get("answer", ""))[:2000]
        ids = [min(ord(c), 30000) for c in s][:self.tok_len]
        ids += [0] * (self.tok_len - len(ids))
        return {"input_ids": torch.tensor(ids, dtype=torch.long), "labels": torch.tensor(ids, dtype=torch.long)}

def main():
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--epochs", type=int, default=3)
    ap.add_argument("--out", default="outputs/tiny-50M")
    ap.add_argument("--cpu", action="store_true", help="force CPU")
    a = ap.parse_args()
    cfg = LlamaConfig(hidden_size=512, num_hidden_layers=8, num_attention_heads=8,
                      num_key_value_heads=4, intermediate_size=1376, vocab_size=32000, max_position_embeddings=2048)
    model = LlamaForCausalLM(cfg)
    ds = JsonlDS("data/processed/train.jsonl")
    use_cuda = torch.cuda.is_available() and not a.cpu
    import os
    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
    args = TrainingArguments(output_dir=a.out, per_device_train_batch_size=1,
                             gradient_accumulation_steps=32,
                             num_train_epochs=a.epochs, logging_steps=5, save_steps=200,
                             save_total_limit=1, fp16=use_cuda, use_cpu=a.cpu,
                             gradient_checkpointing=True, report_to="none",
                             optim="adafactor",
                             dataloader_num_workers=0, dataloader_pin_memory=False,
                             max_grad_norm=0.0)
    Trainer(model=model, args=args, train_dataset=ds).train()
    model.save_pretrained(a.out)
    print(f"saved -> {a.out}")

if __name__ == "__main__":
    main()