"""Focused regression tests for causal-LM evaluation and report output."""

from __future__ import annotations

import json
import math
from types import SimpleNamespace
from unittest.mock import MagicMock, Mock

import pytest
import torch
from torch import nn

from obliteratus.evaluation.evaluator import Evaluator


class _MapDataset:
    def __init__(self, texts: list[object]):
        self.texts = texts
        self.read_indices: list[int] = []
        self.selected_indices: list[int] | None = None

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

    def __getitem__(self, key):
        if isinstance(key, str):
            raise TypeError("the evaluator must not materialize a full dataset column")
        if isinstance(key, slice):
            return {"text": self.texts[key]}
        self.read_indices.append(key)
        return {"text": self.texts[key]}

    def select(self, indices):
        self.selected_indices = list(indices)
        return _MapDataset([self.texts[index] for index in self.selected_indices])


class _ChangingDataset(_MapDataset):
    """Simulate a transformed dataset changing after index selection."""

    def select(self, indices):
        self.selected_indices = list(indices)
        return _MapDataset([""])


class _Encoding(dict):
    def to(self, _device):
        return self


class _ScriptedTokenizer:
    def __init__(self, encodings: list[_Encoding]):
        self.encodings = list(encodings)
        self.calls: list[list[str]] = []

    def __call__(self, texts, **_kwargs):
        self.calls.append(list(texts))
        return self.encodings.pop(0)


class _LossModel(nn.Module):
    def __init__(self, losses: list[float | None]):
        super().__init__()
        self.anchor = nn.Parameter(torch.zeros(()))
        self.losses = list(losses)
        self.labels: list[torch.Tensor] = []

    def forward(self, *, input_ids, attention_mask, labels):
        del input_ids, attention_mask
        self.labels.append(labels.detach().clone())
        loss = self.losses.pop(0)
        return SimpleNamespace(loss=None if loss is None else torch.tensor(loss))


def _encoding(input_ids, attention_mask) -> _Encoding:
    return _Encoding(
        input_ids=torch.as_tensor(input_ids, dtype=torch.long),
        attention_mask=torch.as_tensor(attention_mask, dtype=torch.long),
    )


def _evaluator(
    texts: list[object],
    encodings: list[_Encoding],
    losses: list[float | None],
    *,
    batch_size: int = 8,
    max_samples: int | None = None,
):
    dataset = _MapDataset(texts)
    tokenizer = _ScriptedTokenizer(encodings)
    model = _LossModel(losses)
    handle = SimpleNamespace(model=model, tokenizer=tokenizer, task="causal_lm")
    evaluator = Evaluator(
        handle=handle,
        dataset=dataset,
        batch_size=batch_size,
        max_samples=max_samples,
    )
    return evaluator, dataset, tokenizer, model


def test_causal_lm_skips_blank_rows_bounds_scan_and_masks_padding():
    evaluator, dataset, tokenizer, model = _evaluator(
        ["", "  \t", "alpha", "beta", "must not be read"],
        [_encoding([[1, 2, 0], [3, 4, 5]], [[1, 1, 0], [1, 1, 1]])],
        [math.log(4.0)],
        batch_size=2,
        max_samples=2,
    )

    result = evaluator.evaluate()

    assert result["perplexity"] == pytest.approx(4.0)
    assert dataset.read_indices == [0, 1, 2, 3]
    assert dataset.selected_indices == [2, 3]
    assert tokenizer.calls == [["alpha", "beta"]]
    assert model.labels[0].tolist() == [[1, 2, -100], [3, 4, 5]]


def test_causal_lm_rejects_all_empty_dataset():
    evaluator, dataset, tokenizer, model = _evaluator(["", " \n", None], [], [])

    with pytest.raises(ValueError, match="no non-empty text"):
        evaluator.evaluate()

    assert dataset.read_indices == [0, 1, 2]
    assert tokenizer.calls == []
    assert model.labels == []


