# Current capabilities
nx.memory.store(content, scope="user", memory_type="fact", importance=0.8)
nx.memory.query(scope="user", memory_type="preference")
nx.memory.search(query="Python best practices")

class MemoryType:
    # Existing
    FACT = "fact"               # Factual knowledge
    PREFERENCE = "preference"   # User/agent preferences
    EXPERIENCE = "experience"   # Past experiences

    # NEW: ACE-specific types
    STRATEGY = "strategy"       # Successful approaches (✓ helpful)
    ANTI_PATTERN = "anti_pattern"  # Failed patterns (✗ harmful)
    OBSERVATION = "observation"    # Neutral observations (○)
    TRAJECTORY = "trajectory"      # Execution trace
    REFLECTION = "reflection"      # Analysis of outcomes
    CONSOLIDATED = "consolidated"  # Merged memories

from dataclasses import dataclass
from typing import Literal

@dataclass
class StrategyEntry:
    """Single strategy in playbook."""
    category: Literal["helpful", "harmful", "neutral"]  # ✓, ✗, ○
    pattern: str  # Description of the pattern
    context: str  # When to use/avoid this
    confidence: float  # 0.0-1.0
    evidence_count: int  # Number of trajectories supporting this
    created_at: str
    updated_at: str
    examples: list[str]  # Example trajectories

@dataclass
class PlaybookContent:
    """Structured playbook content."""
    version: int
    strategies: list[StrategyEntry]

    # Categorized for quick access
    helpful: list[StrategyEntry]  # ✓ Do these
    harmful: list[StrategyEntry]  # ✗ Avoid these
    neutral: list[StrategyEntry]  # ○ Context-dependent

    # Meta
    last_consolidated: str | None  # ISO timestamp
    consolidation_count: int
    total_trajectories_analyzed: int

# Example playbook
playbook = PlaybookContent(
    version=5,
    strategies=[
        StrategyEntry(
            category="helpful",
            pattern="When processing API errors, always check rate limits first",
            context="API integration tasks with 429 errors",
            confidence=0.95,
            evidence_count=12,
            created_at="2025-10-15T10:00:00Z",
            updated_at="2025-10-28T14:30:00Z",
            examples=["traj_abc123", "traj_def456"]
        ),
        StrategyEntry(
            category="harmful",
            pattern="Retrying immediately without backoff causes cascading failures",
            context="Error handling in API calls",
            confidence=0.88,
            evidence_count=7,
            created_at="2025-10-20T11:00:00Z",
            updated_at="2025-10-28T14:30:00Z",
            examples=["traj_xyz789"]
        )
    ],
    helpful=[...],
    harmful=[...],
    neutral=[...],
    last_consolidated="2025-10-28T14:30:00Z",
    consolidation_count=3,
    total_trajectories_analyzed=45
)

