"""BT-7274 LoRA — Qwen3.5-27B, bf16 LoRA (NOT QLoRA).

Unsloth docs: use FastLanguageModel for text-only training on Qwen3.5 VL models.
bf16 LoRA 27B ≈ 56GB VRAM — fits on H100 80GB.

Usage:
    pip install --upgrade unsloth unsloth_zoo
    python train_qwen35_27b.py
"""

from unsloth import FastLanguageModel
from trl import SFTTrainer, SFTConfig
from datasets import load_dataset
import torch

# --- Config ---
MODEL = "Qwen/Qwen3.5-27B"
MAX_SEQ = 4096
RANK = 16
ALPHA = 16
DATA = "./bt7274_v3.jsonl"
OUT = "./bt7274-qwen35-27b-lora-v1"
EPOCHS = 3
BATCH = 1
GRAD_ACCUM = 8
LR = 1e-4  # slightly lower for larger model

# --- Load model (bf16, NOT 4-bit) ---
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=MODEL,
    max_seq_length=MAX_SEQ,
    load_in_4bit=False,      # QLoRA not recommended for Qwen3.5
    load_in_16bit=True,      # bf16 LoRA
    full_finetuning=False,
    dtype=torch.bfloat16,
)

# --- LoRA adapter ---
model = FastLanguageModel.get_peft_model(
    model,
    r=RANK,
    lora_alpha=ALPHA,
    lora_dropout=0,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=42,
    max_seq_length=MAX_SEQ,
)

# --- Dataset ---
ds = load_dataset("json", data_files=DATA, split="train")

def fix_tool_calls(messages):
    """Parse tool_call arguments from JSON strings to dicts for Qwen3.5 template."""
    import json as _json
    fixed = []
    for msg in messages:
        msg = dict(msg)
        if msg.get("tool_calls"):
            new_tcs = []
            for tc in msg["tool_calls"]:
                tc = dict(tc)
                if "function" in tc:
                    fn = dict(tc["function"])
                    if isinstance(fn.get("arguments"), str):
                        try:
                            fn["arguments"] = _json.loads(fn["arguments"])
                        except (ValueError, TypeError):
                            fn["arguments"] = {"raw": fn["arguments"]}
                    tc["function"] = fn
                new_tcs.append(tc)
            msg["tool_calls"] = new_tcs
        fixed.append(msg)
    return fixed

def to_chatml(ex):
    messages = fix_tool_calls(ex["messages"])
    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=False
    )
    return {"text": text}

ds = ds.map(to_chatml)

steps = (len(ds) * EPOCHS) // (BATCH * GRAD_ACCUM)
print(f"Dataset: {len(ds)} examples")
print(f"Epochs: {EPOCHS}, effective batch: {BATCH * GRAD_ACCUM}")
print(f"Estimated steps: {steps}")
print(f"LoRA: r={RANK}, alpha={ALPHA}")
print(f"Max seq: {MAX_SEQ}")
print(f"Model: {MODEL}")

# --- Train ---
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=ds,
    args=SFTConfig(
        output_dir=OUT,
        per_device_train_batch_size=BATCH,
        gradient_accumulation_steps=GRAD_ACCUM,
        num_train_epochs=EPOCHS,
        learning_rate=LR,
        bf16=True,
        logging_steps=5,
        save_steps=50,
        save_total_limit=2,
        warmup_steps=10,
        optim="adamw_8bit",
        seed=42,
        report_to="none",
        max_seq_length=MAX_SEQ,
        dataset_num_proc=1,
    ),
)

trainer.train()

# --- Save LoRA adapter ---
model.save_pretrained(OUT)
tokenizer.save_pretrained(OUT)
print(f"\nSaved LoRA adapter to {OUT}/")