"""Tests for PEF schema: Span, Entity, Relationship, PEFState."""

import pytest

from aurora_lens.pef.span import Span
from aurora_lens.pef.entity import Entity, EchoTrace, TurnBindings
from aurora_lens.pef.state import (
    PEFState,
    Relationship,
    canonicalize_relation,
)


# ── Span ─────────────────────────────────────────────────────────────

class TestSpan:
    def test_values(self):
        assert Span.PRESENT.value == "present"
        assert Span.PAST.value == "past"

    def test_enum_identity(self):
        assert Span("present") is Span.PRESENT
        assert Span("past") is Span.PAST


# ── Entity ───────────────────────────────────────────────────────────

class TestEntity:
    def test_create(self):
        e = Entity.create("Emma", turn=1)
        assert e.name == "Emma"
        assert "Emma" in e.aliases
        assert e.turn_introduced == 1
        assert e.resolved is True
        assert len(e.id) > 0  # UUID

    def test_create_unresolved(self):
        e = Entity.create("someone", turn=2, resolved=False)
        assert e.resolved is False

    def test_unique_ids(self):
        """Different names yield different IDs; same name in same session yields same base ID."""
        e1 = Entity.create("Emma", turn=1, session_id="s1")
        e2 = Entity.create("Lucy", turn=1, session_id="s1")
        assert e1.id != e2.id
        e3 = Entity.create("Emma", turn=2, session_id="s1")
        assert e1.id == e3.id  # Same logical entity (base)

    def test_entity_id_session_scoped(self):
        """Same name in two different session_ids yields different IDs."""
        e1 = Entity.create("John Smith", turn=1, session_id="session-a")
        e2 = Entity.create("John Smith", turn=1, session_id="session-b")
        assert e1.id != e2.id

    def test_entity_id_same_session_siblings(self):
        """Two entities named 'john smith' in same session yields base id and #2 id deterministically."""
        pef = PEFState(session_id="s1")
        e1, created1 = pef.get_or_create_entity("John Smith", resolved=True)
        assert created1 is True
        base_id = e1.id
        e2, created2 = pef.get_or_create_entity("John Smith", force_new=True, resolved=True)
        assert created2 is True
        assert e1.id != e2.id
        assert e2.id != base_id
        # Ordinal #2 should be deterministic
        expected_id = Entity.create("John Smith", turn=pef.current_turn, session_id="s1", ordinal=2).id
        assert e2.id == expected_id


# ── EchoTrace ────────────────────────────────────────────────────────

class TestEchoTrace:
    def test_creation(self):
        trace = EchoTrace(turn=3, context="Emma picked up a book", span=Span.PAST)
        assert trace.turn == 3
        assert trace.span == Span.PAST


# ── TurnBindings ─────────────────────────────────────────────────────

class TestTurnBindings:
    def test_per_turn_isolation(self):
        tb1 = TurnBindings(turn=1, bindings={"she@token_3": "id_emma"})
        tb2 = TurnBindings(turn=2, bindings={"she@token_5": "id_lucy"})
        assert tb1.bindings["she@token_3"] == "id_emma"
        assert tb2.bindings["she@token_5"] == "id_lucy"
        # Different turns, different bindings
        assert "she@token_3" not in tb2.bindings

    def test_occurrence_based_keys(self):
        """Pronouns are occurrence-based, not surface-form-based."""
        tb = TurnBindings(turn=1, bindings={
            "she@token_3": "id_emma",
            "she@token_12": "id_lucy",
        })
        # Same surface form, different occurrences, different entities
        assert tb.bindings["she@token_3"] != tb.bindings["she@token_12"]


# ── canonicalize_relation ────────────────────────────────────────────

class TestCanonicalizeRelation:
    def test_known_aliases(self):
        assert canonicalize_relation("has") == "HAS"
        assert canonicalize_relation("have") == "HAS"
        assert canonicalize_relation("had") == "HAS"
        assert canonicalize_relation("is") == "IS"
        assert canonicalize_relation("was") == "IS"
        assert canonicalize_relation("gave") == "GIVE"
        assert canonicalize_relation("sent") == "SEND"

    def test_case_insensitive(self):
        assert canonicalize_relation("Has") == "HAS"
        assert canonicalize_relation("HAS") == "HAS"

    def test_unknown_passes_through(self):
        assert canonicalize_relation("eats") == "EATS"
        assert canonicalize_relation("CUSTOM") == "CUSTOM"


# ── Relationship ─────────────────────────────────────────────────────