class Memory:
    """Enhanced Memory API with ACE capabilities."""

    # ========== Existing methods (unchanged) ==========
    def store(self, content, scope="user", memory_type="fact", importance=0.8) -> str: ...
    def query(self, user_id=None, scope=None, memory_type=None) -> list[dict]: ...
    def search(self, query, scope=None, limit=10) -> list[dict]: ...

    # ========== NEW: Trajectory Management ==========

    def start_trajectory(
        self,
        task_description: str,
        task_type: str | None = None,
        parent_trajectory_id: str | None = None
    ) -> str:
        """Start tracking an execution trajectory.

        Args:
            task_description: Description of the task
            task_type: Type of task ('api_call', 'reasoning', etc.)
            parent_trajectory_id: Parent trajectory for multi-step workflows

        Returns:
            trajectory_id: ID for tracking this execution

        Example:
            >>> traj_id = nx.memory.start_trajectory(
            ...     task_description="Fetch user data from API",
            ...     task_type="api_call"
            ... )
        """

    def log_step(
        self,
        trajectory_id: str,
        step_type: str,
        description: str,
        metadata: dict | None = None
    ) -> None:
        """Log a step in the trajectory.

        Args:
            trajectory_id: Trajectory ID from start_trajectory()
            step_type: Type of step ('decision', 'action', 'observation')
            description: What happened
            metadata: Additional context (tool calls, LLM responses, etc.)

        Example:
            >>> nx.memory.log_step(
            ...     traj_id,
            ...     step_type="decision",
            ...     description="Decided to use exponential backoff",
            ...     metadata={"retry_count": 3, "delay_ms": 2000}
            ... )
        """

    def complete_trajectory(
        self,
        trajectory_id: str,
        status: Literal["success", "failure", "partial"],
        success_score: float | None = None,
        error_message: str | None = None,
        metrics: dict | None = None
    ) -> None:
        """Mark trajectory as complete with outcome.

        Args:
            trajectory_id: Trajectory ID
            status: Final status
            success_score: 0.0-1.0 success rating
            error_message: Error if failed
            metrics: Performance metrics (tokens, cost, duration)

        Example:
            >>> nx.memory.complete_trajectory(
            ...     traj_id,
            ...     status="success",
            ...     success_score=0.95,
            ...     metrics={"tokens": 1500, "cost_usd": 0.03, "duration_ms": 2340}
            ... )
        """

    # ========== NEW: Reflection & Analysis ==========

    def reflect(
        self,
        trajectory_id: str,
        llm_client: Any | None = None
    ) -> dict[str, Any]:
        """Analyze a trajectory to extract insights.

        Uses LLM to analyze what worked, what failed, and what patterns emerged.
        Stores reflection as a memory with type='reflection'.

        Args:
            trajectory_id: Trajectory to analyze
            llm_client: LLM client (uses default if None)

        Returns:
            Reflection analysis with extracted insights

        Example:
            >>> reflection = nx.memory.reflect(traj_id)
            >>> print(reflection['insights'])
            {
                'successful_patterns': [
                    'Used exponential backoff for retries',
                    'Validated input before API call'
                ],
                'failure_patterns': [
                    'Did not check rate limit headers'
                ],
                'recommendations': [
                    'Add rate limit monitoring',
                    'Cache API responses for 5 minutes'
                ]
            }
        """

    def batch_reflect(
        self,
        agent_id: str | None = None,
        since: str | None = None,
        task_type: str | None = None,
        min_trajectories: int = 5
    ) -> list[dict[str, Any]]:
        """Reflect on multiple trajectories to find patterns.

        Args:
            agent_id: Filter by agent
            since: ISO timestamp - only trajectories after this
            task_type: Filter by task type
            min_trajectories: Minimum trajectories needed for pattern detection

        Returns:
            List of pattern insights across trajectories

        Example:
            >>> patterns = nx.memory.batch_reflect(
            ...     agent_id="agent_123",
            ...     since="2025-10-01T00:00:00Z",
            ...     min_trajectories=10
            ... )
        """

    # ========== NEW: Playbook Management ==========

    def get_playbook(
        self,
        name: str = "default",
        agent_id: str | None = None
    ) -> dict[str, Any] | None:
        """Get playbook for agent.

        Args:
            name: Playbook name
            agent_id: Agent ID (uses current if None)

        Returns:
            Playbook content with strategies

        Example:
            >>> playbook = nx.memory.get_playbook("default")
            >>> for strategy in playbook['strategies']['helpful']:
            ...     print(f"✓ {strategy['pattern']}")
        """

    def update_playbook(
        self,
        strategies: list[dict],
        name: str = "default",
        merge: bool = True
    ) -> str:
        """Update playbook with new strategies.

        Args:
            strategies: New strategies to add
            name: Playbook name
            merge: If True, merge with existing; if False, replace

        Returns:
            playbook_id: Updated playbook ID

        Example:
            >>> nx.memory.update_playbook([
            ...     {
            ...         'category': 'helpful',
            ...         'pattern': 'Always validate inputs before API calls',
            ...         'context': 'API integration tasks',
            ...         'confidence': 0.9,
            ...         'evidence_count': 5
            ...     }
            ... ])
        """

    def curate_playbook(
        self,
        reflections: list[str],  # reflection memory_ids
        playbook_name: str = "default",
        llm_client: Any | None = None
    ) -> dict[str, Any]:
        """Automatically curate playbook from reflections.

        Uses LLM to extract strategies from reflection memories and
        update the playbook, removing duplicates and merging similar patterns.

        Args:
            reflections: List of reflection memory IDs to analyze
            playbook_name: Target playbook
            llm_client: LLM client

        Returns:
            Curation result with diff of what changed

        Example:
            >>> result = nx.memory.curate_playbook(
            ...     reflections=["mem_123", "mem_456", "mem_789"]
            ... )
            >>> print(f"Added: {len(result['added'])}")
            >>> print(f"Updated: {len(result['updated'])}")
            >>> print(f"Removed: {len(result['removed'])}")
        """

    # ========== NEW: Consolidation ==========

    def consolidate(
        self,
        memory_type: str | None = None,
        scope: str | None = None,
        min_importance: float = 0.5,
        preserve_high_importance: bool = True,
        importance_threshold: float = 0.8,
        llm_client: Any | None = None
    ) -> dict[str, Any]:
        """Consolidate memories to prevent context collapse.

        Implements importance-based preservation from ACE paper.
        Merges similar low-importance memories while preserving
        high-importance details.

        Args:
            memory_type: Filter by type
            scope: Filter by scope
            min_importance: Only consolidate memories below this
            preserve_high_importance: Never consolidate high-importance memories
            importance_threshold: Threshold for "high importance"
            llm_client: LLM for semantic merging

        Returns:
            Consolidation report

        Example:
            >>> report = nx.memory.consolidate(
            ...     memory_type="experience",
            ...     scope="agent",
            ...     min_importance=0.5,
            ...     importance_threshold=0.8
            ... )
            >>> print(f"Consolidated {report['merged_count']} memories")
            >>> print(f"Preserved {report['preserved_count']} high-importance")
        """

    # ========== NEW: ACE Learning Loop ==========

    def execute_with_learning(
        self,
        task_fn: callable,
        task_description: str,
        task_type: str | None = None,
        auto_reflect: bool = True,
        auto_curate: bool = True,
        playbook_name: str = "default"
    ) -> tuple[Any, str]:
        """Execute a task with automatic learning loop.

        This is the main ACE integration method that wraps task execution
        with trajectory tracking, reflection, and playbook curation.

        Args:
            task_fn: Function to execute (should return result and status)
            task_description: Description of task
            task_type: Type of task
            auto_reflect: Automatically reflect after completion
            auto_curate: Automatically update playbook from reflection
            playbook_name: Playbook to use/update

        Returns:
            (result, trajectory_id): Task result and trajectory ID

        Example:
            >>> def my_task():
            ...     # Load playbook strategies
            ...     playbook = nx.memory.get_playbook()
            ...     strategies = playbook['strategies']['helpful']
            ...
            ...     # Execute with learned strategies
            ...     result = call_api_with_strategies(strategies)
            ...     return result, "success"
            ...
            >>> result, traj_id = nx.memory.execute_with_learning(
            ...     my_task,
            ...     task_description="Call external API",
            ...     task_type="api_call"
            ... )
        """

