"""
Opus-Candid v2 — Qwen 3 Fine-Tuning Script
============================================
Trains 8B and 32B sequentially on RTX PRO 6000 Blackwell (96GB VRAM)
8B: Full 16-bit LoRA (fits in VRAM easily)
32B: 4-bit QLoRA

Usage:
  python train_qwen3.py
"""

import json
import os
import sys
import torch
import gc
from datetime import datetime
from pathlib import Path

# ============================================================
# CONFIG
# ============================================================

DATASET_PATH = "/workspace/opus_candid_full_dataset.json"
OUTPUT_BASE = "/workspace/opus_candid_v2"
HF_ORG = "Verdugie"

MODELS = [
    {
        "name": "Opus-Candid-8B-v2",
        "base_model": "Qwen/Qwen3-8B",
        "load_in_4bit": False,       # 16-bit — fits on 96GB easily
        "lora_r": 16,
        "lora_alpha": 32,
        "per_device_batch_size": 4,
        "gradient_accumulation": 4,   # effective batch = 16
        "num_epochs": 3,
        "learning_rate": 2e-4,
        "max_seq_length": 4096,
        "warmup_ratio": 0.05,
    },
    {
        "name": "Opus-Candid-32B-v2",
        "base_model": "Qwen/Qwen3-32B",
        "load_in_4bit": True,         # QLoRA — 32B needs quantization
        "lora_r": 16,
        "lora_alpha": 32,
        "per_device_batch_size": 2,
        "gradient_accumulation": 8,   # effective batch = 16
        "num_epochs": 3,
        "learning_rate": 1e-4,        # slightly lower for larger model
        "max_seq_length": 4096,
        "warmup_ratio": 0.05,
    },
]

LORA_TARGET_MODULES = [
    "q_proj", "k_proj", "v_proj", "o_proj",
    "gate_proj", "up_proj", "down_proj",
]

# ============================================================
# DATASET LOADING
# ============================================================

def load_and_prepare_dataset(tokenizer, max_seq_length):
    """Load ShareGPT-format dataset and convert to ChatML for Qwen 3."""
    from datasets import Dataset

    print(f"  Loading dataset from {DATASET_PATH}...")
    with open(DATASET_PATH, "r", encoding="utf-8") as f:
        raw_data = json.load(f)

    print(f"  Total conversations: {len(raw_data)}")

    # Convert ShareGPT format (from/value) to ChatML (role/content)
    role_map = {"human": "user", "gpt": "assistant", "system": "system"}

    formatted = []
    skipped = 0
    for conv in raw_data:
        messages = []
        for turn in conv["conversations"]:
            role = role_map.get(turn["from"], turn["from"])
            messages.append({"role": role, "content": turn["value"]})
        
        # Apply chat template WITHOUT thinking mode
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=False,
            enable_thinking=False,  # Critical: disable thinking for Opus-Candid
        )
        
        # Check length
        token_count = len(tokenizer.encode(text))
        if token_count > max_seq_length:
            skipped += 1
            continue
        
        formatted.append({"text": text})

    print(f"  Formatted: {len(formatted)} | Skipped (too long): {skipped}")

    # Train/eval split
    dataset = Dataset.from_list(formatted)
    split = dataset.train_test_split(test_size=0.05, seed=42)

    print(f"  Train: {len(split['train'])} | Eval: {len(split['test'])}")
    return split


# ============================================================
# TRAINING
# ============================================================

