"""
tests/test_eval_loaders.py — Unit tests for LoCoMo and LongMemEval loaders.

Tests:
  - LoCoMoLoader: load sample, load_dict, data access
  - LongMemEvalLoader: load sample, load_dict, data access
"""
from __future__ import annotations

import pytest
from cortex.eval.locomo import LoCoMoLoader, LoCoMoDataset, LoCoMoSession, LoCoMoQAPair
from cortex.eval.longmemeval import LongMemEvalLoader, LongMemEvalDataset, LongMemEvalSession, LongMemEvalQuestion


# ---------------------------------------------------------------------------
# LoCoMo
# ---------------------------------------------------------------------------

class TestLoCoMoLoader:
    def test_load_sample(self):
        loader = LoCoMoLoader()
        dataset = loader.load()  # bundled sample
        assert isinstance(dataset, LoCoMoDataset)
        assert len(dataset.sessions) > 0

    def test_sample_has_five_sessions(self):
        dataset = LoCoMoLoader().load()
        assert len(dataset.sessions) == 5

    def test_session_structure(self):
        dataset = LoCoMoLoader().load()
        s = dataset.sessions[0]
        assert isinstance(s, LoCoMoSession)
        assert s.session_id
        assert len(s.conversation) > 0
        assert len(s.qa_pairs) > 0

    def test_qa_pair_structure(self):
        dataset = LoCoMoLoader().load()
        qa = dataset.sessions[0].qa_pairs[0]
        assert isinstance(qa, LoCoMoQAPair)
        assert qa.question
        assert qa.answer
        assert qa.question_type

    def test_conversation_text(self):
        dataset = LoCoMoLoader().load()
        text = dataset.sessions[0].conversation_text
        assert "User:" in text or "user" in text.lower()

    def test_messages_property(self):
        dataset = LoCoMoLoader().load()
        msgs = dataset.sessions[0].messages
        assert isinstance(msgs, list)
        for m in msgs:
            assert "role" in m
            assert "content" in m

    def test_all_qa_pairs(self):
        dataset = LoCoMoLoader().load()
        all_qa = dataset.all_qa_pairs
        # 5 sessions × 4 pairs each = 20
        assert len(all_qa) == 20

    def test_total_questions(self):
        dataset = LoCoMoLoader().load()
        assert dataset.total_questions == 20

    def test_load_dict_bare_list(self):
        raw = [
            {
                "session_id": "test_s1",
                "conversation": [{"role": "user", "content": "My name is Bob."}],
                "qa_pairs": [
                    {"question_id": "q1", "question": "What is my name?", "answer": "Bob"}
                ],
            }
        ]
        dataset = LoCoMoLoader().load_dict(raw)
        assert len(dataset.sessions) == 1
        assert dataset.sessions[0].session_id == "test_s1"
        assert dataset.sessions[0].qa_pairs[0].answer == "Bob"

    def test_load_dict_envelope(self):
        raw = {
            "dataset": "test_locomo",
            "sessions": [
                {
                    "session_id": "s1",
                    "conversation": [{"role": "user", "content": "Hello"}],
                    "qa_pairs": [],
                }
            ],
        }
        dataset = LoCoMoLoader().load_dict(raw)
        assert dataset.name == "test_locomo"
        assert len(dataset.sessions) == 1

    def test_load_nonexistent_file(self):
        with pytest.raises(FileNotFoundError):
            LoCoMoLoader().load("/tmp/nonexistent_locomo_dataset.json")

    def test_questions_by_type(self):
        dataset = LoCoMoLoader().load()
        single_hop = dataset.questions_by_type("single_hop")
        assert len(single_hop) > 0
        multi_hop = dataset.questions_by_type("multi_hop")
        # Sample only has single_hop
        assert len(multi_hop) == 0


# ---------------------------------------------------------------------------
# LongMemEval
# ---------------------------------------------------------------------------

class TestLongMemEvalLoader:
    def test_load_sample(self):
        loader = LongMemEvalLoader()
        dataset = loader.load()
        assert isinstance(dataset, LongMemEvalDataset)
        assert len(dataset.sessions) > 0

    def test_sample_has_five_sessions(self):
        dataset = LongMemEvalLoader().load()
        assert len(dataset.sessions) == 5

    def test_session_structure(self):
        dataset = LongMemEvalLoader().load()
        s = dataset.sessions[0]
        assert isinstance(s, LongMemEvalSession)
        assert s.session_id
        assert len(s.turns) > 0
        assert len(s.questions) > 0

    def test_question_structure(self):
        dataset = LongMemEvalLoader().load()
        q = dataset.sessions[0].questions[0]
        assert isinstance(q, LongMemEvalQuestion)
        assert q.question
        assert q.answer
        assert q.question_type

    def test_conversation_property(self):
        dataset = LongMemEvalLoader().load()
        conv = dataset.sessions[0].conversation
        assert isinstance(conv, list)
        for m in conv:
            assert "role" in m
            assert "content" in m

    def test_conversation_text(self):
        dataset = LongMemEvalLoader().load()
        text = dataset.sessions[0].conversation_text
        assert len(text) > 0

    def test_turns_by_ids(self):
        dataset = LongMemEvalLoader().load()
        session = dataset.sessions[0]
        turns = session.turns_by_ids([1])
        assert len(turns) == 1
        assert turns[0].turn_id == 1

    def test_all_questions(self):
        dataset = LongMemEvalLoader().load()
        all_q = dataset.all_questions
        assert len(all_q) > 0

    def test_total_questions(self):
        dataset = LongMemEvalLoader().load()
        # sample has varying q counts — just verify it's positive
        assert dataset.total_questions > 0

    def test_get_session(self):
        dataset = LongMemEvalLoader().load()
        session_id = dataset.sessions[0].session_id
        s = dataset.get_session(session_id)
        assert s is not None
        assert s.session_id == session_id

    def test_get_session_missing(self):
        dataset = LongMemEvalLoader().load()
        assert dataset.get_session("nonexistent") is None

    def test_load_dict_bare_list(self):
        raw = [
            {
                "session_id": "test_s1",
                "turns": [
                    {"turn_id": 1, "role": "user", "content": "I live in Tokyo."}
                ],
                "questions": [
                    {
                        "question_id": "q1",
                        "question": "Where does the user live?",
                        "answer": "Tokyo",
                        "evidence_turn_ids": [1],
                        "question_type": "single_session",
                    }
                ],
            }
        ]
        dataset = LongMemEvalLoader().load_dict(raw)
        assert len(dataset.sessions) == 1
        assert dataset.sessions[0].questions[0].answer == "Tokyo"

    def test_load_nonexistent_file(self):
        with pytest.raises(FileNotFoundError):
            LongMemEvalLoader().load("/tmp/nonexistent_longmemeval.json")

    def test_load_dict_plain_conversation(self):
        """Test loading sessions that use plain conversation format (no turn_id)."""
        raw = {
            "sessions": [
                {
                    "session_id": "s1",
                    "turns": [
                        {"role": "user", "content": "Hello"},
                        {"role": "assistant", "content": "Hi there"},
                    ],
                    "questions": [],
                }
            ]
        }
        dataset = LongMemEvalLoader().load_dict(raw)
        assert len(dataset.sessions[0].turns) == 2
        assert dataset.sessions[0].turns[0].turn_id == 1
        assert dataset.sessions[0].turns[1].turn_id == 2
