import argparse
import os
from dataclasses import dataclass

import pandas as pd
import torch
from sklearn.metrics import accuracy_score, f1_score
from torch.utils.data import DataLoader, Dataset
from torch.utils.tensorboard import SummaryWriter
from transformers import (
    RobertaForSequenceClassification,
    RobertaTokenizer,
    Trainer,
    TrainingArguments,
)
from transformers.integrations import TensorBoardCallback

MODEL_DIR = "roberta_prompt_extraction_model"
LOG_DIR = "prompt_extraction_logs"

if not os.path.exists(MODEL_DIR):
    os.makedirs(MODEL_DIR, exist_ok=True)
if not os.path.exists(LOG_DIR):
    os.makedirs(LOG_DIR, exist_ok=True)


class AttackDataset(Dataset):
    def __init__(self, dataframe, tokenizer, max_length=128):
        self.data = dataframe
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        text = (
            self.data.iloc[index]["access_code"]
            + "[SEP]"
            + self.data.iloc[index]["llm_output"]
            + "[SEP]"
        )
        label = torch.tensor(
            self.data.iloc[index]["is_prompt_extraction"], dtype=torch.long
        )

        encoding = self.tokenizer(
            text,
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
            return_tensors="pt",
        )

        return {
            "input_ids": encoding["input_ids"].squeeze(),
            "attention_mask": encoding["attention_mask"].squeeze(),
            "labels": label,
        }


@dataclass
class Metrics:
    def __init__(self, f1, accuracy):
        self.f1 = f1
        self.accuracy = accuracy

    def __str__(self):
        return f"F1: {self.f1:.4f}, Accuracy: {self.accuracy:.4f}"


def evaluate(model, val_loader):
    model.eval()
    y_true = []
    y_pred = []
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    with torch.no_grad():
        for batch in val_loader:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels = batch["labels"].to(device)

            outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
            logits = outputs.logits

            _, predicted = torch.max(logits, dim=1)
            y_true.extend(labels.cpu().numpy())
            y_pred.extend(predicted.cpu().numpy())

    f1 = f1_score(y_true, y_pred)
    accuracy = accuracy_score(y_true, y_pred)
    return Metrics(f1, accuracy)


def train_and_evaluate(retrain=False):
    device = (
        "mps"
        if torch.backends.mps.is_available()
        else "cuda" if torch.cuda.is_available() else "cpu"
    )
    device = torch.device(device)

    # train and evaluate model
    if retrain:
        tokenizer = RobertaTokenizer.from_pretrained("roberta-large")
        model = RobertaForSequenceClassification.from_pretrained(
            "roberta-large", num_labels=2
        )
        model.to(device)

        train_df = pd.read_parquet(
            "../outdata/prompt_extraction_detection_training.parquet"
        )
        test_df = pd.read_parquet(
            "../outdata/prompt_extraction_detection_validation.parquet"
        )

        train_dataset = AttackDataset(train_df, tokenizer)
        val_dataset = AttackDataset(test_df, tokenizer)

        val_loader = DataLoader(val_dataset, batch_size=16)
        training_args = TrainingArguments(
            output_dir=MODEL_DIR,
            save_total_limit=1,
            learning_rate=2e-5,
            weight_decay=0.01,
            logging_strategy="steps",
            logging_steps=10,
        )

        tb_callback = TensorBoardCallback(SummaryWriter(log_dir=LOG_DIR))

        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=train_dataset,
            eval_dataset=val_dataset,
            callbacks=[tb_callback],
        )

        trainer.train()
        trainer.save_model()
    else:
        # you should set this to a directory if you're executing code in your home directory
        model = RobertaForSequenceClassification.from_pretrained(
            "prompt_extraction_model"
        ).to(device)
        tokenizer = RobertaTokenizer.from_pretrained("roberta-large")

    # Evaluate the model on the validation set
    test_df = pd.read_parquet(
        "../outdata/prompt_extraction_detection_validation.parquet"
    )
    val_dataset = AttackDataset(test_df, tokenizer)
    val_loader = DataLoader(val_dataset, batch_size=16)

    metrics = evaluate(model, val_loader)
    print(metrics)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="RoBERTa model for prompt extraction detection"
    )
    parser.add_argument(
        "--retrain",
        action="store_const",
        const=True,
        default=False,
        help="Retrain model",
    )
    args = parser.parse_args()
    train_and_evaluate(retrain=args.retrain)
