#!/usr/bin/env python3
"""
LoRA fine-tuning for editorai:jtafull using Unsloth + TRL SFTTrainer.
Base model: Qwen/Qwen3-8B (fallback: Qwen/Qwen2.5-7B-Instruct)
"""
import json
import torch
from datasets import Dataset
from trl import SFTTrainer, SFTConfig
from unsloth import FastLanguageModel
from unsloth.chat_templates import get_chat_template

BASE_MODEL = "Qwen/Qwen3-8B"
MAX_SEQ_LENGTH = 4096
LORA_R = 16
LORA_ALPHA = 32
LORA_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
TRAIN_DATA = "data/train.jsonl"
OUTPUT_ADAPTER = "./lora-adapter"
OUTPUT_MERGED = "./merged-model"


def load_dataset(path: str) -> Dataset:
    examples = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line:
                examples.append(json.loads(line))
    print(f"Loaded {len(examples)} training examples")
    return Dataset.from_list(examples)


def format_example(example, tokenizer):
    return tokenizer.apply_chat_template(
        example["messages"],
        tokenize=False,
        add_generation_prompt=False,
    )


def main():
    print(f"Loading base model: {BASE_MODEL}")
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=BASE_MODEL,
        max_seq_length=MAX_SEQ_LENGTH,
        dtype=None,
        load_in_4bit=True,
    )

    tokenizer = get_chat_template(tokenizer, chat_template="chatml")

    model = FastLanguageModel.get_peft_model(
        model,
        r=LORA_R,
        lora_alpha=LORA_ALPHA,
        target_modules=LORA_TARGETS,
        lora_dropout=0.05,
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=42,
    )

    dataset = load_dataset(TRAIN_DATA)
    dataset = dataset.map(
        lambda ex: {"text": format_example(ex, tokenizer)},
        remove_columns=dataset.column_names,
    )

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=dataset,
        args=SFTConfig(
            dataset_text_field="text",
            max_seq_length=MAX_SEQ_LENGTH,
            per_device_train_batch_size=4,
            gradient_accumulation_steps=4,
            num_train_epochs=3,
            learning_rate=2e-4,
            lr_scheduler_type="cosine",
            warmup_ratio=0.05,
            fp16=not torch.cuda.is_bf16_supported(),
            bf16=torch.cuda.is_bf16_supported(),
            logging_steps=10,
            save_steps=100,
            output_dir=OUTPUT_ADAPTER,
            report_to="none",
        ),
    )

    print("Starting training...")
    trainer.train()

    print(f"Saving LoRA adapter to {OUTPUT_ADAPTER}")
    model.save_pretrained(OUTPUT_ADAPTER)
    tokenizer.save_pretrained(OUTPUT_ADAPTER)

    print(f"Merging LoRA into full model → {OUTPUT_MERGED}")
    model.save_pretrained_merged(OUTPUT_MERGED, tokenizer, save_method="merged_16bit")

    print("Done. Run convert_and_push.sh to build GGUF and push to Ollama.")


if __name__ == "__main__":
    main()