# trajectory.py
class TrajectoryManager:
    def start(self, task_description, task_type=None) -> str: ...
    def log_step(self, trajectory_id, step_type, description, metadata=None): ...
    def complete(self, trajectory_id, status, success_score, error_message=None): ...
    def get_trajectory(self, trajectory_id) -> dict: ...
    def query_trajectories(self, filters) -> list[dict]: ...

# reflection.py
class Reflector:
    def reflect(self, trajectory_id, llm_client=None) -> dict: ...
    def batch_reflect(self, trajectories, llm_client=None) -> list[dict]: ...
    def extract_patterns(self, reflections) -> dict: ...

# curation.py
class Curator:
    def curate(self, reflections, playbook_name) -> dict: ...
    def merge_strategies(self, existing, new) -> list: ...
    def remove_duplicates(self, strategies) -> list: ...
    def calculate_confidence(self, strategy, evidence) -> float: ...

# playbook.py
class PlaybookManager:
    def get(self, name, agent_id=None) -> dict: ...
    def create(self, name, content, scope="agent") -> str: ...
    def update(self, playbook_id, content, increment_version=True) -> str: ...
    def delete(self, playbook_id) -> bool: ...
    def list_versions(self, name, agent_id=None) -> list: ...

# consolidation.py
class ConsolidationEngine:
    def consolidate(
        self,
        memory_type=None,
        scope=None,
        strategy="importance_based"
    ) -> dict: ...
    def identify_similar_memories(self, memories) -> list[list]: ...
    def merge_memories(self, memory_ids, llm_client=None) -> str: ...
    def preserve_high_importance(self, memories, threshold=0.8) -> list: ...

