import inspect
import json
import re
from pathlib import Path

import numpy as np
from datasets import load_dataset, load_from_disk
from transformers import AutoTokenizer, Trainer, TrainingArguments


LABELS = {
    "id2label": {0: "negative", 1: "positive"},
    "label2id": {"negative": 0, "positive": 1},
}


def make_training_arguments(**kwargs):
    params = inspect.signature(TrainingArguments.__init__).parameters
    if "eval_strategy" in params and "evaluation_strategy" in kwargs:
        kwargs["eval_strategy"] = kwargs.pop("evaluation_strategy")
    if "evaluation_strategy" in params and "eval_strategy" in kwargs:
        kwargs["evaluation_strategy"] = kwargs.pop("eval_strategy")
    return TrainingArguments(**kwargs)


def make_trainer(tokenizer, **kwargs):
    try:
        return Trainer(processing_class=tokenizer, **kwargs)
    except TypeError:
        return Trainer(tokenizer=tokenizer, **kwargs)


def compute_binary_metrics(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)

    tp = int(((preds == 1) & (labels == 1)).sum())
    fp = int(((preds == 1) & (labels == 0)).sum())
    fn = int(((preds == 0) & (labels == 1)).sum())
    accuracy = float((preds == labels).mean())
    precision = tp / (tp + fp) if tp + fp else 0.0
    recall = tp / (tp + fn) if tp + fn else 0.0
    f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0

    return {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1": f1,
    }


def load_or_prepare_imdb(
    checkpoint="bert-base-uncased",
    data_dir="./data/imdb_tokenized",
    max_length=512,
    padding="max_length",
    max_train_samples=None,
    max_eval_samples=None,
):
    data_dir = Path(data_dir)
    if max_train_samples is not None or max_eval_samples is not None:
        train_tag = "all" if max_train_samples is None else str(max_train_samples)
        eval_tag = "all" if max_eval_samples is None else str(max_eval_samples)
        data_dir = data_dir / f"sample_train_{train_tag}_eval_{eval_tag}"
    train_path = data_dir / "train"
    eval_path = data_dir / "test"
    tokenizer_path = data_dir / "tokenizer"

    if train_path.exists() and eval_path.exists() and tokenizer_path.exists():
        return (
            load_from_disk(str(train_path)),
            load_from_disk(str(eval_path)),
            AutoTokenizer.from_pretrained(str(tokenizer_path)),
        )

    raw = load_dataset("imdb")
    train_ds = raw["train"]
    eval_ds = raw["test"]

    if max_train_samples is not None:
        train_ds = train_ds.select(range(max_train_samples))
    if max_eval_samples is not None:
        eval_ds = eval_ds.select(range(max_eval_samples))

    tokenizer = AutoTokenizer.from_pretrained(checkpoint)

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

    train_ds = train_ds.map(preprocess, batched=True)
    eval_ds = eval_ds.map(preprocess, batched=True)
    train_ds = train_ds.remove_columns(["text"]).rename_column("label", "labels")
    eval_ds = eval_ds.remove_columns(["text"]).rename_column("label", "labels")

    data_dir.mkdir(parents=True, exist_ok=True)
    train_ds.save_to_disk(str(train_path))
    eval_ds.save_to_disk(str(eval_path))
    tokenizer.save_pretrained(str(tokenizer_path))
    return train_ds, eval_ds, tokenizer


def get_best_checkpoint(output_dir):
    state_path = Path(output_dir) / "trainer_state.json"
    if not state_path.exists():
        return None, None

    state = json.loads(state_path.read_text(encoding="utf-8"))
    best_checkpoint = state.get("best_model_checkpoint")
    if best_checkpoint is None:
        return None, None

    match = re.search(r"checkpoint-(\d+)", best_checkpoint)
    best_step = int(match.group(1)) if match else None
    return best_checkpoint, best_step


def find_backbone(model):
    if hasattr(model, "backbone"):
        return model.backbone
    if hasattr(model, "base_model"):
        return model.base_model
    raise AttributeError("backbone or base_model was not found")


def set_backbone_trainable(model, trainable):
    backbone = find_backbone(model)
    for param in backbone.parameters():
        param.requires_grad = trainable