import argparse, os, sys
from typing import List
from datasets import load_dataset
from transformers import (
    AutoTokenizer, AutoModelForCausalLM,
    DataCollatorForLanguageModeling, TrainingArguments, Trainer
)

def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--dataset", required=True, help="JSON/JSONL (.jsonl or .jsonl.gz)")
    p.add_argument("--output", default="trained_model")
    p.add_argument("--model_name", default="distilgpt2")
    p.add_argument("--epochs", type=float, default=0.5)
    p.add_argument("--batch_size", type=int, default=2)
    p.add_argument("--block_size", type=int, default=256)
    p.add_argument("--learning_rate", type=float, default=5e-5)
    # quick mode:
    p.add_argument("--quick", type=int, default=0)        # 1 => tiny model + fast
    p.add_argument("--max_steps", type=int, default=0)    # >0 overrides epochs
    p.add_argument("--subset", type=int, default=0)       # use first N rows
    return p.parse_args()

def main():
    a = parse_args()

    if a.quick:
        a.model_name = "sshleifer/tiny-gpt2"  # ultra-tiny, very fast
        if a.max_steps <= 0: a.max_steps = 8
        if a.subset <= 0: a.subset = 32
        a.epochs = 1.0

    print(f"📥 Loading dataset: {a.dataset}", flush=True)
    ds = load_dataset("json", data_files=a.dataset, split="train")
    cols = ds.column_names
    print("🧾 Columns:", cols, flush=True)

    if a.subset and a.subset > 0:
        ds = ds.select(range(min(a.subset, len(ds))))
        print(f"✂ Using subset: {len(ds)} rows", flush=True)

    tok = AutoTokenizer.from_pretrained(a.model_name)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    model = AutoModelForCausalLM.from_pretrained(a.model_name)

    def build_texts(batch) -> List[str]:
        if "text" in batch:
            return [str(t) for t in batch["text"]]
        if "prompt" in batch and "completion" in batch:
            return [f"{str(p).rstrip()}\n{str(c)}" for p, c in zip(batch["prompt"], batch["completion"])]
        raise ValueError("Dataset must contain 'text' OR both 'prompt' and 'completion'.")

    def tokenize(batch):
        texts = build_texts(batch)
        return tok(texts, padding="max_length", truncation=True, max_length=a.block_size)

    print("🔁 Tokenizing…", flush=True)
    tokds = ds.map(tokenize, batched=True, remove_columns=cols)
    collator = DataCollatorForLanguageModeling(tokenizer=tok, mlm=False)

    print("⚙ Trainer…", flush=True)
    args = TrainingArguments(
        output_dir=a.output,
        overwrite_output_dir=True,
        per_device_train_batch_size=a.batch_size,
        num_train_epochs=a.epochs if a.max_steps == 0 else 1,
        learning_rate=a.learning_rate,
        logging_steps=1,
        save_steps=50,
        save_total_limit=1,
        report_to=[],
        fp16=False,
        max_steps=a.max_steps if a.max_steps > 0 else -1,
    )
    trainer = Trainer(model=model, args=args, train_dataset=tokds, tokenizer=tok, data_collator=collator)

    print("🚀 Training…", flush=True)
    trainer.train()
    print(f"💾 Saving to {a.output}", flush=True)
    os.makedirs(a.output, exist_ok=True)
    trainer.save_model(a.output)
    tok.save_pretrained(a.output)
    print("✅ Done.", flush=True)

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"❌ Training failed: {e}", flush=True)
        sys.exit(1)