"""Supervised Fine-Tuning (SFT) for Bee using TRL + Accelerate."""

import argparse
import logging
import sys
from pathlib import Path

from datasets import load_dataset
from transformers import AutoTokenizer, TrainingArguments, set_seed
from trl import SFTTrainer, SFTConfig

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from bee.register import register
from bee.config import BeeConfig
from bee.modeling_bee import BeeForCausalLM

register()

logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s")
logger = logging.getLogger("bee.sft")


def get_args():
    parser = argparse.ArgumentParser(description="SFT train Bee")
    parser.add_argument("--model_path", type=str, required=True, help="Path to pretrained Bee checkpoint")
    parser.add_argument("--dataset", type=str, default="tatsu-lab/alpaca", help="HF dataset for SFT")
    parser.add_argument("--output_dir", type=str, required=True)
    parser.add_argument("--max_seq_length", type=int, default=2048)
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=4)
    parser.add_argument("--learning_rate", type=float, default=2e-5)
    parser.add_argument("--num_train_epochs", type=int, default=3)
    parser.add_argument("--warmup_ratio", type=float, default=0.03)
    parser.add_argument("--save_steps", type=int, default=500)
    parser.add_argument("--logging_steps", type=int, default=50)
    parser.add_argument("--bf16", action="store_true", default=True)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--push_to_hub", action="store_true", default=False)
    parser.add_argument("--hub_model_id", type=str, default=None)
    return parser.parse_args()


def formatting_alpaca(examples):
    texts = []
    for instruction, input_text, output in zip(examples["instruction"], examples.get("input", []), examples["output"]):
        if input_text:
            text = f"### Instruction:\n{instruction}\n### Input:\n{input_text}\n### Response:\n{output}"
        else:
            text = f"### Instruction:\n{instruction}\n### Response:\n{output}"
        texts.append(text)
    return {"text": texts}


def main():
    args = get_args()
    set_seed(args.seed)

    logger.info("Loading model from %s", args.model_path)
    model = BeeForCausalLM.from_pretrained(args.model_path)
    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    logger.info("Loading SFT dataset: %s", args.dataset)
    ds = load_dataset(args.dataset, split="train")
    if "alpaca" in args.dataset.lower():
        ds = ds.map(formatting_alpaca, batched=True)

    training_args = SFTConfig(
        output_dir=args.output_dir,
        num_train_epochs=args.num_train_epochs,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.learning_rate,
        warmup_ratio=args.warmup_ratio,
        logging_steps=args.logging_steps,
        save_steps=args.save_steps,
        save_strategy="steps",
        bf16=args.bf16,
        max_seq_length=args.max_seq_length,
        dataset_text_field="text",
        report_to=["tensorboard"],
        push_to_hub=args.push_to_hub,
        hub_model_id=args.hub_model_id,
    )

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=ds,
        args=training_args,
    )

    logger.info("Starting SFT training...")
    trainer.train()
    logger.info("SFT complete. Saving to %s", args.output_dir)
    trainer.save_model(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)


if __name__ == "__main__":
    main()