"""Tests for LongMemEval question type mapping (#1278)."""
from __future__ import annotations

import pytest

from cortex.eval.longmemeval import (
    QUESTION_TYPE_MAP,
    ABSTENTION_TYPES,
    LongMemEvalLoader,
    LongMemEvalQuestion,
)


class TestQuestionTypeMapping:
    def test_native_types_preserved(self):
        """Native LongMemEval types should map to themselves."""
        for t in ["knowledge_update", "knowledge_retention", "single_session",
                   "cross_session", "temporal_reasoning", "abstention"]:
            assert QUESTION_TYPE_MAP[t] == t

    def test_aliases_mapped(self):
        """Common aliases should map to the correct canonical type."""
        assert QUESTION_TYPE_MAP["temporal"] == "temporal_reasoning"
        assert QUESTION_TYPE_MAP["multi_session"] == "cross_session"
        assert QUESTION_TYPE_MAP["update"] == "knowledge_update"
        assert QUESTION_TYPE_MAP["refuse"] == "abstention"

    def test_loader_maps_types(self):
        """Loader should use QUESTION_TYPE_MAP for question types."""
        loader = LongMemEvalLoader()
        dataset = loader.load_dict({
            "sessions": [{
                "session_id": "s1",
                "turns": [{"role": "user", "content": "hello"}],
                "questions": [
                    {"question_id": "q1", "question": "test", "answer": "a",
                     "question_type": "temporal"},
                    {"question_id": "q2", "question": "test", "answer": "a",
                     "question_type": "abstention"},
                    {"question_id": "q3", "question": "test", "answer": "a",
                     "type": "multi_session"},
                ],
            }],
        })

        qs = dataset.sessions[0].questions
        assert qs[0].question_type == "temporal_reasoning"
        assert qs[1].question_type == "abstention"
        assert qs[2].question_type == "cross_session"

    def test_unknown_types_passed_through(self):
        """Unknown types should be passed through unchanged."""
        loader = LongMemEvalLoader()
        dataset = loader.load_dict({
            "sessions": [{
                "session_id": "s1",
                "turns": [{"role": "user", "content": "hello"}],
                "questions": [
                    {"question_id": "q1", "question": "test", "answer": "a",
                     "question_type": "custom_type"},
                ],
            }],
        })
        assert dataset.sessions[0].questions[0].question_type == "custom_type"


class TestAbstention:
    def test_is_abstention_property(self):
        q = LongMemEvalQuestion(
            question_id="q1", question="test", answer="",
            question_type="abstention",
        )
        assert q.is_abstention is True

    def test_non_abstention(self):
        q = LongMemEvalQuestion(
            question_id="q1", question="test", answer="x",
            question_type="single_session",
        )
        assert q.is_abstention is False

    def test_abstention_types(self):
        assert "abstention" in ABSTENTION_TYPES
