"""SFT LoRA for text->SVG. Runs on a Runpod pod. Wall-clock bounded; pushes adapter to HF; caller kills pod."""
import os, sys, time, json, random, math
import torch
from datasets import load_dataset
from transformers import (AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments,
                          TrainerCallback)
from peft import LoraConfig, get_peft_model

BASE = os.environ.get("BASE", "Qwen/Qwen3.5-4B")
DATA = os.environ.get("DATA", "aixonlab/svg-sft-v1")
OUT_REPO = os.environ.get("OUT_REPO", "aixonlab/svg-slm-4b-lora")
MAX_LEN = int(os.environ.get("MAX_LEN", 6400))
MAX_SECONDS = float(os.environ.get("TRAIN_SECONDS", 3 * 3600))   # hard wall-clock for training loop
MAX_STEPS = int(os.environ.get("MAX_STEPS", -1))
LR = float(os.environ.get("LR", 1e-4))
TOKEN = os.environ["HF_TOKEN"]
t0 = time.time()

tok = AutoTokenizer.from_pretrained(BASE, token=TOKEN)
model = AutoModelForCausalLM.from_pretrained(BASE, torch_dtype=torch.bfloat16, token=TOKEN,
                                             attn_implementation="sdpa", device_map={"": 0})
model.gradient_checkpointing_enable()
model.enable_input_require_grads()
cfg = LoraConfig(r=32, lora_alpha=64, lora_dropout=0.0, task_type="CAUSAL_LM",
                 target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"])
model = get_peft_model(model, cfg)
model.print_trainable_parameters()

ds = load_dataset(DATA, token=TOKEN)
train = ds["train"].shuffle(seed=0)

def encode(ex):
    msgs = ex["messages"]
    try:
        prompt = tok.apply_chat_template(msgs[:1], tokenize=False, add_generation_prompt=True, enable_thinking=False)
    except TypeError:
        prompt = tok.apply_chat_template(msgs[:1], tokenize=False, add_generation_prompt=True)
    p = tok(prompt, add_special_tokens=False)["input_ids"]
    a = tok(msgs[1]["content"] + tok.eos_token if False else msgs[1]["content"] + "<|im_end|>\n",
            add_special_tokens=False)["input_ids"]
    ids = (p + a)[:MAX_LEN]
    labels = ([-100] * len(p) + a)[:MAX_LEN]
    return {"input_ids": ids, "labels": labels}

train = train.map(encode, remove_columns=train.column_names, num_proc=4)
train = train.filter(lambda e: len(e["input_ids"]) < MAX_LEN)  # drop truncated

def collate(batch):
    m = max(len(b["input_ids"]) for b in batch)
    pad = tok.pad_token_id or 0
    ids = torch.tensor([b["input_ids"] + [pad] * (m - len(b["input_ids"])) for b in batch])
    lab = torch.tensor([b["labels"] + [-100] * (m - len(b["labels"])) for b in batch])
    att = torch.tensor([[1] * len(b["input_ids"]) + [0] * (m - len(b["input_ids"])) for b in batch])
    return {"input_ids": ids, "labels": lab, "attention_mask": att}

class TimeLimit(TrainerCallback):
    def on_step_end(self, args, state, control, **kw):
        if time.time() - t0 > MAX_SECONDS:
            print("TIME LIMIT HIT -> stopping", flush=True)
            control.should_training_stop = True
        return control

# Estimate steps so cosine schedule completes inside the time budget: use max_steps after a probe.
bs, ga = int(os.environ.get("BS",1)), int(os.environ.get("GA",16))
n = len(train)
steps_full = n // (bs * ga)
args = TrainingArguments(
    output_dir="/workspace/out", per_device_train_batch_size=bs, gradient_accumulation_steps=ga,
    learning_rate=LR, lr_scheduler_type="cosine", warmup_steps=20, bf16=True, logging_steps=5,
    save_strategy="steps", save_steps=100, save_total_limit=1, max_steps=MAX_STEPS if MAX_STEPS > 0 else steps_full,
    report_to="none", dataloader_num_workers=2, remove_unused_columns=False,
    gradient_checkpointing=True, optim="adamw_torch_fused", weight_decay=0.0, max_grad_norm=1.0)
trainer = Trainer(model=model, args=args, train_dataset=train, data_collator=collate, callbacks=[TimeLimit()])

def push(tag):
    try:
        from huggingface_hub import HfApi
        api = HfApi(token=TOKEN)
        api.create_repo(OUT_REPO, exist_ok=True, private=False)
        model.save_pretrained("/workspace/adapter")
        tok.save_pretrained("/workspace/adapter")
        api.upload_folder(folder_path="/workspace/adapter", repo_id=OUT_REPO, commit_message=f"adapter {tag}")
        print("PUSHED", tag, flush=True)
    except Exception as e:
        print("PUSH FAILED", e, flush=True)

class PushCkpt(TrainerCallback):
    def on_save(self, args, state, control, **kw):
        if state.global_step % 300 == 0:
            push(f"step{state.global_step}")
trainer.add_callback(PushCkpt())

print(f"examples={n} steps_planned={args.max_steps} setup_s={time.time()-t0:.0f}", flush=True)
if os.environ.get("PROBE"):
    t1 = time.time(); trainer.train(); sps = (time.time() - t1) / max(1, args.max_steps)
    json.dump({"sec_per_step": sps, "full_steps": steps_full}, open("/workspace/probe.json", "w"))
    print("PROBE", sps, steps_full, flush=True); sys.exit(0)
trainer.train()
push("final")
print("TRAINING DONE", flush=True)