REFLECTION_PROMPT = """
Analyze the following task execution trajectory and extract insights:

Task: {task_description}
Status: {status}
Success Score: {success_score}

Execution Steps:
{steps}

Please identify:
1. What strategies were successful? (✓ Helpful patterns)
2. What approaches failed? (✗ Harmful patterns)
3. What observations are context-dependent? (○ Neutral)

Format your response as:
{
  "successful_patterns": [...],
  "failure_patterns": [...],
  "neutral_observations": [...],
  "recommendations": [...]
}
"""

import nexus

nx = nexus.connect()

# Define your agent's task
def api_task():
    # 1. Load playbook to get learned strategies
    playbook = nx.memory.get_playbook("api_client")
    strategies = playbook['strategies']['helpful']

    # 2. Apply strategies
    for strategy in strategies:
        print(f"Applying: {strategy['pattern']}")

    # 3. Execute task
    response = call_external_api()

    # 4. Return result and status
    return response, "success"

# Execute with automatic learning
result, traj_id = nx.memory.execute_with_learning(
    api_task,
    task_description="Call external API for user data",
    task_type="api_call",
    auto_reflect=True,
    auto_curate=True
)

# Playbook automatically updated with insights!

# 1. Start trajectory
traj_id = nx.memory.start_trajectory(
    task_description="Process invoice PDFs",
    task_type="document_processing"
)

# 2. Execute task with logging
try:
    nx.memory.log_step(traj_id, "decision", "Checking PDF format")
    pdf_format = detect_format(pdf_file)

    nx.memory.log_step(traj_id, "action", f"Parsing as {pdf_format}")
    data = parse_pdf(pdf_file, format=pdf_format)

    nx.memory.log_step(traj_id, "observation", f"Extracted {len(data)} fields")

    # Task succeeded
    nx.memory.complete_trajectory(
        traj_id,
        status="success",
        success_score=0.95,
        metrics={"fields_extracted": len(data), "duration_ms": 1500}
    )

except Exception as e:
    nx.memory.log_step(traj_id, "observation", f"Error: {str(e)}")
    nx.memory.complete_trajectory(
        traj_id,
        status="failure",
        success_score=0.0,
        error_message=str(e)
    )

# 3. Reflect on what happened
reflection = nx.memory.reflect(traj_id)
print(reflection['insights'])

# 4. Update playbook
nx.memory.curate_playbook([reflection['memory_id']])

# After running 50 API calls, reflect on patterns
patterns = nx.memory.batch_reflect(
    agent_id="api_agent",
    since="2025-10-01T00:00:00Z",
    task_type="api_call",
    min_trajectories=10
)

# Get all reflection memory IDs
reflection_ids = [p['memory_id'] for p in patterns]

# Curate playbook from patterns
result = nx.memory.curate_playbook(
    reflections=reflection_ids,
    playbook_name="api_client"
)

print(f"✓ Added {len(result['added'])} new strategies")
print(f"↻ Updated {len(result['updated'])} existing strategies")
print(f"✗ Removed {len(result['removed'])} outdated strategies")

