import torch
from dataclasses import dataclass
from torch.utils.data import Dataset, DataLoader
from transformers import AutoModelForCausalLM, AutoTokenizer


class TextDataset(Dataset):
    def __init__(self, texts: list[str], tokenizer, max_length: int = 512):
        self.texts = texts
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self) -> int:
        return len(self.texts)

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        enc = self.tokenizer(
            self.texts[idx],
            truncation=True,
            max_length=self.max_length,
            padding="max_length",
            return_tensors="pt",
        )
        item = {k: v.squeeze(0) for k, v in enc.items()}
        item["labels"] = item["input_ids"].clone()
        item["labels"][item["attention_mask"] == 0] = -100
        return item


def causal_lm_loss(model, batch) -> torch.Tensor:
    outputs = model(**batch)
    return outputs.loss


def kl_retain_loss(current_model, reference_model, batch, temperature: float = 1.0) -> torch.Tensor:
    with torch.no_grad():
        ref_logits = reference_model(
            input_ids=batch["input_ids"],
            attention_mask=batch["attention_mask"],
        ).logits / temperature

    cur_logits = current_model(
        input_ids=batch["input_ids"],
        attention_mask=batch["attention_mask"],
    ).logits / temperature

    ref_probs = torch.softmax(ref_logits, dim=-1)
    cur_log_probs = torch.log_softmax(cur_logits, dim=-1)
    ref_log_probs = torch.log_softmax(ref_logits, dim=-1)

    mask = batch["attention_mask"].unsqueeze(-1)
    kl = ref_probs * (ref_log_probs - cur_log_probs)
    return (kl * mask).sum() / mask.sum().clamp(min=1)


@dataclass
class UnlearningConfig:
    lambda_retain: float = 1.0
    lr: float = 1e-5
    epochs: int = 3
    use_kl_retain: bool = True


def unlearn_gradient_diff(
    model,
    reference_model,
    forget_loader: DataLoader,
    retain_loader: DataLoader,
    cfg: UnlearningConfig,
    device: str = "cuda",
):
    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr)
    model.train()
    reference_model.eval()

    for epoch in range(cfg.epochs):
        retain_iter = iter(retain_loader)
        for forget_batch in forget_loader:
            try:
                retain_batch = next(retain_iter)
            except StopIteration:
                retain_iter = iter(retain_loader)
                retain_batch = next(retain_iter)

            forget_batch = {k: v.to(device) for k, v in forget_batch.items()}
            retain_batch = {k: v.to(device) for k, v in retain_batch.items()}

            forget_loss = causal_lm_loss(model, forget_batch)
            if cfg.use_kl_retain:
                retain_loss = kl_retain_loss(model, reference_model, retain_batch)
            else:
                retain_loss = causal_lm_loss(model, retain_batch)

            loss = -forget_loss + cfg.lambda_retain * retain_loss
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()

        print(
            f"epoch={epoch} forget_loss={forget_loss.item():.4f} "
            f"retain_loss={retain_loss.item():.4f}"
        )

    return model


def lcs_len(a: list[str], b: list[str]) -> int:
    dp = [[0] * (len(b) + 1) for _ in range(len(a) + 1)]
    for i, x in enumerate(a, 1):
        for j, y in enumerate(b, 1):
            dp[i][j] = dp[i - 1][j - 1] + 1 if x == y else max(dp[i - 1][j], dp[i][j - 1])
    return dp[-1][-1]


def rouge_l(prediction: str, reference: str) -> float:
    pred_tokens = prediction.split()
    ref_tokens = reference.split()
    if not pred_tokens or not ref_tokens:
        return 0.0
    lcs = lcs_len(pred_tokens, ref_tokens)
    precision = lcs / len(pred_tokens)
    recall = lcs / len(ref_tokens)
    if precision + recall == 0:
        return 0.0
    return 2 * precision * recall / (precision + recall)


@torch.no_grad()
def extraction_eval(model, tokenizer, prompts: list[str], references: list[str], device: str = "cuda") -> dict:
    model.eval()
    scores = []
    for prompt, ref in zip(prompts, references):
        inputs = tokenizer(prompt, return_tensors="pt").to(device)
        output_ids = model.generate(**inputs, max_new_tokens=128, do_sample=False)
        text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
        scores.append(rouge_l(text, ref))
    return {
        "mean_rouge_l": sum(scores) / len(scores),
        "max_rouge_l": max(scores),
        "n": len(scores),
    }