class TestRelationship:
    def test_object_entity(self):
        r = Relationship(
            subject_id="s1", relation="GIVE",
            object_entity_id="o1", object_literal=None,
            span=Span.PRESENT, source_turn=1, evidence="gave it",
        )
        assert r.object_entity_id == "o1"
        assert r.object_literal is None

    def test_object_literal(self):
        r = Relationship(
            subject_id="s1", relation="HAS",
            object_entity_id=None, object_literal="a red book",
            span=Span.PRESENT, source_turn=1, evidence="has a red book",
        )
        assert r.object_literal == "a red book"
        assert r.object_entity_id is None

    def test_both_null_raises(self):
        with pytest.raises(ValueError, match="Exactly one"):
            Relationship(
                subject_id="s1", relation="HAS",
                object_entity_id=None, object_literal=None,
                span=Span.PRESENT, source_turn=1, evidence="bad",
            )

    def test_both_set_raises(self):
        with pytest.raises(ValueError, match="Exactly one"):
            Relationship(
                subject_id="s1", relation="HAS",
                object_entity_id="o1", object_literal="also a book",
                span=Span.PRESENT, source_turn=1, evidence="bad",
            )

    def test_negated_default(self):
        r = Relationship(
            subject_id="s1", relation="HAS",
            object_entity_id=None, object_literal="key",
            span=Span.PRESENT, source_turn=1, evidence="has a key",
        )
        assert r.negated is False


# ── PEFState ─────────────────────────────────────────────────────────