def train_model(config):
    """Train a single model with the given config."""
    
    print("\n" + "=" * 60)
    print(f"TRAINING: {config['name']}")
    print(f"Base model: {config['base_model']}")
    print(f"{'QLoRA (4-bit)' if config['load_in_4bit'] else 'LoRA (16-bit)'}")
    print("=" * 60)

    start_time = datetime.now()
    output_dir = os.path.join(OUTPUT_BASE, config["name"])
    os.makedirs(output_dir, exist_ok=True)

    # --- Load model ---
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, TaskType
    from trl import SFTTrainer, SFTConfig

    print("\n  Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(
        config["base_model"],
        trust_remote_code=True,
    )

    # Qwen 3 has no pad token — use eos
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id

    print(f"  EOS token: {tokenizer.eos_token} (id: {tokenizer.eos_token_id})")
    print(f"  PAD token: {tokenizer.pad_token} (id: {tokenizer.pad_token_id})")

    # --- Prepare dataset ---
    dataset = load_and_prepare_dataset(tokenizer, config["max_seq_length"])

    # --- Load model ---
    print(f"\n  Loading model: {config['base_model']}...")

    model_kwargs = {
        "trust_remote_code": True,
        "torch_dtype": torch.bfloat16,
        "device_map": "auto",
        "attn_implementation": "flash_attention_2",
    }

    if config["load_in_4bit"]:
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )
        model_kwargs["quantization_config"] = bnb_config

    model = AutoModelForCausalLM.from_pretrained(
        config["base_model"],
        **model_kwargs,
    )

    if config["load_in_4bit"]:
        model = prepare_model_for_kbit_training(model)

    # --- LoRA ---
    print("  Applying LoRA config...")
    lora_config = LoraConfig(
        r=config["lora_r"],
        lora_alpha=config["lora_alpha"],
        lora_dropout=0.05,
        target_modules=LORA_TARGET_MODULES,
        bias="none",
        task_type=TaskType.CAUSAL_LM,
    )

    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # --- Training args ---
    print("  Configuring trainer...")

    effective_batch = config["per_device_batch_size"] * config["gradient_accumulation"]
    steps_per_epoch = len(dataset["train"]) // effective_batch
    total_steps = steps_per_epoch * config["num_epochs"]
    eval_steps = max(steps_per_epoch // 4, 10)  # eval ~4x per epoch
    save_steps = steps_per_epoch  # save each epoch

    print(f"  Effective batch size: {effective_batch}")
    print(f"  Steps per epoch: {steps_per_epoch}")
    print(f"  Total steps: {total_steps}")
    print(f"  Eval every: {eval_steps} steps")

    training_args = SFTConfig(
        output_dir=output_dir,
        num_train_epochs=config["num_epochs"],
        per_device_train_batch_size=config["per_device_batch_size"],
        per_device_eval_batch_size=1,
        gradient_accumulation_steps=config["gradient_accumulation"],
        learning_rate=config["learning_rate"],
        lr_scheduler_type="cosine",
        warmup_ratio=config["warmup_ratio"],
        weight_decay=0.01,
        bf16=True,
        logging_steps=10,
        eval_strategy="steps",
        eval_steps=eval_steps,
        save_strategy="steps",
        save_steps=save_steps,
        save_total_limit=2,
        max_seq_length=config["max_seq_length"],
        dataset_text_field="text",
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        optim="adamw_torch_fused",
        report_to="none",
        padding="right",  # Qwen 3 prefers right padding
    )

    # --- Train ---
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        processing_class=tokenizer,
    )

    print(f"\n  🚀 Training {config['name']}...\n")
    train_result = trainer.train()

    # --- Save ---
    print(f"\n  Saving LoRA adapters to {output_dir}...")
    trainer.save_model(output_dir)
    tokenizer.save_pretrained(output_dir)

    # Save training stats
    stats = {
        "model": config["name"],
        "base_model": config["base_model"],
        "quantization": "4-bit QLoRA" if config["load_in_4bit"] else "16-bit LoRA",
        "dataset_size": len(dataset["train"]),
        "eval_size": len(dataset["test"]),
        "epochs": config["num_epochs"],
        "learning_rate": config["learning_rate"],
        "lora_r": config["lora_r"],
        "lora_alpha": config["lora_alpha"],
        "effective_batch_size": effective_batch,
        "max_seq_length": config["max_seq_length"],
        "train_loss": train_result.training_loss,
        "train_runtime_seconds": train_result.metrics["train_runtime"],
        "train_samples_per_second": train_result.metrics["train_samples_per_second"],
        "timestamp": datetime.now().isoformat(),
    }

    with open(os.path.join(output_dir, "training_stats.json"), "w") as f:
        json.dump(stats, f, indent=2)

    elapsed = datetime.now() - start_time
    print(f"\n  ✅ {config['name']} complete!")
    print(f"  Training loss: {train_result.training_loss:.4f}")
    print(f"  Time: {elapsed}")
    print(f"  Adapters saved: {output_dir}")

    # --- Cleanup ---
    del model, trainer
    gc.collect()
    torch.cuda.empty_cache()

    return stats


# ============================================================
# MAIN
# ============================================================

def main():
    print("=" * 60)
    print("OPUS-CANDID v2 — Qwen 3 Training Pipeline")
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB")
    print(f"Models to train: {', '.join(m['name'] for m in MODELS)}")
    print("=" * 60)

    # Verify dataset exists
    if not os.path.exists(DATASET_PATH):
        print(f"\n❌ Dataset not found: {DATASET_PATH}")
        sys.exit(1)

    with open(DATASET_PATH, "r", encoding="utf-8") as f:
        data = json.load(f)
    print(f"\nDataset: {len(data)} conversations")

    all_stats = []

    for config in MODELS:
        try:
            stats = train_model(config)
            all_stats.append(stats)
        except Exception as e:
            print(f"\n❌ FAILED: {config['name']}")
            print(f"   Error: {e}")
            import traceback
            traceback.print_exc()
            # Continue to next model
            gc.collect()
            torch.cuda.empty_cache()
            continue

    # Summary
    print("\n" + "=" * 60)
    print("TRAINING COMPLETE — SUMMARY")
    print("=" * 60)
    for s in all_stats:
        print(f"\n  {s['model']}:")
        print(f"    Loss: {s['train_loss']:.4f}")
        print(f"    Time: {s['train_runtime_seconds']:.0f}s")
        print(f"    Saved: {OUTPUT_BASE}/{s['model']}")

    print(f"\nNext steps:")
    print(f"  1. Merge adapters into base models")
    print(f"  2. Convert to GGUF")
    print(f"  3. Quantize Q4_K_M and Q8_0")
    print(f"  4. Upload to HuggingFace")


if __name__ == "__main__":
    main()