"""
train.py — Unsloth QLoRA fine-tuning for Qwen2.5-Coder-7B-Instruct.

Trains on synthetic NL→SQL pairs from train.jsonl using:
  - 4-bit QLoRA (bitsandbytes)
  - Unsloth acceleration (2x faster, 50% less VRAM)
  - SFTTrainer from TRL
  - ~3 epochs on A10G (24 GB)

Output: lora-adapter/ directory with adapter weights.
"""

import json
import torch
from pathlib import Path
from datasets import Dataset
from unsloth import FastLanguageModel
from unsloth import is_bfloat16_supported
from trl import SFTTrainer
from transformers import TrainingArguments


# ── Configuration ──────────────────────────────────────────────────────

MODEL_NAME = "unsloth/Qwen2.5-Coder-7B-Instruct"
MAX_SEQ_LENGTH = 2048
LORA_R = 16
LORA_ALPHA = 16
LORA_DROPOUT = 0.0

TRAIN_JSONL = Path(__file__).parent / "train.jsonl"
OUTPUT_DIR = Path(__file__).parent / "lora-adapter"

# Training hyperparams (tuned for ~1200 pairs, A10G, ~3-6 hrs)
NUM_EPOCHS = 3
BATCH_SIZE = 4
GRAD_ACCUM = 4
LEARNING_RATE = 2e-4
WARMUP_RATIO = 0.1
LOGGING_STEPS = 10
SAVE_STEPS = 200


# ── Chat template ──────────────────────────────────────────────────────

SYSTEM_PROMPT = (
    "You are an expert DuckDB SQL developer for school district administration. "
    "Generate ONLY valid DuckDB SQL queries wrapped in ```sql ``` markdown blocks. "
    "Never produce INSERT, UPDATE, DELETE, DROP, or ALTER statements."
)

def format_chat(question: str, sql: str) -> str:
    """
    Format a training example using Qwen2.5 chat template.

    <|im_start|>system
    {system}
    <|im_end|>
    <|im_start|>user
    Question: {question}
    <|im_end|>
    <|im_start|>assistant
    ```sql
    {sql}
    ```<|im_end|>
    """
    return (
        f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
        f"<|im_start|>user\nQuestion: {question}<|im_end|>\n"
        f"<|im_start|>assistant\n```sql\n{sql}\n```<|im_end|>"
    )


# ── Data loading ───────────────────────────────────────────────────────

def load_training_data(path: Path = TRAIN_JSONL) -> Dataset:
    """Load JSONL training pairs and format with chat template."""
    if not path.exists():
        raise FileNotFoundError(
            f"train.jsonl not found at {path}. Run generate_synthetic.py first."
        )

    texts = []
    with open(path) as f:
        for line in f:
            pair = json.loads(line)
            text = format_chat(pair["question"], pair["sql"])
            texts.append({"text": text})

    dataset = Dataset.from_list(texts)
    print(f"📊 Loaded {len(dataset)} training examples from {path}")
    return dataset


# ── Model loading ──────────────────────────────────────────────────────

def load_model_and_tokenizer():
    """Load Qwen2.5-Coder-7B in 4-bit with Unsloth optimization."""
    print(f"🦙 Loading {MODEL_NAME} (4-bit QLoRA)...")

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=MODEL_NAME,
        max_seq_length=MAX_SEQ_LENGTH,
        dtype=None,  # auto-detect
        load_in_4bit=True,
    )

    # Apply QLoRA adapters
    model = FastLanguageModel.get_peft_model(
        model,
        r=LORA_R,
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
        ],
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=42,
    )

    # Unsloth patching for faster training
    FastLanguageModel.for_training(model)

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print(f"✅ Model loaded. Trainable: {trainable:,} / {total:,} params "
          f"({100*trainable/total:.1f}%)")

    return model, tokenizer


# ── Training ───────────────────────────────────────────────────────────

def train(model, tokenizer, dataset: Dataset):
    """Run QLoRA fine-tuning with SFTTrainer."""
    print(f"\n🏋️ Starting training: {NUM_EPOCHS} epochs, "
          f"batch={BATCH_SIZE}, grad_accum={GRAD_ACCUM}")

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=dataset,
        dataset_text_field="text",
        max_seq_length=MAX_SEQ_LENGTH,
        args=TrainingArguments(
            output_dir=str(OUTPUT_DIR / "checkpoints"),
            per_device_train_batch_size=BATCH_SIZE,
            gradient_accumulation_steps=GRAD_ACCUM,
            warmup_ratio=WARMUP_RATIO,
            num_train_epochs=NUM_EPOCHS,
            learning_rate=LEARNING_RATE,
            fp16=not is_bfloat16_supported(),
            bf16=is_bfloat16_supported(),
            logging_steps=LOGGING_STEPS,
            save_strategy="no",  # manual save only — prevent pickle error
            optim="adamw_8bit",
            seed=42,
            report_to="none",  # No wandb on Modal
        ),
    )

    import traceback
    try:
        trainer.train()
    except Exception as e:
        if "Pickle" in str(e) or "pickle" in str(e):
            print("\n⚠️  Trainer save failed (pickle error on args) — continuing with manual save...")
            traceback.print_exc()
        else:
            raise

    # Save final adapter (manual — bypasses trainer's pickle-prone save)
    print("\n💾 Saving adapter weights...")
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(str(OUTPUT_DIR))
    tokenizer.save_pretrained(str(OUTPUT_DIR))
    print(f"✅ Adapter saved to {OUTPUT_DIR}")

    return trainer


# ── Main ───────────────────────────────────────────────────────────────

def main():
    # Skip training if adapter already exists on volume
    if (Path("/data") / "lora-adapter" / "adapter_config.json").exists():
        print("✅ LoRA adapter already exists in volume — skipping training")
        return
    
    if OUTPUT_DIR.is_dir() and (OUTPUT_DIR / "adapter_config.json").exists():
        print("✅ LoRA adapter already exists — skipping training")
        return
    
    dataset = load_training_data()
    model, tokenizer = load_model_and_tokenizer()
    train(model, tokenizer, dataset)
    
    # Free GPU memory before merge step
    import gc
    del model
    del tokenizer
    gc.collect()
    import torch
    torch.cuda.empty_cache()
    print("🧹 GPU memory cleared for merge step")
    
    print("\n🎉 Training complete!")


if __name__ == "__main__":
    main()