@pytest.mark.parametrize("max_samples", [0, -1])
def test_causal_lm_requires_positive_max_samples(max_samples):
    evaluator, dataset, _tokenizer, _model = _evaluator(
        ["alpha"],
        [],
        [],
        max_samples=max_samples,
    )

    with pytest.raises(ValueError, match="max_samples must be a positive integer"):
        evaluator.evaluate()

    assert dataset.read_indices == []


def test_causal_lm_rejects_tokenless_and_one_token_sequences():
    evaluator, _dataset, _tokenizer, model = _evaluator(
        ["tokenless", "single"],
        [
            _encoding(torch.empty((1, 0)), torch.empty((1, 0))),
            _encoding([[1]], [[1]]),
        ],
        [],
        batch_size=1,
    )

    with pytest.raises(RuntimeError, match="zero valid prediction tokens"):
        evaluator.evaluate()

    assert model.labels == []


def test_causal_lm_skips_batch_emptied_by_dataset_transform():
    dataset = _ChangingDataset(["alpha"])
    tokenizer = _ScriptedTokenizer([])
    model = _LossModel([])
    handle = SimpleNamespace(model=model, tokenizer=tokenizer, task="causal_lm")

    with pytest.raises(RuntimeError, match="zero valid prediction tokens"):
        Evaluator(handle=handle, dataset=dataset).evaluate()

    assert tokenizer.calls == []
    assert model.labels == []


def test_causal_lm_skips_zero_target_batch_then_reports_valid_perplexity():
    evaluator, _dataset, _tokenizer, model = _evaluator(
        ["padded target", "valid target"],
        [
            _encoding([[1, 0]], [[1, 0]]),
            _encoding([[2, 3]], [[1, 1]]),
        ],
        [math.log(3.0)],
        batch_size=1,
    )

    result = evaluator.evaluate()

    assert result["perplexity"] == pytest.approx(3.0)
    assert len(model.labels) == 1


def test_causal_lm_rejects_missing_loss():
    evaluator, _dataset, _tokenizer, _model = _evaluator(
        ["alpha"],
        [_encoding([[1, 2]], [[1, 1]])],
        [None],
    )

    with pytest.raises(RuntimeError, match="returned no loss"):
        evaluator.evaluate()


@pytest.mark.parametrize("loss", [float("nan"), float("inf"), float("-inf")])
def test_causal_lm_rejects_non_finite_loss(loss):
    evaluator, _dataset, _tokenizer, _model = _evaluator(
        ["alpha"],
        [_encoding([[1, 2]], [[1, 1]])],
        [loss],
    )

    with pytest.raises(RuntimeError, match="Non-finite evaluation loss"):
        evaluator.evaluate()


def test_causal_lm_perplexity_is_weighted_by_prediction_tokens():
    evaluator, _dataset, _tokenizer, _model = _evaluator(
        ["three tokens", "two tokens"],
        [
            _encoding([[1, 2, 3]], [[1, 1, 1]]),
            _encoding([[4, 5]], [[1, 1]]),
        ],
        [math.log(2.0), math.log(8.0)],
        batch_size=1,
    )

    result = evaluator.evaluate()

    expected = math.exp((2 * math.log(2.0) + math.log(8.0)) / 3)
    assert result["perplexity"] == pytest.approx(expected)


def test_report_creates_nested_output_directory(monkeypatch, tmp_path):
    import obliteratus.reporting.report
    from obliteratus import cli

    results_path = tmp_path / "results.json"
    results_path.write_text(
        json.dumps(
            {
                "model_name": "fixture",
                "baseline_metrics": {"perplexity": 4.0},
                "results": [],
            }
        )
    )
    report = MagicMock()
    monkeypatch.setattr(
        obliteratus.reporting.report,
        "AblationReport",
        Mock(return_value=report),
    )
    output_dir = tmp_path / "nested" / "plots"

    cli._cmd_report(
        SimpleNamespace(results_json=str(results_path), output_dir=str(output_dir))
    )

    assert output_dir.is_dir()
    report.plot_impact.assert_called_once_with(
        metric="perplexity",
        output_path=output_dir / "impact.png",
    )
    report.plot_heatmap.assert_called_once_with(output_path=output_dir / "heatmap.png")