from __future__ import annotations

import torch
from torch.utils.data import DataLoader

from src.losses import compute_loss_for_phase
from src.training_data import move_batch_to_device

def evaluate_validation_loss(
    phase: str,
    model,
    dataloader: DataLoader,
    device: torch.device,
    alpha: float,
    temperature: float,
    online_kd_token_chunk_size: int = 2048,
    teacher_model=None,
    max_batches: int = -1,
) -> dict[str, float | int]:
    was_training = model.training
    model.eval()

    total_loss = 0.0
    total_ce = 0.0
    total_kd = 0.0
    batches = 0

    with torch.inference_mode():
        for batch in dataloader:
            if max_batches > 0 and batches >= max_batches:
                break
            batch = move_batch_to_device(batch, device)
            input_ids = batch["input_ids"]
            attention_mask = batch["attention_mask"]
            labels = batch["labels"]
            loss_mask = batch["loss_mask"]

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

            if phase == "online_kd" and teacher_model is not None:
                teacher_logits = teacher_model(input_ids=input_ids, attention_mask=attention_mask).logits
            else:
                teacher_logits = None

            loss, ce, kd = compute_loss_for_phase(
                phase,
                logits,
                labels,
                loss_mask,
                batch,
                alpha,
                temperature,
                teacher_logits=teacher_logits,
                online_kd_token_chunk_size=online_kd_token_chunk_size,
            )
            total_loss += float(loss.detach().item())
            total_ce += float(ce.detach().item())
            total_kd += float(kd.detach().item())
            batches += 1

    if was_training:
        model.train()

    denom = max(batches, 1)
    return {
        "loss": total_loss / denom,
        "ce": total_ce / denom,
        "kd": total_kd / denom,
        "batches": batches,
    }