"""
tests/test_classifier.py
────────────────────────────────────────────────────────────────
Unit tests for src/nlp/classifier.py

The HuggingFace model and tokenizer are mocked throughout so
the test suite runs without GPU or internet access.
────────────────────────────────────────────────────────────────
"""

from __future__ import annotations

from unittest.mock import MagicMock

import pytest

from src.nlp.classifier import ClassificationResult, ClinicalClassifier


class TestClassificationResult:
    """ClassificationResult serialisation."""

    def test_to_dict_rounds_values(self):
        result = ClassificationResult(
            label         = "urgent",
            confidence    = 0.845678,
            probabilities = {"routine": 0.1, "urgent": 0.845678, "critical": 0.05},
            task          = "severity",
        )
        d = result.to_dict()
        assert d["label"]      == "urgent"
        assert d["confidence"] == 0.846
        assert d["task"]       == "severity"
        assert all(isinstance(v, float) for v in d["probabilities"].values())

    def test_repr(self):
        result = ClassificationResult("routine", 0.9, {}, "severity")
        assert "routine" in repr(result)
        assert "0.900" in repr(result)


class TestClinicalClassifier:
    """ClinicalClassifier with mocked transformers."""

    def _make_classifier(self, task="severity"):
        clf = ClinicalClassifier(task=task)
        return clf

    def test_init_sets_labels_from_config(self):
        clf = self._make_classifier(task="severity")
        assert clf._labels == ["routine", "urgent", "critical"]
        assert clf._label2id["routine"] == 0
        assert clf._id2label[2]         == "critical"

    def test_init_readmission_task(self):
        clf = self._make_classifier(task="readmission")
        assert clf._labels == ["not_readmitted", "readmitted"]

    def test_unknown_task_raises_on_labels(self):
        from src.utils.config import ClassifierConfig
        original = ClassifierConfig.task
        ClassifierConfig.task = "unknown_task"
        try:
            with pytest.raises(ValueError, match="Unknown classifier task"):
                ClassifierConfig.active_labels()
        finally:
            ClassifierConfig.task = original

    def test_predict_raises_if_not_loaded(self):
        clf = self._make_classifier()
        with pytest.raises(RuntimeError, match="not loaded"):
            clf.predict("Some clinical text.")

    def test_predict_returns_classification_result(self):
        """predict() with mocked model returns a ClassificationResult."""
        import torch

        clf = self._make_classifier(task="severity")

        # Mock tokenizer
        mock_tokenizer = MagicMock()
        mock_encoding  = {
            "input_ids":      torch.zeros(1, 512, dtype=torch.long),
            "attention_mask": torch.ones(1, 512, dtype=torch.long),
        }
        mock_tokenizer.return_value = mock_encoding
        clf._tokenizer = mock_tokenizer

        # Mock model — returns logits [critical, urgent, routine] scores
        mock_model  = MagicMock()
        mock_output = MagicMock()
        # logits: critical has highest score
        mock_output.logits = torch.tensor([[0.1, 0.2, 3.5]])
        mock_model.return_value = mock_output
        mock_model.parameters   = MagicMock(
            return_value=iter([torch.zeros(1)])
        )
        clf._model = mock_model

        result = clf.predict("Patient in ICU following cardiac arrest.")
        assert isinstance(result, ClassificationResult)
        assert result.task == "severity"
        assert result.label in ["routine", "urgent", "critical"]
        assert 0.0 <= result.confidence <= 1.0

    def test_predict_batch_empty_list(self):
        clf = self._make_classifier()
        assert clf.predict_batch([]) == []

    def test_save_raises_if_not_loaded(self):
        clf = self._make_classifier()
        with pytest.raises(RuntimeError, match="No model to save"):
            clf.save()