class TestPEFState:
    def test_add_entity(self):
        pef = PEFState()
        e = Entity.create("Emma", turn=0)
        pef.add_entity(e)
        assert e.id in pef.entities
        assert pef.entities[e.id].name == "Emma"

    def test_find_entity_by_name(self):
        pef = PEFState()
        e = Entity.create("Emma", turn=0)
        e.aliases.add("Ms Smith")
        pef.add_entity(e)

        assert pef.find_entity_by_name("Emma") is e
        assert pef.find_entity_by_name("emma") is e  # case insensitive
        assert pef.find_entity_by_name("Ms Smith") is e
        assert pef.find_entity_by_name("Lucy") is None

    def test_get_or_create_entity(self):
        pef = PEFState()
        e1, created1 = pef.get_or_create_entity("Emma")
        assert created1 is True

        e2, created2 = pef.get_or_create_entity("Emma")
        assert created2 is False
        assert e1.id == e2.id

    def test_add_relationship_indexes(self):
        pef = PEFState()
        e = Entity.create("Emma", turn=0)
        pef.add_entity(e)

        rel = Relationship(
            subject_id=e.id, relation="HAS",
            object_entity_id=None, object_literal="red book",
            span=Span.PRESENT, source_turn=0, evidence="Emma has a red book",
        )
        idx = pef.add_relationship(rel)

        assert idx == 0
        assert pef.rel_by_subject[e.id] == [0]
        assert pef.rel_by_relation["HAS"] == [0]
        assert e.id not in pef.rel_by_object_entity  # literal, not entity

    def test_relationship_entity_object_index(self):
        pef = PEFState()
        emma = Entity.create("Emma", turn=0)
        lucy = Entity.create("Lucy", turn=0)
        pef.add_entity(emma)
        pef.add_entity(lucy)

        rel = Relationship(
            subject_id=emma.id, relation="GIVE",
            object_entity_id=lucy.id, object_literal=None,
            span=Span.PRESENT, source_turn=0, evidence="gave to Lucy",
        )
        pef.add_relationship(rel)

        assert pef.rel_by_object_entity[lucy.id] == [0]

    def test_get_relationships_for_subject(self):
        pef = PEFState()
        e = Entity.create("Emma", turn=0)
        pef.add_entity(e)

        r1 = Relationship(
            subject_id=e.id, relation="HAS",
            object_entity_id=None, object_literal="book",
            span=Span.PRESENT, source_turn=0, evidence="has book",
        )
        r2 = Relationship(
            subject_id=e.id, relation="IS",
            object_entity_id=None, object_literal="teacher",
            span=Span.PRESENT, source_turn=0, evidence="is teacher",
        )
        pef.add_relationship(r1)
        pef.add_relationship(r2)

        rels = pef.get_relationships_for_subject(e.id)
        assert len(rels) == 2

    def test_resolve_pronoun(self):
        pef = PEFState()
        e = Entity.create("Emma", turn=0)
        pef.add_entity(e)
        pef.current_turn = 1

        pef.resolve_pronoun("she@token_3", e.id)

        assert pef.get_binding("she@token_3") == e.id
        assert pef.get_binding("she@token_3", turn=1) == e.id
        assert pef.get_binding("she@token_3", turn=0) is None  # different turn

    def test_record_echo_trace(self):
        pef = PEFState()
        e = Entity.create("Emma", turn=0)
        pef.add_entity(e)
        pef.current_turn = 1

        pef.record_echo_trace(e.id, "Emma picked up the book")

        assert len(e.echo_traces) == 1
        assert e.echo_traces[0].turn == 1
        assert e.echo_traces[0].span == Span.PRESENT
        assert e.turn_last_active == 1

    def test_record_echo_trace_unknown_entity(self):
        pef = PEFState()
        with pytest.raises(KeyError):
            pef.record_echo_trace("nonexistent", "text")

    def test_advance_turn(self):
        pef = PEFState()
        assert pef.current_turn == 0
        assert pef.advance_turn() == 1
        assert pef.advance_turn() == 2

    def test_to_context_summary_empty(self):
        pef = PEFState()
        assert pef.to_context_summary() == "No established facts."

    def test_to_context_summary_with_data(self):
        pef = PEFState()
        e = Entity.create("Emma", turn=0)
        pef.add_entity(e)

        rel = Relationship(
            subject_id=e.id, relation="HAS",
            object_entity_id=None, object_literal="red book",
            span=Span.PRESENT, source_turn=0, evidence="has a red book",
        )
        pef.add_relationship(rel)

        summary = pef.to_context_summary()
        assert "Emma" in summary
        assert "HAS" in summary
        assert "red book" in summary

    def test_find_relationship(self):
        pef = PEFState()
        e = Entity.create("Emma", turn=0)
        pef.add_entity(e)

        rel = Relationship(
            subject_id=e.id, relation="HAS",
            object_entity_id=None, object_literal="book",
            span=Span.PRESENT, source_turn=0, evidence="has book",
        )
        pef.add_relationship(rel)

        found = pef.find_relationship(e.id, "HAS", object_literal="book")
        assert found is not None
        assert found.object_literal == "book"

        not_found = pef.find_relationship(e.id, "HAS", object_literal="car")
        assert not_found is None

    def test_to_dict_from_dict_round_trip(self):
        """Phase C: PEFState serialization round-trip preserves all data."""
        pef = PEFState()
        emma = Entity.create("Emma", turn=1)
        emma.aliases.add("Ms Smith")
        emma.attributes["role"] = "teacher"
        pef.add_entity(emma)
        lucy = Entity.create("Lucy", turn=2)
        pef.add_entity(lucy)
        pef.advance_turn()
        pef.advance_turn()
        rel = Relationship(
            subject_id=emma.id, relation="HAS",
            object_entity_id=None, object_literal="book",
            span=Span.PRESENT, source_turn=1, evidence="Emma has a book",
        )
        pef.add_relationship(rel)
        pef.resolve_pronoun("she@1", emma.id)

        data = pef.to_dict()
        restored = PEFState.from_dict(data)
        assert len(restored.entities) == len(pef.entities)
        assert len(restored.relationships) == len(pef.relationships)
        assert restored.current_turn == pef.current_turn
        assert restored.active_span == pef.active_span
        e = restored.find_entity_by_name("Emma")
        assert e is not None
        assert "Ms Smith" in e.aliases
        assert e.attributes.get("role") == "teacher"
        assert restored.get_relationships_for_subject(e.id)[0].object_literal == "book"
        assert restored.get_binding("she@1", 2) == e.id

    def test_to_dict_deterministic_for_state_hash(self):
        """Same logical state yields identical JSON bytes regardless of add order (replay verification)."""
        import json

        def build_pef_a() -> PEFState:
            pef = PEFState()
            emma = Entity.create("Emma", turn=1)
            lucy = Entity.create("Lucy", turn=1)
            pef.add_entity(emma)
            pef.add_entity(lucy)
            pef.add_relationship(Relationship(
                subject_id=emma.id, relation="HAS",
                object_entity_id=None, object_literal="book",
                span=Span.PRESENT, source_turn=1, evidence="Emma has a book",
            ))
            pef.add_relationship(Relationship(
                subject_id=lucy.id, relation="HAS",
                object_entity_id=None, object_literal="pen",
                span=Span.PRESENT, source_turn=1, evidence="Lucy has a pen",
            ))
            return pef

        def build_pef_b() -> PEFState:
            pef = PEFState()
            lucy = Entity.create("Lucy", turn=1)
            emma = Entity.create("Emma", turn=1)
            pef.add_entity(lucy)
            pef.add_entity(emma)
            pef.add_relationship(Relationship(
                subject_id=lucy.id, relation="HAS",
                object_entity_id=None, object_literal="pen",
                span=Span.PRESENT, source_turn=1, evidence="Lucy has a pen",
            ))
            pef.add_relationship(Relationship(
                subject_id=emma.id, relation="HAS",
                object_entity_id=None, object_literal="book",
                span=Span.PRESENT, source_turn=1, evidence="Emma has a book",
            ))
            return pef

        data_a = build_pef_a().to_dict()
        data_b = build_pef_b().to_dict()
        json_a = json.dumps(data_a, sort_keys=True, separators=(",", ":"))
        json_b = json.dumps(data_b, sort_keys=True, separators=(",", ":"))
        assert json_a == json_b, "state_hash must be stable across add-order variation"
