"""PEFState: persistent present-state container with indexed relationships."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from .span import Span
from .entity import Entity, EchoTrace, TurnBindings


# ── Canonical relation mapping ───────────────────────────────────────

_RELATION_ALIASES: dict[str, str] = {
    "has": "HAS",
    "have": "HAS",
    "had": "HAS",
    "owns": "HAS",
    "own": "HAS",
    "possessed": "HAS",
    "possesses": "HAS",
    "be": "IS",    # spaCy lemma for is/was/are/were — must precede surface forms
    "is": "IS",
    "was": "IS",
    "are": "IS",
    "were": "IS",
    "at": "AT",
    "in": "AT",
    "on": "AT",
    "located": "AT",
    "lives": "AT",
    "lived": "AT",
    "gave": "GIVE",
    "gives": "GIVE",
    "give": "GIVE",
    "given": "GIVE",
    "sent": "SEND",
    "sends": "SEND",
    "send": "SEND",
    "told": "TELL",
    "tells": "TELL",
    "tell": "TELL",
    "showed": "SHOW",
    "shows": "SHOW",
    "show": "SHOW",
    "returned": "RETURN",
    "returns": "RETURN",
    "return": "RETURN",
    "because": "BECAUSE",
    "likes": "LIKES",
    "like": "LIKES",
    "liked": "LIKES",
    "loves": "LOVES",
    "love": "LOVES",
    "loved": "LOVES",
    "wants": "WANTS",
    "want": "WANTS",
    "wanted": "WANTS",
    "knows": "KNOWS",
    "know": "KNOWS",
    "knew": "KNOWS",
}

CANONICAL_RELATIONS = frozenset({
    "HAS", "IS", "AT", "GIVE", "SEND", "TELL", "SHOW",
    "RETURN", "BECAUSE", "LIKES", "LOVES", "WANTS", "KNOWS",
})


def canonicalize_relation(surface: str) -> str:
    """Map a surface verb form to its canonical relation string.

    Returns the uppercase canonical form, or the uppercased input
    if no alias is defined (allowing extension).
    """
    return _RELATION_ALIASES.get(surface.lower().strip(), surface.upper().strip())


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

@dataclass
class Relationship:
    """A single fact linking a subject to an object (entity or literal).

    Exactly one of object_entity_id or object_literal must be non-null.

    Provenance: where the fact came from. user_input | pre_populated | llm_output | system.
    extractor_backend: spacy | llm | rule | manual | unknown.
    """
    subject_id: str
    relation: str                       # Must be canonical (pass through canonicalize_relation)

    object_entity_id: str | None        # References another entity
    object_literal: Any | None          # Literal value ("red", 42, "a book")

    span: Span
    source_turn: int                    # When established
    evidence: str                       # Surface text supporting this

    negated: bool = False               # Explicit negation ("doesn't have")

    provenance: str = "user_input"       # user_input | pre_populated | llm_output | system
    extractor_backend: str = "unknown"  # spacy | llm | rule | manual | unknown
    document_id: str | None = None
    document_locator: str | None = None  # e.g. page=3, line=120-140, section=Employment

    def __post_init__(self):
        has_entity = self.object_entity_id is not None
        has_literal = self.object_literal is not None
        if has_entity == has_literal:
            raise ValueError(
                "Exactly one of object_entity_id or object_literal must be non-null. "
                f"Got entity={self.object_entity_id!r}, literal={self.object_literal!r}"
            )


# ── Serialization helpers (Phase C) ───────────────────────────────────

def _entity_to_dict(e: Entity) -> dict[str, Any]:
    return {
        "id": e.id,
        "name": e.name,
        "aliases": sorted(e.aliases),
        "attributes": dict(e.attributes),
        "echo_traces": [
            {"turn": et.turn, "context": et.context, "span": et.span.value}
            for et in e.echo_traces
        ],
        "turn_introduced": e.turn_introduced,
        "turn_last_active": e.turn_last_active,
        "resolved": e.resolved,
    }


def _entity_from_dict(d: dict[str, Any]) -> Entity:
    aliases = set(d.get("aliases", []))
    if not aliases and d.get("name"):
        aliases = {d["name"]}
    echo_traces = [
        EchoTrace(
            turn=et["turn"],
            context=et.get("context", ""),
            span=Span(et["span"]) if isinstance(et.get("span"), str) else Span.PRESENT,
        )
        for et in d.get("echo_traces", [])
    ]
    return Entity(
        id=d["id"],
        name=d["name"],
        aliases=aliases,
        attributes=dict(d.get("attributes", {})),
        echo_traces=echo_traces,
        turn_introduced=int(d.get("turn_introduced", 0)),
        turn_last_active=int(d.get("turn_last_active", 0)),
        resolved=bool(d.get("resolved", True)),
    )


def _rel_to_dict(r: Relationship) -> dict[str, Any]:
    out: dict[str, Any] = {
        "subject_id": r.subject_id,
        "relation": r.relation,
        "object_entity_id": r.object_entity_id,
        "object_literal": r.object_literal,
        "span": r.span.value,
        "source_turn": r.source_turn,
        "evidence": r.evidence,
        "negated": r.negated,
        "provenance": r.provenance,
        "extractor_backend": r.extractor_backend,
    }
    if r.document_id is not None:
        out["document_id"] = r.document_id
    if r.document_locator is not None:
        out["document_locator"] = r.document_locator
    return out


def _rel_from_dict(d: dict[str, Any]) -> Relationship:
    span_val = d.get("span", "present")
    span = Span(span_val) if isinstance(span_val, str) else Span.PRESENT
    return Relationship(
        subject_id=d["subject_id"],
        relation=d["relation"],
        object_entity_id=d.get("object_entity_id"),
        object_literal=d.get("object_literal"),
        span=span,
        source_turn=int(d.get("source_turn", 0)),
        evidence=d.get("evidence", ""),
        negated=bool(d.get("negated", False)),
        provenance=d.get("provenance", "user_input"),
        extractor_backend=d.get("extractor_backend", "unknown"),
        document_id=d.get("document_id"),
        document_locator=d.get("document_locator"),
    )


def _turn_bindings_to_dict(tb: TurnBindings) -> dict[str, Any]:
    return {"turn": tb.turn, "bindings": dict(tb.bindings)}


def _turn_bindings_from_dict(d: dict[str, Any]) -> TurnBindings:
    return TurnBindings(
        turn=int(d.get("turn", 0)),
        bindings=dict(d.get("bindings", {})),
    )


# ── PEF State ────────────────────────────────────────────────────────

@dataclass
class PEFState:
    """The persistent present-state container.

    Maintains entities, relationships, and indexes for O(1) lookup.
    """
    entities: dict[str, Entity] = field(default_factory=dict)
    relationships: list[Relationship] = field(default_factory=list)

    # Indexes into relationships list
    rel_by_subject: dict[str, list[int]] = field(default_factory=dict)
    rel_by_object_entity: dict[str, list[int]] = field(default_factory=dict)
    rel_by_relation: dict[str, list[int]] = field(default_factory=dict)

    turn_bindings: dict[int, TurnBindings] = field(default_factory=dict)
    current_turn: int = 0
    active_span: Span = Span.PRESENT

    # Session-scoped entity IDs: name_index[canon_name] = next ordinal (2,3,...)
    session_id: str = ""
    name_index: dict[str, int] = field(default_factory=dict)

    # ── Entity operations ────────────────────────────────────────

    def add_entity(self, entity: Entity) -> Entity:
        """Register an entity. Returns the entity (for chaining)."""
        self.entities[entity.id] = entity
        return entity

    def find_entity_by_name(self, name: str) -> Entity | None:
        """Look up an entity by canonical name or alias (case-insensitive)."""
        name_lower = name.lower()
        for entity in self.entities.values():
            if entity.name.lower() == name_lower:
                return entity
            if any(a.lower() == name_lower for a in entity.aliases):
                return entity
        return None

    def get_or_create_entity(
        self,
        name: str,
        *,
        resolved: bool = True,
        entity_id: str | None = None,
        force_new: bool = False,
    ) -> tuple[Entity, bool]:
        """Find entity by name or create a new one.

        Unification: reuse if entity_id supplied (caller binding) or find by name.
        Sibling: when force_new=True and canon_name exists, mint ordinal #2, #3, ...
        Returns (entity, created) where created is True if new.
        """
        canon_name = name.strip().lower()
        if entity_id is not None and entity_id in self.entities:
            e = self.entities[entity_id]
            e.turn_last_active = self.current_turn
            return e, False
        if not force_new:
            existing = self.find_entity_by_name(name)
            if existing is not None:
                existing.turn_last_active = self.current_turn
                return existing, False
        # Create new: base (first) or sibling (ordinal)
        existing = self.find_entity_by_name(name)
        if existing is None:
            entity = Entity.create(
                name,
                self.current_turn,
                resolved=resolved,
                session_id=self.session_id,
                ordinal=None,
            )
        else:
            k = self.name_index.get(canon_name, 2)
            self.name_index[canon_name] = k + 1
            entity = Entity.create(
                name,
                self.current_turn,
                resolved=resolved,
                session_id=self.session_id,
                ordinal=k,
            )
        self.add_entity(entity)
        return entity, True

    # ── Relationship operations ──────────────────────────────────

    def add_relationship(self, rel: Relationship) -> int:
        """Add a relationship and update all indexes. Returns the index."""
        idx = len(self.relationships)
        self.relationships.append(rel)

        self.rel_by_subject.setdefault(rel.subject_id, []).append(idx)
        if rel.object_entity_id is not None:
            self.rel_by_object_entity.setdefault(rel.object_entity_id, []).append(idx)
        self.rel_by_relation.setdefault(rel.relation, []).append(idx)

        return idx

    def get_relationships_for_subject(self, subject_id: str) -> list[Relationship]:
        """All relationships where this entity is the subject."""
        indices = self.rel_by_subject.get(subject_id, [])
        return [self.relationships[i] for i in indices]

    def get_relationships_for_object(self, entity_id: str) -> list[Relationship]:
        """All relationships where this entity is the object."""
        indices = self.rel_by_object_entity.get(entity_id, [])
        return [self.relationships[i] for i in indices]

    def get_relationships_by_relation(self, relation: str) -> list[Relationship]:
        """All relationships of a specific type."""
        indices = self.rel_by_relation.get(relation, [])
        return [self.relationships[i] for i in indices]

    def find_relationship(
        self,
        subject_id: str,
        relation: str,
        *,
        object_entity_id: str | None = None,
        object_literal: Any | None = None,
    ) -> Relationship | None:
        """Find a specific relationship matching the given criteria."""
        for rel in self.get_relationships_for_subject(subject_id):
            if rel.relation != relation:
                continue
            if object_entity_id is not None and rel.object_entity_id == object_entity_id:
                return rel
            if object_literal is not None and rel.object_literal == object_literal:
                return rel
        return None

    # ── Pronoun binding operations ───────────────────────────────

    def resolve_pronoun(self, pronoun_key: str, entity_id: str) -> None:
        """Bind a pronoun occurrence to an entity for the current turn.

        pronoun_key should be occurrence-based, e.g. "she@token_14".
        """
        if self.current_turn not in self.turn_bindings:
            self.turn_bindings[self.current_turn] = TurnBindings(
                turn=self.current_turn
            )
        self.turn_bindings[self.current_turn].bindings[pronoun_key] = entity_id

    def get_binding(self, pronoun_key: str, turn: int | None = None) -> str | None:
        """Look up what entity a pronoun resolved to in a specific turn."""
        t = turn if turn is not None else self.current_turn
        tb = self.turn_bindings.get(t)
        if tb is None:
            return None
        return tb.bindings.get(pronoun_key)

    # ── Echo trace operations ────────────────────────────────────

    def record_echo_trace(self, entity_id: str, context: str) -> None:
        """Record an echo trace for an entity at the current turn and span."""
        entity = self.entities.get(entity_id)
        if entity is None:
            raise KeyError(f"Entity {entity_id!r} not found")
        trace = EchoTrace(
            turn=self.current_turn,
            context=context,
            span=self.active_span,
        )
        entity.echo_traces.append(trace)
        entity.turn_last_active = self.current_turn

    # ── Turn management ──────────────────────────────────────────

    def advance_turn(self) -> int:
        """Advance to the next turn. Returns the new turn number."""
        self.current_turn += 1
        return self.current_turn

    # ── Serialization (Phase C) ─────────────────────────────────────────────

    def to_dict(self) -> dict[str, Any]:
        """Serialize to JSON-serializable dict. Round-trip with from_dict()."""
        # Sort relationships for deterministic state_hash (claims/LLM order may vary).
        def _rel_sort_key(r: Relationship) -> tuple:
            return (
                r.subject_id,
                r.relation,
                r.object_entity_id or "",
                str(r.object_literal) if r.object_literal is not None else "",
                r.source_turn,
                r.evidence,
                r.provenance,
                r.extractor_backend,
                r.document_id or "",
                r.document_locator or "",
            )
        return {
            "entities": {
                eid: _entity_to_dict(e) for eid, e in self.entities.items()
            },
            "relationships": [
                _rel_to_dict(r)
                for r in sorted(self.relationships, key=_rel_sort_key)
            ],
            "turn_bindings": {
                f"{turn:05d}": _turn_bindings_to_dict(tb)
                for turn, tb in sorted(self.turn_bindings.items())
            },
            "current_turn": self.current_turn,
            "active_span": self.active_span.value,
            "session_id": self.session_id,
            "name_index": dict(sorted(self.name_index.items())),
        }

    @staticmethod
    def from_dict(data: dict[str, Any]) -> "PEFState":
        """Deserialize from dict. Rebuilds indexes from relationships."""
        entities = {
            eid: _entity_from_dict(d) for eid, d in data.get("entities", {}).items()
        }
        relationships = [_rel_from_dict(d) for d in data.get("relationships", [])]
        turn_bindings = {
            int(k): _turn_bindings_from_dict(v)
            for k, v in data.get("turn_bindings", {}).items()
        }
        current_turn = int(data.get("current_turn", 0))
        active_span_val = data.get("active_span", "present")
        active_span = Span(active_span_val) if isinstance(active_span_val, str) else Span.PRESENT
        session_id = str(data.get("session_id", ""))
        name_index = dict(data.get("name_index", {}))

        state = PEFState(
            entities=entities,
            relationships=relationships,
            turn_bindings=turn_bindings,
            current_turn=current_turn,
            active_span=active_span,
            session_id=session_id,
            name_index=name_index,
        )
        # Rebuild indexes
        for idx, rel in enumerate(state.relationships):
            state.rel_by_subject.setdefault(rel.subject_id, []).append(idx)
            if rel.object_entity_id is not None:
                state.rel_by_object_entity.setdefault(rel.object_entity_id, []).append(idx)
            state.rel_by_relation.setdefault(rel.relation, []).append(idx)
        return state

    # ── Summary for LLM injection ────────────────────────────────

    def to_context_summary(self) -> str:
        """Generate a ground-truth summary suitable for LLM system prompt injection."""
        if not self.entities and not self.relationships:
            return "No established facts."

        lines: list[str] = []

        # Entities
        if self.entities:
            lines.append("## Known Entities")
            for entity in self.entities.values():
                status = "" if entity.resolved else " [UNRESOLVED]"
                attrs = ""
                if entity.attributes:
                    attr_parts = [f"{k}={v}" for k, v in entity.attributes.items()]
                    attrs = f" ({', '.join(attr_parts)})"
                lines.append(f"- {entity.name}{status}{attrs}")

        # Relationships (grouped by subject)
        present_rels: list[str] = []
        past_rels: list[str] = []

        for rel in self.relationships:
            subj = self.entities.get(rel.subject_id)
            subj_name = subj.name if subj else rel.subject_id

            obj_str: str
            if rel.object_entity_id:
                obj_ent = self.entities.get(rel.object_entity_id)
                obj_str = obj_ent.name if obj_ent else rel.object_entity_id
            else:
                obj_str = str(rel.object_literal)

            neg = "NOT " if rel.negated else ""
            fact = f"- {subj_name} {rel.relation} {neg}{obj_str}"

            if rel.span == Span.PRESENT:
                present_rels.append(fact)
            else:
                past_rels.append(fact)

        if present_rels:
            lines.append("\n## Current Facts (present)")
            lines.extend(present_rels)

        if past_rels:
            lines.append("\n## Past Facts")
            lines.extend(past_rels)

        return "\n".join(lines)
