from praisonaiagents.session.hierarchy import HierarchicalSessionStore

# Create a hierarchical session store
store = HierarchicalSessionStore(session_dir="./sessions")

# Create a session
session_id = store.create_session(title="My Project")

# Add messages
store.add_message(session_id, "user", "Hello!")
store.add_message(session_id, "assistant", "Hi there!")

# Create a snapshot
snapshot_id = store.create_snapshot(session_id, label="Checkpoint 1")

# Continue working...
store.add_message(session_id, "user", "Do something risky")

# Revert to snapshot if needed
store.revert_to_snapshot(session_id, snapshot_id)

class HierarchicalSessionStore(DefaultSessionStore):
    def create_session(
        self,
        title: Optional[str] = None,
        parent_id: Optional[str] = None
    ) -> str:
        """Create a new session, optionally as child of another."""
    
    def fork_session(
        self,
        session_id: str,
        from_message_index: Optional[int] = None,
        title: Optional[str] = None
    ) -> str:
        """Fork a session from a specific message point."""
    
    def create_snapshot(
        self,
        session_id: str,
        label: Optional[str] = None
    ) -> str:
        """Create a labeled snapshot of the current session state."""
    
    def revert_to_snapshot(
        self,
        session_id: str,
        snapshot_id: str
    ) -> bool:
        """Revert session to a previous snapshot."""
    
    def get_snapshots(self, session_id: str) -> List[SessionSnapshot]:
        """Get all snapshots for a session."""
    
    def export_session(self, session_id: str) -> Dict[str, Any]:
        """Export session data for transfer."""
    
    def import_session(self, data: Dict[str, Any]) -> str:
        """Import a session from exported data."""

store = HierarchicalSessionStore()

# Create main session
main_id = store.create_session(title="Main Branch")
store.add_message(main_id, "user", "Start project")
store.add_message(main_id, "assistant", "Project initialized")

# Fork from message index 1
fork_id = store.fork_session(
    main_id,
    from_message_index=1,
    title="Experimental Branch"
)

# Fork has messages up to index 1
# Can now diverge independently
store.add_message(fork_id, "user", "Try experimental approach")

store = HierarchicalSessionStore()
session_id = store.create_session()

# Work and create snapshots
store.add_message(session_id, "user", "Phase 1")
snap1 = store.create_snapshot(session_id, label="After Phase 1")

store.add_message(session_id, "user", "Phase 2")
snap2 = store.create_snapshot(session_id, label="After Phase 2")

# List all snapshots
snapshots = store.get_snapshots(session_id)
for snap in snapshots:
    print(f"{snap.label}: {snap.message_count} messages")

# Revert to Phase 1
store.revert_to_snapshot(session_id, snap1)

# Export from one store
store1 = HierarchicalSessionStore(session_dir="./store1")
session_id = store1.create_session(title="Portable Session")
store1.add_message(session_id, "user", "Important data")

exported = store1.export_session(session_id)

# Import to another store
store2 = HierarchicalSessionStore(session_dir="./store2")
new_id = store2.import_session(exported)