"""sakthai_lora_train_1.5b.py
LoRA fine-tune SakThai on Qwen2.5-1.5B-Instruct with v4 curated dataset.
Runs on HF Jobs (t4-small / L4).

After training the adapter is pushed to:
    https://huggingface.co/Nanthasit/sakthai-context-1.5b-tools
"""

import os, sys

try:
    from datasets import load_dataset
    from transformers import (
        AutoTokenizer,
        TrainingArguments,
        Trainer,
        DataCollatorForLanguageModeling,
    )
    from transformers import Qwen2ForCausalLM
    from peft import LoraConfig, get_peft_model
except ImportError as e:
    print(f"❌ Missing dependency: {e}")
    sys.exit(1)

# ── Config (optimised for 16GB T4) ────────────────────────────────────────────
BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
DATASET = "Nanthasit/sakthai-combined-v4"
TARGET_REPO = "Nanthasit/sakthai-context-1.5b-tools"
MERGE_REPO = "Nanthasit/sakthai-context-1.5b-merged"
OUTPUT_DIR = "/tmp/lora-adapter"

MAX_LENGTH = 768       # longer context for tool definitions
LR = 2e-4
EPOCHS = 4             # more epochs on cleaner v4
BATCH_SIZE = 1         # 1.5B is bigger — 1 per device
GRAD_ACCUM = 16        # effective batch = 16
WARMUP_RATIO = 0.1
WEIGHT_DECAY = 0.01

PUSH_TO_HUB = "--no-push" not in sys.argv

import transformers as _tf
_TF_MAJOR = int(_tf.__version__.split(".")[0])
print(f"📊 transformers v{_tf.__version__}")

# ── 1. Dataset ───────────────────────────────────────────────────────────────
print(f"\n📦 Loading dataset: {DATASET}")
ds = load_dataset(DATASET, split="train")
print(f"   {len(ds)} examples loaded")

# ── 2. Tokenizer + model ─────────────────────────────────────────────────────
print(f"\n📥 Loading base model: {BASE_MODEL}")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = Qwen2ForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype="auto",
    device_map="auto",
)

# Enable gradient checkpointing to save memory
model.gradient_checkpointing_enable()
model.config.use_cache = False
print(f"   Model loaded ({sum(p.numel() for p in model.parameters()):,} params)")

# ── 3. LoRA ──────────────────────────────────────────────────────────────────
print("\n🔧 Applying LoRA (r=16, alpha=32)")
lora_config = LoraConfig(
    r=16, lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.1, bias="none", task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# ── 4. Format ────────────────────────────────────────────────────────────────
def format_example(ex):
    msgs = ex.get("messages", [])
    tools = ex.get("tools", [])
    text = tokenizer.apply_chat_template(
        msgs, tools=tools or None,
        tokenize=False, add_generation_prompt=False,
    )
    return {"text": text}

print("\n🔄 Formatting...")
ds = ds.map(format_example)

def tok_fn(examples):
    return tokenizer(
        examples["text"], truncation=True,
        max_length=MAX_LENGTH, padding="max_length",
    )

ds = ds.map(tok_fn, batched=True, remove_columns=ds.column_names)
ds = ds.train_test_split(test_size=0.1, seed=42)
print(f"   Train: {len(ds['train'])}  |  Eval: {len(ds['test'])}")

# ── 5. Training ──────────────────────────────────────────────────────────────
print(f"\n🏋️  Training ({EPOCHS} epochs, LR={LR}, batch={BATCH_SIZE}×{GRAD_ACCUM})")
args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    per_device_train_batch_size=BATCH_SIZE,
    gradient_accumulation_steps=GRAD_ACCUM,
    learning_rate=LR,
    num_train_epochs=EPOCHS,
    warmup_ratio=WARMUP_RATIO,
    weight_decay=WEIGHT_DECAY,
    fp16=True,
    logging_steps=5,
    save_strategy="no",
    report_to="none",
    remove_unused_columns=False,
    ddp_find_unused_parameters=None,
    optim="adamw_torch",
)

kw = dict(model=model, args=args, train_dataset=ds["train"],
          eval_dataset=ds["test"],
          data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False))
if _TF_MAJOR < 5:
    kw["tokenizer"] = tokenizer

trainer = Trainer(**kw)
trainer.train()

# ── 6. Save ──────────────────────────────────────────────────────────────────
print(f"\n💾 Saving adapter to {OUTPUT_DIR}")
model.save_pretrained(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)

# ── 7. Push ──────────────────────────────────────────────────────────────────
if PUSH_TO_HUB:
    print(f"\n☁️  Pushing to {TARGET_REPO}...")
    try:
        from huggingface_hub import HfApi
        HfApi().upload_folder(
            repo_id=TARGET_REPO,
            folder_path=OUTPUT_DIR,
            repo_type="model",
            commit_message=f"sakthai-lora-1.5b r=16 alpha=32 epoch={EPOCHS} v4-dataset",
        )
        print(f"✅ Adapter at https://huggingface.co/{TARGET_REPO}")
    except Exception as e:
        print(f"❌ Push failed: {e}")
else:
    print(f"\n⏭️  Push skipped. Adapter at {OUTPUT_DIR}")

print(f"""
{'='*50}
✅ TRAINING COMPLETE
   Base:    {BASE_MODEL}
   Dataset: {DATASET}
   Adapter: {TARGET_REPO}
   Epochs:  {EPOCHS}
{'='*50}
""")