"""
tests/test_locomo_native.py — Tests for native LoCoMo dataset loader + converter.
"""
from __future__ import annotations

import pathlib
from collections import Counter

import pytest

from cortex.eval.locomo import (
    CATEGORY_MAP,
    LoCoMoConversation,
    LoCoMoDataset,
    LoCoMoLoader,
    LoCoMoQAPair,
    LoCoMoSession,
)
from cortex.eval.locomo_converter import (
    conversations_to_remember_payloads,
    session_to_remember_payloads,
)

NATIVE_PATH = pathlib.Path(__file__).parent.parent / "cortex" / "eval" / "datasets" / "locomo10.json"
SAMPLE_PATH = pathlib.Path(__file__).parent.parent / "cortex" / "eval" / "datasets" / "locomo_sample.json"


@pytest.fixture
def loader():
    return LoCoMoLoader()


@pytest.fixture
def native_dataset(loader):
    return loader.load_native(NATIVE_PATH)


# ── Test 1: Parses all 10 conversations ──


class TestNativeLoader:
    def test_loads_all_conversations(self, native_dataset):
        """load_native() parses all 10 conversations."""
        convs = native_dataset.conversations
        assert len(convs) == 10

    def test_all_conversations_have_sessions(self, native_dataset):
        """Each conversation has at least one session."""
        for conv in native_dataset.conversations:
            assert len(conv.sessions) > 0, f"{conv.sample_id} has no sessions"

    def test_sessions_flattened(self, native_dataset):
        """All sessions are flattened in dataset.sessions."""
        total_sessions = sum(len(c.sessions) for c in native_dataset.conversations)
        assert len(native_dataset.sessions) == total_sessions

    # ── Test 2: QA pair counts ──

    def test_total_qa_pairs(self, native_dataset):
        """Total QA pairs should be 1,986."""
        total = sum(len(c.qa_pairs) for c in native_dataset.conversations)
        assert total == 1986

    def test_total_qa_via_metadata(self, native_dataset):
        """Metadata tracks total QA count."""
        assert native_dataset.metadata["total_qa_pairs"] == 1986

    # ── Test 3: Category distribution ──

    def test_category_distribution(self, native_dataset):
        """Category counts match expected distribution."""
        all_qa = []
        for conv in native_dataset.conversations:
            all_qa.extend(conv.qa_pairs)

        cats = Counter(qa.category for qa in all_qa)
        assert cats[1] == 282, f"multi_hop: expected 282, got {cats[1]}"
        assert cats[2] == 321, f"single_hop: expected 321, got {cats[2]}"
        assert cats[3] == 96, f"temporal: expected 96, got {cats[3]}"
        assert cats[4] == 841, f"open_domain: expected 841, got {cats[4]}"
        assert cats[5] == 446, f"adversarial: expected 446, got {cats[5]}"

    def test_question_types_mapped(self, native_dataset):
        """Question types are mapped from category numbers."""
        all_qa = []
        for conv in native_dataset.conversations:
            all_qa.extend(conv.qa_pairs)

        types = Counter(qa.question_type for qa in all_qa)
        assert types["multi_hop"] == 282
        assert types["single_hop"] == 321
        assert types["temporal"] == 96
        assert types["open_domain"] == 841
        assert types["adversarial"] == 446

    # ── Test 4: Evidence fields preserved ──

    def test_evidence_preserved(self, native_dataset):
        """Evidence fields are preserved on QA pairs."""
        conv = native_dataset.conversations[0]
        # First QA in first conversation should have evidence
        qa_with_evidence = [qa for qa in conv.qa_pairs if qa.evidence]
        assert len(qa_with_evidence) > 0, "No QA pairs have evidence"

        # Check evidence format (should be like "D1:3")
        for qa in qa_with_evidence:
            for ev in qa.evidence:
                assert ":" in ev or ev.startswith("D"), f"Unexpected evidence format: {ev}"

    # ── Test 5: load_conversation ──

    def test_load_conversation_by_id(self, loader):
        """load_conversation returns single conversation by sample_id."""
        dataset = loader.load_conversation("conv-26", path=NATIVE_PATH)
        convs = dataset.conversations
        assert len(convs) == 1
        assert convs[0].sample_id == "conv-26"

    def test_load_conversation_by_index(self, loader):
        """load_conversation supports numeric index."""
        dataset = loader.load_conversation("0", path=NATIVE_PATH)
        convs = dataset.conversations
        assert len(convs) == 1
        # First conversation in the file
        assert convs[0].sample_id == "conv-26"

    def test_load_conversation_not_found(self, loader):
        """load_conversation raises ValueError for unknown sample_id."""
        with pytest.raises(ValueError, match="not found"):
            loader.load_conversation("nonexistent", path=NATIVE_PATH)

    def test_conversation_has_speakers(self, native_dataset):
        """Conversations have speaker_a and speaker_b."""
        for conv in native_dataset.conversations:
            assert conv.speaker_a, f"{conv.sample_id} missing speaker_a"
            assert conv.speaker_b, f"{conv.sample_id} missing speaker_b"

    def test_session_has_date(self, native_dataset):
        """Sessions from native loader have session_date set."""
        # At least the first session of each conversation should have a date
        for conv in native_dataset.conversations:
            first = conv.sessions[0]
            assert first.session_date, f"{conv.sample_id} session 1 missing date"

    def test_conversation_turns_have_dia_id(self, native_dataset):
        """Dialog turns include dia_id."""
        conv = native_dataset.conversations[0]
        first_session = conv.sessions[0]
        assert len(first_session.conversation) > 0
        first_turn = first_session.conversation[0]
        assert "dia_id" in first_turn
        assert first_turn["dia_id"], "dia_id should not be empty"


