echo/core/tree.py

"""
echo/core/tree.py
-----------------
The branching tree of alternate lives. Holds every WorldState node and the
parent/child links the UI renders (gold for flourishing, dark for struggling).

Pure data structure + traversal helpers. The orchestrator mutates it; the
Gradio front-end reads it to draw the D3/vis-network graph.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional

from .world_state import WorldState


@dataclass
class LifeTree:
    seed_choice: str
    nodes: dict[str, WorldState] = field(default_factory=dict)
    root_id: Optional[str] = None

    # ---------------------------------------------------------------- build
    def add(self, state: WorldState) -> None:
        self.nodes[state.node_id] = state
        if state.parent_id is None:
            self.root_id = state.node_id

    def children(self, node_id: str) -> list[WorldState]:
        return [n for n in self.nodes.values() if n.parent_id == node_id]

    def path_to_root(self, node_id: str) -> list[WorldState]:
        """Ancestor chain from root -> node (inclusive). The branch's history."""
        chain: list[WorldState] = []
        cur = self.nodes.get(node_id)
        while cur is not None:
            chain.append(cur)
            cur = self.nodes.get(cur.parent_id) if cur.parent_id else None
        return list(reversed(chain))

    def branch_narrative(self, node_id: str) -> str:
        """Human-readable history of a branch, fed to the Curator as context."""
        chain = self.path_to_root(node_id)
        lines = []
        for n in chain:
            if n.depth == 0:
                lines.append(f"Origin: {n.divergence}")
            else:
                lines.append(
                    f"+{n.years_elapsed}y -> {n.divergence} "
                    f"(now: {n.facts.constraints_text()})"
                )
        return "\n".join(lines)

    # --------------------------------------------------------------- export
    def to_graph(self) -> dict:
        """Serialize to nodes/edges for the front-end visualization."""
        graph_nodes = []
        graph_edges = []
        for n in self.nodes.values():
            graph_nodes.append({
                "id": n.node_id,
                "label": n.facts.occupation or "…",
                "summary": n.summary,
                "valence": n.tone.valence,
                "flourishing": n.tone.is_flourishing,
                "depth": n.depth,
                "has_voice": n.voice_audio_path is not None,
            })
            if n.parent_id:
                graph_edges.append({
                    "from": n.parent_id, "to": n.node_id,
                    "label": n.divergence[:40],
                })
        return {"nodes": graph_nodes, "edges": graph_edges,
                "root": self.root_id, "seed": self.seed_choice}

    @property
    def size(self) -> int:
        return len(self.nodes)

    @property
    def max_depth(self) -> int:
        return max((n.depth for n in self.nodes.values()), default=0)