# Prevent context collapse by consolidating old memories
report = nx.memory.consolidate(
    memory_type="experience",
    scope="agent",
    min_importance=0.5,
    preserve_high_importance=True,
    importance_threshold=0.8
)

print(f"Consolidated {report['merged_count']} memories")
print(f"Preserved {report['preserved_count']} high-importance memories")
print(f"Saved {report['space_saved_bytes']} bytes")

# Agent A learns from task execution
agent_a = nexus.connect(config={"agent_id": "agent_a"})
result, traj_id = agent_a.memory.execute_with_learning(
    task_fn=process_documents,
    task_description="Process PDF invoices"
)

# Agent B can access Agent A's playbook (with ReBAC permissions)
agent_b = nexus.connect(config={"agent_id": "agent_b"})

# Grant permission for Agent B to read Agent A's playbook
agent_a.rebac.create(
    subject=("agent", "agent_b"),
    relation="viewer",
    object=("playbook", f"agent_a/default")
)

# Agent B uses Agent A's learned strategies
playbook = agent_b.memory.get_playbook(
    name="default",
    agent_id="agent_a"  # Read from Agent A
)

# Apply learned strategies
for strategy in playbook['strategies']['helpful']:
    print(f"Learned from Agent A: {strategy['pattern']}")

def consolidate(self, preserve_high_importance=True, importance_threshold=0.8):
    """
    Never consolidate memories above importance threshold.
    For consolidation candidates:
    - Track information loss
    - Preserve named entities, numbers, specific examples
    - Store consolidation lineage for rollback
    """

def calculate_confidence(strategy, evidence_trajectories):
    """
    Confidence = f(
        success_rate across trajectories,
        number of supporting examples,
        recency of evidence,
        consistency of results
    )
    """
    success_count = sum(1 for t in evidence if t.status == "success")
    total = len(evidence_trajectories)
    base_confidence = success_count / total

    # Boost for more evidence
    evidence_boost = min(0.2, len(evidence_trajectories) * 0.01)

    # Decay for old evidence
    recency_factor = calculate_recency_decay(evidence_trajectories)

    return min(1.0, base_confidence + evidence_boost) * recency_factor

def remove_duplicates(strategies, similarity_threshold=0.85):
    """
    Use semantic embeddings to detect duplicate strategies.
    Merge similar strategies and combine evidence counts.
    """
    embeddings = compute_embeddings([s.pattern for s in strategies])

    # Find similar pairs
    similar_pairs = find_similar_above_threshold(
        embeddings,
        threshold=similarity_threshold
    )

    # Merge strategies
    for s1, s2 in similar_pairs:
        merged = merge_strategies(s1, s2)
        strategies.append(merged)
        strategies.remove(s1)
        strategies.remove(s2)

# Create private playbook (default)
playbook_id = nx.memory.create_playbook(
    name="api_strategies",
    scope="agent",
    visibility="private"
)

# Share playbook with team
nx.rebac.create(
    subject=("user", "alice"),
    relation="viewer",
    object=("playbook", playbook_id)
)

# Create tenant-wide playbook
team_playbook = nx.memory.create_playbook(
    name="team_best_practices",
    scope="tenant",
    visibility="shared"
)

# Configuration
TRAJECTORY_RETENTION_DAYS = 30
TRAJECTORY_SAMPLE_RATE = 0.1  # 10% of successes
TRAJECTORY_ALWAYS_STORE_FAILURES = True

# Auto-consolidation triggers
AUTO_CONSOLIDATE_THRESHOLD = 10000  # memories
CONSOLIDATION_BATCH_SIZE = 1000
RUN_CONSOLIDATION_ASYNC = True

# Keep only:
# - Latest 5 versions
# - All major versions
# - Versions from last 90 days
PLAYBOOK_VERSION_RETENTION = {
    "keep_latest": 5,
    "keep_all_major": True,
    "keep_days": 90
}