"""
Fine-tune CodeT5 on a buggy-code -> fixed-code dataset.

Dataset: CodeXGLUE "code-refinement" (Java, small subset) — pairs of
(buggy_function, fixed_function). This is the standard benchmark for this
exact task, so it's a strong, defensible choice for a project report.

Swap `DATASET_NAME` for "google/code_x_glue_cc_code_refinement" (small) or
use the "Bugs2Fix" dataset if you want Python-only pairs instead.

Run:
    pip install -r requirements.txt
    python train.py
"""

from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    DataCollatorForSeq2Seq,
)

MODEL_NAME = "Salesforce/codet5-base"
DATASET_NAME = "google/code_x_glue_cc_code_refinement"
DATASET_CONFIG = "small"  # or "medium"
OUTPUT_DIR = "./checkpoints/codet5-bugfix-finetuned"

MAX_INPUT_LEN = 256
MAX_TARGET_LEN = 256


def main():
    print("Loading dataset...")
    raw_dataset = load_dataset(DATASET_NAME, DATASET_CONFIG)

    print("Loading tokenizer & model...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

    def preprocess(examples):
        inputs = ["fix bug: " + code for code in examples["buggy"]]
        targets = examples["fixed"]
        model_inputs = tokenizer(
            inputs, max_length=MAX_INPUT_LEN, truncation=True, padding="max_length"
        )
        labels = tokenizer(
            targets, max_length=MAX_TARGET_LEN, truncation=True, padding="max_length"
        )
        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

    print("Tokenizing...")
    tokenized_dataset = raw_dataset.map(
        preprocess, batched=True, remove_columns=raw_dataset["train"].column_names
    )

    data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)

    training_args = Seq2SeqTrainingArguments(
        output_dir=OUTPUT_DIR,
        eval_strategy="epoch",
        save_strategy="epoch",
        learning_rate=3e-5,
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        num_train_epochs=5,
        weight_decay=0.01,
        predict_with_generate=True,
        fp16=True,  # requires GPU
        logging_steps=50,
        save_total_limit=2,
        load_best_model_at_end=True,
    )

    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset["train"],
        eval_dataset=tokenized_dataset["validation"],
        data_collator=data_collator,
        tokenizer=tokenizer,
    )

    print("Starting fine-tuning...")
    trainer.train()

    print(f"Saving final model to {OUTPUT_DIR}")
    trainer.save_model(OUTPUT_DIR)
    tokenizer.save_pretrained(OUTPUT_DIR)


if __name__ == "__main__":
    main()