#!/usr/bin/env python3
"""
Harbour Fine-tuning Script for Qwen3.6-35B-A3B (MoE)
Uses Unsloth + LoRA with GGUF quantized model
Optimized for CPU with 121GB RAM
"""

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

# Configuration
MODEL_NAME = "unsloth/Qwen3.6-35B-A3B-GGUF"
MODEL_FILE = "Qwen3.6-35B-A3B-UD-Q4_K_M.gguf"
TRAIN_FILE = Path("/home/fivetech/finetune/harbour_train.jsonl")
VAL_FILE = Path("/home/fivetech/finetune/harbour_val.jsonl")
OUTPUT_DIR = Path("/home/fivetech/finetune/output")
MAX_SEQ_LENGTH = 2048

print("=" * 60)
print("Harbour Fine-tuning - Qwen3.6-35B-A3B (MoE) with Unsloth + LoRA")
print("=" * 60)

# 1. Load model from GGUF
print("\n1. Loading model from GGUF (Q4_K_M)...")
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=MODEL_NAME,
    gguf_file=MODEL_FILE,
    max_seq_length=MAX_SEQ_LENGTH,
    load_in_4bit=True,
    dtype=None,
)

# 2. LoRA configuration
print("2. Configuring LoRA...")
model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    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,
)

# 3. Load dataset
print("3. Loading dataset...")

def load_jsonl(path):
    data = []
    with open(path) as f:
        for line in f:
            data.append(json.loads(line))
    return data

train_data = load_jsonl(TRAIN_FILE)
val_data = load_jsonl(VAL_FILE)

print(f"   Train: {len(train_data)} entries")
print(f"   Val: {len(val_data)} entries")

# 4. Format conversations
print("4. Formatting conversations...")

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

train_dataset = Dataset.from_list([format_conversation(e) for e in train_data])
val_dataset = Dataset.from_list([format_conversation(e) for e in val_data])

# 5. Tokenize
print("5. Tokenizing...")

def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=MAX_SEQ_LENGTH,
        padding=False,
    )

train_dataset = train_dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=["text"],
    desc="Tokenizing train",
)
val_dataset = val_dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=["text"],
    desc="Tokenizing val",
)

print(f"   Train tokens: {sum(len(x) for x in train_dataset['input_ids']):,}")
print(f"   Val tokens: {sum(len(x) for x in val_dataset['input_ids']):,}")

# 6. Training arguments
print("6. Setting up training...")
training_args = TrainingArguments(
    output_dir=str(OUTPUT_DIR),
    num_train_epochs=3,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,
    learning_rate=1e-4,
    weight_decay=0.01,
    warmup_ratio=0.1,
    lr_scheduler_type="cosine",
    logging_steps=5,
    save_steps=50,
    save_total_limit=3,
    eval_strategy="steps",
    eval_steps=50,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    bf16=False,
    fp16=False,
    dataloader_num_workers=1,
    report_to="none",
    remove_unused_columns=False,
    max_grad_norm=1.0,
    optim="adamw_8bit",
)

# 7. Create trainer
print("7. Creating trainer...")
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    max_seq_length=MAX_SEQ_LENGTH,
    dataset_text_field="text",
)

# 8. Train
print("\n8. Starting training...")
print("=" * 60)
trainer.train()

# 9. Save LoRA adapter
print("\n9. Saving LoRA adapter...")
trainer.save_model(str(OUTPUT_DIR / "final"))
tokenizer.save_pretrained(str(OUTPUT_DIR / "final"))

# 10. Export to GGUF (optional)
print("\n10. Exporting to GGUF...")
model.save_pretrained_gguf(
    str(OUTPUT_DIR / "gguf"),
    tokenizer,
    quantization_method="q4_k_m",
)

print("\n" + "=" * 60)
print("Training complete!")
print(f"LoRA adapter saved to: {OUTPUT_DIR / 'final'}")
print(f"GGUF model saved to: {OUTPUT_DIR / 'gguf'}")
print("=" * 60)