"""EvalHook and HuggingFace ForgettingCallback."""
from __future__ import annotations

import math
from typing import TYPE_CHECKING, Any, Optional

from llm_refresh.models import EvalResult

try:
    import torch
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False

try:
    from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments
    HAS_TRANSFORMERS = True
except ImportError:
    HAS_TRANSFORMERS = False
    # Stubs so type hints work at runtime
    TrainerCallback = object  # type: ignore[misc,assignment]
    TrainerControl = Any  # type: ignore[misc,assignment]
    TrainerState = Any  # type: ignore[misc,assignment]
    TrainingArguments = Any  # type: ignore[misc,assignment]

if TYPE_CHECKING:
    from llm_refresh.metrics import ForgettingTracker


def compute_perplexity(
    model: Any,
    tokenizer: Any,
    dataset: list[dict[str, Any]],
    batch_size: int = 8,
    max_length: int = 512,
) -> EvalResult:
    """
    Compute perplexity on `dataset` (list of {"text": "..."} dicts).
    Requires torch + transformers.
    """
    if not HAS_TORCH or not HAS_TRANSFORMERS:
        raise RuntimeError(
            "torch and transformers are required for perplexity eval. "
            "Install with: pip install llm-refresh-wheel[torch]"
        )

    import torch

    model.eval()
    total_loss = 0.0
    total_tokens = 0
    device = next(model.parameters()).device

    texts = [ex["text"] for ex in dataset if "text" in ex]
    if not texts:
        raise ValueError("Dataset contains no 'text' fields.")

    for i in range(0, len(texts), batch_size):
        batch_texts = texts[i : i + batch_size]
        enc = tokenizer(
            batch_texts,
            return_tensors="pt",
            truncation=True,
            max_length=max_length,
            padding=True,
        )
        input_ids = enc["input_ids"].to(device)
        attention_mask = enc["attention_mask"].to(device)

        with torch.no_grad():
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=input_ids,
            )
        # outputs.loss is mean loss over non-padding tokens
        n_tokens = attention_mask.sum().item()
        total_loss += outputs.loss.item() * n_tokens
        total_tokens += n_tokens

    avg_loss = total_loss / max(total_tokens, 1)
    ppl = math.exp(avg_loss)
    return EvalResult.now(perplexity=ppl, loss=avg_loss, dataset_size=len(texts))


class EvalHook:
    """Callable that evaluates perplexity on a fixed anchor dataset."""

    def __init__(
        self,
        model: Any,
        tokenizer: Any,
        dataset: list[dict[str, Any]],
        batch_size: int = 8,
        label: str = "anchor",
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.dataset = dataset
        self.batch_size = batch_size
        self.label = label

    def __call__(self) -> EvalResult:
        result = compute_perplexity(
            self.model, self.tokenizer, self.dataset, self.batch_size
        )
        result.label = self.label
        return result


if HAS_TRANSFORMERS:

    class ForgettingCallback(TrainerCallback):  # type: ignore[misc]
        """
        HuggingFace Trainer callback that measures forgetting each cycle.

        - on_train_begin: eval on anchor set (pre), record to tracker
        - on_train_end: eval on anchor set (post), compute BWT, optionally rollback
        """

        def __init__(
            self,
            tracker: "ForgettingTracker",
            eval_hook: EvalHook,
            checkpoint_path: Optional[str] = None,
            auto_rollback: bool = True,
            forgetting_threshold: float = -0.1,
        ):
            self.tracker = tracker
            self.eval_hook = eval_hook
            self.checkpoint_path = checkpoint_path
            self.auto_rollback = auto_rollback
            self.forgetting_threshold = forgetting_threshold
            self._pre_result: Optional[EvalResult] = None

        def on_train_begin(
            self,
            args: TrainingArguments,
            state: TrainerState,
            control: TrainerControl,
            **kwargs: Any,
        ) -> None:
            self._pre_result = self.eval_hook()
            self.tracker.record_pre_eval(self._pre_result)

        def on_train_end(
            self,
            args: TrainingArguments,
            state: TrainerState,
            control: TrainerControl,
            model: Any = None,
            **kwargs: Any,
        ) -> None:
            post_result = self.eval_hook()
            self.tracker.record_post_eval(post_result)

            pre = self._pre_result
            post = post_result
            if pre is not None:
                bwt = pre.perplexity - post.perplexity
                if bwt < self.forgetting_threshold and self.auto_rollback and self.checkpoint_path:
                    # Reload weights from checkpoint
                    if model is not None and HAS_TRANSFORMERS:
                        try:
                            from transformers import AutoModelForCausalLM
                            loaded = AutoModelForCausalLM.from_pretrained(self.checkpoint_path)
                            model.load_state_dict(loaded.state_dict())
                        except Exception:
                            pass

else:
    # Stub when transformers not installed
    class ForgettingCallback:  # type: ignore[no-redef]
        def __init__(self, *args: Any, **kwargs: Any) -> None:
            raise RuntimeError(
                "transformers is required for ForgettingCallback. "
                "Install with: pip install llm-refresh-wheel[transformers]"
            )