# ── Test 6: Converter ──


class TestConverter:
    def test_produces_payloads(self, native_dataset):
        """Converter produces remember() payloads from dataset."""
        payloads = conversations_to_remember_payloads(native_dataset)
        assert len(payloads) > 0

    def test_payload_structure(self, native_dataset):
        """Each payload has required fields."""
        payloads = conversations_to_remember_payloads(native_dataset)
        for p in payloads[:10]:  # Check first 10
            assert "content" in p
            assert "owner_id" in p
            assert "conversation_id" in p
            assert "metadata" in p
            assert p["owner_id"] == "locomo_eval"
            assert "dia_id" in p["metadata"]
            assert "session_date" in p["metadata"]
            assert "speaker" in p["metadata"]

    def test_payload_content_format(self, native_dataset):
        """Payload content is formatted as 'speaker: text'."""
        payloads = conversations_to_remember_payloads(native_dataset)
        for p in payloads[:5]:
            assert ": " in p["content"]

    def test_custom_owner_id(self, native_dataset):
        """Converter respects custom owner_id."""
        payloads = conversations_to_remember_payloads(native_dataset, owner_id="test_owner")
        assert all(p["owner_id"] == "test_owner" for p in payloads[:5])

    def test_session_converter(self, native_dataset):
        """session_to_remember_payloads works for a single session."""
        session = native_dataset.sessions[0]
        payloads = session_to_remember_payloads(session)
        assert len(payloads) == len(session.conversation)


# ── Test 7: Backward compatibility ──


class TestBackwardCompat:
    def test_load_sample_still_works(self, loader):
        """load() with bundled sample still works."""
        if not SAMPLE_PATH.exists():
            pytest.skip("Sample dataset not available")
        dataset = loader.load()
        assert dataset.name
        assert len(dataset.sessions) > 0

    def test_qa_pair_defaults(self):
        """LoCoMoQAPair new fields have safe defaults."""
        qa = LoCoMoQAPair(
            question_id="q1",
            question="What?",
            answer="That.",
        )
        assert qa.question_type == "single_hop"
        assert qa.category == 0
        assert qa.evidence == []

    def test_session_date_default(self):
        """LoCoMoSession.session_date defaults to empty string."""
        session = LoCoMoSession(session_id="s1", conversation=[])
        assert session.session_date == ""

    def test_category_map_completeness(self):
        """CATEGORY_MAP covers all 5 categories."""
        assert set(CATEGORY_MAP.keys()) == {1, 2, 3, 4, 5}

    def test_dataset_conversations_property_empty(self):
        """conversations property returns empty list for non-native datasets."""
        ds = LoCoMoDataset(name="test")
        assert ds.conversations == []
